Deliberation in Latent Space via Differentiable Cache Augmentation

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Liu, Luyang, Pfeiffer, Jonas, Wu, Jiaxing, Xie, Jun, Szlam, Arthur
Formato: Preprint
Publicado: 2024
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866912167483146240
author Liu, Luyang
Pfeiffer, Jonas
Wu, Jiaxing
Xie, Jun
Szlam, Arthur
author_facet Liu, Luyang
Pfeiffer, Jonas
Wu, Jiaxing
Xie, Jun
Szlam, Arthur
contents Techniques enabling large language models (LLMs) to "think more" by generating and attending to intermediate reasoning steps have shown promise in solving complex problems. However, the standard approaches generate sequences of discrete tokens immediately before responding, and so they can incur significant latency costs and be challenging to optimize. In this work, we demonstrate that a frozen LLM can be augmented with an offline coprocessor that operates on the model's key-value (kv) cache. This coprocessor augments the cache with a set of latent embeddings designed to improve the fidelity of subsequent decoding. We train this coprocessor using the language modeling loss from the decoder on standard pretraining data, while keeping the decoder itself frozen. This approach enables the model to learn, in an end-to-end differentiable fashion, how to distill additional computation into its kv-cache. Because the decoder remains unchanged, the coprocessor can operate offline and asynchronously, and the language model can function normally if the coprocessor is unavailable or if a given cache is deemed not to require extra computation. We show experimentally that when a cache is augmented, the decoder achieves lower perplexity on numerous subsequent tokens. Furthermore, even without any task-specific training, our experiments demonstrate that cache augmentation consistently reduces perplexity and improves performance across a range of reasoning-intensive tasks.
format Preprint
id arxiv_https___arxiv_org_abs_2412_17747
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Deliberation in Latent Space via Differentiable Cache Augmentation
Liu, Luyang
Pfeiffer, Jonas
Wu, Jiaxing
Xie, Jun
Szlam, Arthur
Computation and Language
Artificial Intelligence
Machine Learning
Techniques enabling large language models (LLMs) to "think more" by generating and attending to intermediate reasoning steps have shown promise in solving complex problems. However, the standard approaches generate sequences of discrete tokens immediately before responding, and so they can incur significant latency costs and be challenging to optimize. In this work, we demonstrate that a frozen LLM can be augmented with an offline coprocessor that operates on the model's key-value (kv) cache. This coprocessor augments the cache with a set of latent embeddings designed to improve the fidelity of subsequent decoding. We train this coprocessor using the language modeling loss from the decoder on standard pretraining data, while keeping the decoder itself frozen. This approach enables the model to learn, in an end-to-end differentiable fashion, how to distill additional computation into its kv-cache. Because the decoder remains unchanged, the coprocessor can operate offline and asynchronously, and the language model can function normally if the coprocessor is unavailable or if a given cache is deemed not to require extra computation. We show experimentally that when a cache is augmented, the decoder achieves lower perplexity on numerous subsequent tokens. Furthermore, even without any task-specific training, our experiments demonstrate that cache augmentation consistently reduces perplexity and improves performance across a range of reasoning-intensive tasks.
title Deliberation in Latent Space via Differentiable Cache Augmentation
topic Computation and Language
Artificial Intelligence
Machine Learning
url https://arxiv.org/abs/2412.17747