title: Deliberation in Latent Space via Differentiable Cache Augmentation
summary: 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.
paper_authors._id: 676a3ec0b91a321e164a8780
paper_authors.name: Luyang Liu
paper_authors._id: 676a3ec0b91a321e164a8781
paper_authors.name: Jonas Pfeiffer
paper_authors._id: 676a3ec0b91a321e164a8782
paper_authors.name: Jiaxing Wu
paper_authors._id: 676a3ec0b91a321e164a8783
paper_authors.name: Jun Xie
paper_authors._id: 676a3ec0b91a321e164a8784
paper_authors.name: Arthur Szlam
paper_title: Deliberation in Latent Space via Differentiable Cache Augmentation
paper_summary: 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.
paper_ai_keywords: Large language models (LLMs) key-value (kv) cache latent embeddings language modeling loss end-to-end differentiable perplexity