Deliberation in Latent Space via Differentiable Cache Augmentation | Summary

This article explains the key points of Deliberation in Latent Space via Differentiable Cache Augmentation.

• 2024-12-23 (arXiv)
• Liu, Luyang, Pfeiffer, Jonas, Wu, Jiaxing, Xie, Jun, Szlam, Arthur.
• Google DeepMind
• Paper

Read this article in Korean

Summary

• Generating intermediate reasoning text before answering adds latency to large language models (LLMs) and makes the reasoning process difficult to optimize end to end. Without task-specific training, adding 64 learned latent embeddings, continuous representations that extend the model’s context, raises Gemma-2 2B accuracy from 21.38% to 31.43% on GSM8K, a mathematical reasoning benchmark, under 8-shot evaluation, which supplies eight examples in the prompt, while keeping the original model’s weights frozen, or unchanged.
• A coprocessor, a separately trained model, reads the decoder’s key-value (KV) cache, stored attention representations of preceding text, and produces input-dependent embeddings in one forward pass. The decoder, which generates the output text, supplies a language-modeling loss measuring target-token prediction error; training updates only the coprocessor and its trainable input embeddings.
• The gains persist beyond the 16-token training prediction horizon: with 64 latent embeddings, perplexity, a measure of token-prediction uncertainty where lower is better, decreases relative to baseline by 3.94% at the first subsequent token and 1.20% at the 32nd. The experiments demonstrate prediction and benchmark improvements for Gemma-2 2B, but report no wall-clock latency or throughput measurements to establish the practical benefit of asynchronous execution, in which coprocessor work proceeds independently of decoding.

1. Introduction

Generating discrete intermediate reasoning steps makes additional computation sequential and complicates end-to-end optimization. The paper instead processes the model’s existing cache and stores the result as continuous context, allowing ordinary decoding to resume after augmentation.

• Because the base model remains frozen, it can still generate normally when the coprocessor is unavailable.
• Benefits that persist across subsequent tokens motivate performing augmentation independently of immediate response generation.

2. Methodology

The method augments a pretrained, frozen decoder with a coprocessor that consumes its KV cache and generates soft tokens, continuous embeddings without corresponding words or subwords. Training propagates the output loss through the frozen decoder to learn input-dependent augmentations without changing the decoder’s parameters.

2.1. Problem statement

For input x and target y, the frozen model has parameters θ and produces cache (k_{θ,x}, v_{θ,x}); the learned coprocessor maps this cache to latent representations through f(k_{θ,x}, v_{θ,x}) → z. The stated objective is max E_x[log p_θ(y | x, z)], so the coprocessor learns representations that improve target likelihood while θ stays fixed.

2.2. Model architecture

The frozen decoder first processes the input to create its cache. A coprocessor with the same architecture, initialized from the decoder’s pretrained weights, processes that cache together with distinct trainable soft tokens and produces latent embeddings that augment the original context before output generation.

• The output loss backpropagates through the frozen decoder to train the coprocessor and soft-token embeddings.
• All augmentation embeddings are produced in one forward call rather than through sequential text generation.

The frozen decoder and trainable coprocessor communicate through cached representations and generated latent embeddings. This interface allows the output loss to train the additional computation while preserving the original decoder’s weights.

2.3. Pretraining setup

Training samples multiple augmentation positions within each sequence and predicts \(N_A\) future tokens from each position using \(N_L\) latent embeddings. Teacher forcing, supplying the actual preceding target tokens during prediction, lets the loss supervise a segment beyond the augmentation point.

• For the example sequence “a b c d e f,” augmentation after “b” uses the preceding cache to support predictions of “c” and “d”; augmentation after “d” supports “e” and “f.”
• Modified inputs, position indices, targets, and attention masks, which determine which positions can attend to one another, combine multiple augmentation branches into a single decoder forward pass.
• Each branch can use its original prefix and permitted earlier branch positions, while later original tokens and unrelated augmentation branches are masked.

The branches after “b” and “d” share an original sequence but receive separate latent augmentations. The attention matrix shows how branch losses can be computed together while preventing access to later original text and unrelated branches.

3. Experiments

The main experiments use frozen Gemma-2 2B and the same primarily English, 2 trillion-token dataset used for its pretraining, including web documents, code, and scientific articles. The default setup runs for 100,000 steps with batch size 1024, packed sequence length 2048, 16 ahead tokens, meaning future target tokens after augmentation, and 128 randomly sampled augmentation positions, called “traces.”

