PaperlayerView exact revision on arXiv

Source-linked AI summary

Deliberation in Latent Space via Differentiable Cache Augmentation

Luyang Liu, Jonas Pfeiffer, Jiaxing Wu, Jun Xie, Arthur Szlam

arXiv:2412.17747v1cs.CLcs.AIcs.LG

TL;DR

Discrete, just-in-time reasoning can be difficult to optimize end to end and adds computation during generation. This paper trains a separate coprocessor to augment a frozen LLM’s kv-cache with latent embeddings; the authors report improvements across reasoning-intensive tasks without task-specific fine-tuning.

  • Problem

    Many extra-thinking approaches use discrete intermediate outputs and perform computation just in time, making end-to-end training difficult and placing extra computation in the response-generation process.

    Source: §1 · p. 1

  • Method

    A separate coprocessor generates latent embeddings to augment a frozen LLM’s kv-cache, trained end to end using language-modeling loss on pretraining data and multi-position future-token prediction.

    Source: §1 · pp. 1–2, §2.3 · p. 4

  • Results

    Without task-specific fine-tuning, cache augmentation consistently improves reasoning-intensive task performance, including a 10.05% improvement on GSM8K with 64 latents.

    Source: §1 · p. 2, §3.2 · p. 6

  • Takeaways & Limitations

    Because the base transformer remains frozen and improvements extend many tokens beyond augmentation, the coprocessor can operate asynchronously.

    Source: §1 · p. 2

Abstract

from arXiv · show

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.

1. Introduction

The section contrasts discrete, just-in-time reasoning with a differentiable approach that augments a frozen LLM’s kv-cache using a separate coprocessor. The authors report improved performance across reasoning-intensive tasks without task-specific fine-tuning.

1. Introduction

  • Standard extra-thinking approaches use discrete intermediate sequences and perform computation just in time, making end-to-end training difficult and adding computation during response generation.Source: §1 · p. 1
  • The proposed method uses a separate coprocessor to generate latent embeddings in one forward pass and augment the frozen transformer’s kv-cache, enabling end-to-end backpropagation without reinforcement learning.Source: §1 · p. 1
  • Standard language-modeling loss on pretraining data supports scalable coprocessor training, while improvements many tokens after augmentation make asynchronous operation feasible.The base transformer remains frozen during coprocessor training.Source: §1 · p. 2
  • Without task-specific fine-tuning, cache augmentation improves performance across reasoning-intensive tasks, including a 10.05% improvement on GSM8K and a 4.70% improvement on MMLU.The authors also report that increasing the number of injected latent embeddings generally improves performance.Source: §1 · p. 2

2. Methodology

The methodology trains a coprocessor to augment a frozen LLM’s kv-cache with latent embeddings that help predict target text. Its pretraining setup places augmentations at multiple sequence positions and predicts future tokens from each.

2.1. Problem statement

  • The coprocessor maps a frozen LLM’s kv-cache to latent representations, trained to improve the frozen model’s expected log-likelihood of the target given the input and latents.Source: §2.1 · p. 2, §2 · p. 2

2.2. Model architecture

  • The frozen LLM creates the kv-cache; a coprocessor using trainable soft tokens produces latent embeddings that are appended to the cache before the LLM generates output.Training updates the coprocessor and trainable embeddings, not the frozen LLM.Source: §2.2 · pp. 2–3, Figure 1 · p. 3

2.3. Pretraining setup

  • At randomly selected sequence positions, the coprocessor generates configurable latent embeddings and is teacher-forced to predict multiple future tokens from the preceding context.The paper describes this as bridging known preceding context and future context through latent representations.Source: §2.3 · pp. 3–4, Figure 2 · p. 4
  • A single extended input and attention mask let training compute losses from multiple augmentation positions in parallel while restricting each latent group and its target tokens to preceding original tokens.Source: §2.3 · pp. 4–5
  • The embeddings are trained to provide information useful for future-token prediction, a form of latent thinking that differs from standard next-token prediction.The passage says this approach allows the coprocessor to learn more generalizable representations.Source: §2.3 · p. 5