• The main public benchmark results use coprocessors trained on pretraining data without task-specific fine-tuning.
• The ahead-token ablation varies the default prediction horizon, and the appendix separately evaluates adaptation using task training sets.

3.1. Perplexity Evaluation

Evaluation uses Gemma’s proprietary validation dataset and measures prediction at specified distances after augmentation. The 32nd-token result conditions on the original prefix, latent embeddings, and the following 31 actual tokens, measuring sustained conditional prediction benefit.

• For 8, 16, 32, and 64 latents, first-token relative perplexity changes are −1.53%, −2.48%, −3.28%, and −3.94%, respectively.
• At position 32, the corresponding changes are −0.55%, −0.64%, −0.88%, and −1.20%, despite training with only 16 ahead tokens.
• Figure 3 also plots a 4-latent configuration, although its caption lists only 8, 16, 32, and 64.

Perplexity falls during coprocessor training for both the first and 32nd subsequent tokens, with larger latent counts generally producing larger reductions. The second panel shows that conditional prediction benefits persist beyond the 16-token training horizon.

The reductions generally diminish with distance from augmentation but remain present at position 32. With 64 latents, the relative perplexity change is −3.94% at position 1 and −1.20% at position 32.

3.2. Public Benchmark Evaluation

Public benchmark evaluation calls the coprocessor once, at the end of the prompt, with 4, 8, 16, 32, or 64 latents. At 64 latents, MMLU, a language understanding benchmark, rises from 52.00% to 56.70% in 5-shot evaluation, GSM8K rises from 21.38% to 31.43% in 8-shot evaluation, and MATH, a mathematical problem-solving dataset, rises from 16.50% to 18.56% in 4-shot evaluation; each shot count specifies the number of prompt examples.

• The corresponding gains are 4.70, 10.05, and 2.06 percentage points, calculated as absolute score differences.
• All 17 benchmark scores in the 64-latent column exceed baseline, but improvement is not monotonic across every latent count.
• For example, Boolq, one of the evaluated reasoning benchmarks, peaks at 78.17% with 16 latents and falls to 76.91% with 64, against baseline 75.41%, in 0-shot evaluation without prompt examples.

A single augmentation at the prompt boundary improves all 17 reported benchmark scores at 64 latents. Smaller latent counts include regressions, and several benchmarks vary nonmonotonically as latent count increases.

3.3. Comparison with other baselines and variations

3.3.1. Pause Token compares against trainable embeddings placed between input and output whose values are not generated conditionally on the input. With 32 embeddings and the same training data and experimental setup, cache augmentation achieves validation perplexity 10.60 and GSM8K 8-shot accuracy 26.76%, compared with 11.63 and 22.37% for Pause Token and 10.96 and 21.38% for the baseline.

• The comparison supports the trained cache-conditioned coprocessor, although it does not hold additional model computation constant.

At 32 embeddings, cache augmentation improves both metrics, whereas Pause Token increases validation perplexity from 10.96 to 11.63 despite raising GSM8K accuracy. The comparison shows that benchmark accuracy and language-modeling perplexity can respond differently to augmentation.

3.3.2. Zero-shot CoT evaluates Chain-of-Thought (CoT), prompting the model to generate intermediate reasoning steps, by appending “Let’s think step by step” to the prompt. Table 4 reports GSM8K accuracy of 23.20% for this condition, 24.72% with 16 latents, and 26.76% with 32 latents, against baseline 21.38%; its caption identifies the benchmark setting as GSM8K 8-shot.

• The paper labels the intervention “0-shot CoT” while retaining the 8-shot benchmark designation in the caption; it does not further explain this terminology.
• The accuracy comparison includes no timing measurements.

The reported “0-shot CoT” condition reaches 23.20% accuracy, below 24.72% with 16 latents and 26.76% with 32. The caption identifies GSM8K 8-shot evaluation; the table supplies no inference-cost measurements.

3.3.3. Alternative Coprocessor Configurations examines initialization, restricted parameter updates, and the source of input representations. Pretrained initialization outperforms random initialization on GSM8K across the plotted latent counts, while Low-Rank Adaptation (LoRA), training low-rank weight updates with pretrained weights frozen, gives smaller gains than full coprocessor fine-tuning.