3. Experiments

Experiments evaluate cache augmentation with a frozen Gemma-2 2B model and no task-specific training. The method improves perplexity and performance across reasoning-intensive benchmarks, while results vary with design choices and task.

3.1. Perplexity Evaluation

  • Cache augmentation consistently reduces perplexity across latent sizes and subsequent-token positions, with benefits extending to position 32 despite training with 16 ahead tokens.The reduction generally correlates with the number of latents.Source: §3.1 · p. 6, Table 1 · p. 6

3.2. Public Benchmark Evaluation

  • +10.05% GSM8K gain with 64 latents accompanies broad improvements over frozen Gemma-2 2B, though some tasks show dips at higher latent counts.The coprocessor is called once at the end of the prompt.Source: §3.2 · p. 6, Table 2 · p. 7

3.3.1. Pause Token

  • 10.60 first-token perplexity beats the Gemma-2 2B baseline (10.96) and Pause Token (11.63), while 26.76% GSM8K accuracy exceeds 21.38% and 22.37%, respectively.Both latent methods use 32 embeddings, and the Pause Token baseline uses the same training data and experimental setup.Source: §3.3.1 · p. 7, Table 3 · p. 8

3.3.2. Zero-shot CoT

  • 24.72% GSM8K accuracy with 16 latents exceeds the baseline (21.38%) and zero-shot CoT (23.20%); accuracy reaches 26.76% with 32 latents.Zero-shot CoT generates intermediate reasoning tokens sequentially, while this method generates latent embeddings in one forward pass.Source: Table 4 · p. 8, §3.3.2 · p. 8

3.3.3. Alternative Coprocessor Configurations

  • Finetuning from pretrained Gemma-2 2B weights outperforms training the coprocessor from scratch on GSM8K as latent count increases.Training from scratch still improves over the baseline across downstream tasks.Source: §3.3.3 · p. 8, Figure 4 · p. 9
  • 23.35% and 24.03% GSM8K accuracy for LoRA ranks 64 and 128 fall between the baseline (21.38%) and scratch training (25.78%); last-layer activations also trail kv-cache augmentation.With 32 latents, last-layer activations yield 10.81 validation perplexity and 23.20% GSM8K accuracy, versus 10.69 and 26.76% for kv-cache augmentation.Source: Table 5 · p. 9, §3.3.3 · pp. 9–10

3.4. Impact of the number of ahead token in training

  • 16 ahead tokens achieves the highest reported GSM8K accuracy, 24.72%, while larger lookahead often improves later-token perplexity but worsens earlier-token perplexity.The authors chose 16 ahead tokens for most experiments.Source: Table 6 · p. 10, §3.4 · p. 10

4. Related Work

Related work studies explicit and latent reasoning, cache compression, external modules, and hypernetworks. This work instead augments a frozen model’s kv-cache with input-conditioned latent embeddings from a coprocessor.

4.1. Chain-of-Thought Reasoning in LLMs

4.2. Latent Space Reasoning

  • Prior studies identify latent reasoning pathways and interpretable middle-layer embeddings; related methods use intermediate tokens, while this work generates them dynamically from input.Unlike COCONUT, the approach trains only the coprocessor rather than using multi-stage training to internalize reasoning.Source: §4.2 · pp. 10–11

4.3. KV-Cache Compression

  • Cache-compression methods reduce stored past activations, whereas this work adds coprocessor-produced latent embeddings to the cache to improve reasoning without changing the transformer architecture.Examples of compression approaches include ICAE (Ge et al., 2024) and gist tokens (Mu et al., 2024).Source: §4.3 · p. 11