• With 32 latents, GSM8K accuracy is 23.35% for LoRA rank 64, 24.03% for rank 128, 25.78% for training from scratch, and 26.76% for full fine-tuning, versus baseline 21.38%; rank sets the size of the low-rank update.
• Using only the frozen model’s last-layer activations, its final-layer internal representations, gives perplexity 10.81 and GSM8K accuracy 23.20%, compared with 10.69 and 26.76% for KV-cache input in this paragraph.
• The KV-cache perplexity of 10.69 in this ablation differs from the 10.60 reported for 32 latents in Table 3; the paper does not reconcile the values.
• Figure 4 estimates means and 95% confidence intervals from the last 5 checkpoints, saved training states, rather than from independent training runs.

Both scratch-trained and pretrained-initialized coprocessors improve GSM8K accuracy as latent count grows, with pretrained initialization giving higher plotted means. The shaded 95% confidence intervals are estimated from the last 5 checkpoints and therefore do not measure variability across independent training runs.

LoRA rank 128 improves over rank 64, but both trail scratch training and full fine-tuning at 32 latents. The table quantifies accuracy under restricted parameter updates without measuring memory consumption or runtime.

3.4. Impact of the number of ahead token in training

With 16 latent embeddings, training on 4, 8, 16, or 32 ahead tokens yields GSM8K accuracies of 24.03%, 24.11%, 24.72%, and 23.73%, respectively, versus baseline 21.38%. The authors select 16 ahead tokens because it performs best in this comparison; they report that longer prediction horizons improve later-token perplexity while often worsening earlier-token perplexity.

With latent count held at 16, accuracy peaks at a 16-token training horizon and drops when the horizon increases to 32. This supports the selected horizon among the tested configurations.

4. Related Work

The related-work discussion connects cache augmentation to explicit reasoning, computation in continuous internal representations, context compression, external trainable modules, and conditional parameter generation. These connections explain how the coprocessor extends a frozen language model through its cached context.

4.1. Chain-of-Thought Reasoning in LLMs

Prior CoT methods generate, aggregate, verify, or search over intermediate text reasoning. The paper also discusses reinforcement learning, optimizing reasoning using signals such as answer correctness or target likelihood, and identifies sequential reasoning-token generation as a source of training and inference cost.

4.2. Latent Space Reasoning

Prior latent-reasoning work studies intermediate transformer representations and pause or filler tokens, additional positions used for computation without adding ordinary text content. The paper distinguishes its input-dependent embeddings from fixed learned token embeddings and contrasts its coprocessor-only training with COCONUT, a method that uses hidden states as latent input embeddings and employs multi-stage training.

4.3. KV-Cache Compression

KV-cache compression reduces memory requirements by representing context more compactly, including through learned compressed embeddings or “gist tokens,” tokens trained to summarize context. This paper draws inspiration from those interfaces but adds contextual representations to support prediction and reasoning; its experiments do not establish a reduction in cache size.

4.4. Augmenting LLMs with External Modules

The paper relates its approach to prompt tuning and prefix tuning, which learn continuous prompts, and adapters, small trainable modules inserted into a pretrained model, alongside LoRA. It also discusses cross-attention, attention between representations from different components, as an interface for combining models, including CALM, a system connecting specialized and general language models.

4.5. Hypernetworks for Parameter Generation

A hypernetwork is a network that generates another network’s parameters conditionally on contextual input. The authors draw a conceptual analogy: their coprocessor dynamically generates cache augmentations from the frozen model’s cache, although its outputs are latent embeddings rather than parameters of a separate module.

5. Conclusion

The paper demonstrates that differentiable cache augmentation improves a frozen decoder’s future-token prediction and benchmark performance in settings with no prompt examples or a small number of examples. Larger models, multiple modular coprocessors, alternative coprocessor architectures, and broader task coverage remain proposed future investigations.

Appendix

• A. Appendix reports better GSM8K accuracy and validation perplexity with longer coprocessor training, without separating data exposure from optimization duration. Task adaptation updates the baseline decoder with LoRA but only the coprocessor in the augmented model, whose decoder stays frozen; the augmented model improves while the baseline declines after an initial gain, and the authors’ reduced-overfitting explanation remains untested. Random initialization improves most benchmarks; pretrained initialization gives stronger GSM8K results, with a Winogrande exception to its broader advantage.

Brief Thoughts

The benchmark gains and conditional prediction improvements at position 32 show that a frozen decoder can use learned, input-dependent cache additions trained through its prediction loss. Practical efficiency remains unresolved: the default coprocessor has the decoder’s architecture, evaluation centers on Gemma-2 2B, validation data are proprietary, and the paper reports no end-to-end latency, throughput, or total compute comparison. The appendix adds evidence on training duration and downstream adaptation, while checkpoint-based intervals and adaptation curves leave independent-run variability and the proposed overfitting mechanism untested.