4.4. Augmenting LLMs with External Modules

  • External-module approaches include adapters, LoRA, and multimodal cross-attention; CALM uses cross-attention between specialized and general models to enhance the general model.The passage also describes prompt tuning, prefix tuning, and other module-augmentation approaches.Source: §4.4 · p. 11

4.5. Hypernetworks for Parameter Generation

  • The coprocessor resembles a hypernetwork conditioned on the frozen model’s kv-cache, but generates context-specific latent augmentations rather than parameters for another module.Hypernetworks dynamically generate outputs from contextual inputs.Source: §4.5 · p. 12

5. Conclusion

The conclusion presents differentiable cache augmentation as a way to improve frozen decoder-only language models using a learned coprocessor on the kv-cache. It reports lower perplexity and better reasoning-task performance, and identifies asynchronous operation and broader scaling as future directions.

5. Conclusion

  • The learned coprocessor augments a frozen decoder-only model’s kv-cache with latent embeddings.The paper calls this approach differentiable cache augmentation.Source: §5 · p. 12
  • The approach consistently reduces perplexity and improves performance across reasoning-intensive tasks, including MMLU and GSM8K.The reported improvements occur in zero/few-shot settings.Source: §5 · p. 12
  • Offline, asynchronous coprocessor operation motivates future research into more deliberate computation, including deliberation not conditioned on a particular response.Future work also includes larger models, modular coprocessors, other architectures, and more diverse downstream tasks.Source: §5 · p. 12

A.1. Scaling with Training Data

As coprocessor training exposure increases, GSM8K accuracy rises and validation perplexity falls. Figure 5 compares these trends with the frozen Gemma-2 2B baseline across training steps.

A.1. Scaling with Training Data

  • GSM8K accuracy rises and validation perplexity falls as coprocessor training steps increase.The passage describes these trends as consistent and reports a clear benefit from increased training exposure.Source: §A.1 · p. 17, Figure 5 · p. 17
  • The passage attributes the trends to the coprocessor learning more useful latent embeddings and integrating better with the frozen LLM as it sees more data.Source: §A.1 · p. 17
  • Figure 5 plots GSM8K accuracy and validation perplexity against coprocessor training steps, with frozen Gemma-2 2B performance as reference.The reference is shown at corresponding intervals along the training-step axis.Source: Figure 5 · p. 17, §A.1 · p. 17

A.2. Adaptation to Downstream Tasks

With LoRA finetuning on GSM8K and MATH training data, the augmented model substantially outperforms the baseline on GSM8K test accuracy. The comparison applies LoRA to the coprocessor for the augmented model and to the base LLM for the baseline.

A.2. Adaptation to Downstream Tasks

  • The adaptation experiment uses GSM8K and MATH training sets, applying rank-128 LoRA to the coprocessor while keeping the augmented model’s base LLM frozen.For the baseline, LoRA is applied directly to the base LLM.Source: §A.2 · p. 17
  • The authors suggest that freezing the base LLM may regularize adaptation and limit overfitting on the relatively small downstream datasets.They contrast this with possible overfitting when LoRA is applied directly to the LLM.Source: §A.2 · pp. 17–18
  • The augmented model substantially improves GSM8K test accuracy over the baseline after LoRA finetuning.Figure 6 presents the test-set comparison.Source: §A.2 · p. 17, Figure 6 · p. 18

A.3. Training Coprocessor from Scratch

Training the coprocessor from randomly initialized weights produced performance gains across most benchmarks. The supplied section does not state the result of its comparison with pretrained initialization.

A.3. Training Coprocessor from Scratch

  • Training the coprocessor from scratch with randomly initialized weights produced gains across most benchmarks.Source: §A.3 · p. 18
  • The scratch-trained coprocessor uses randomly initialized weights, rather than initialization from the pretrained LLM.Source: §A.3 · p. 18
  • Table 7 covers baseline and augmented model performance across various benchmarks for scratch-trained coprocessors.The caption points to Table 2 for a more detailed description.Source: Table 7 · p. 18
Loading 2412.17747v1…