Deliberation in Latent Space via Differentiable Cache Augmentation
Authors: Luyang Liu
Jonas Pfeiffer
Jiaxing Wu
Jun Xie
Arthur Szlam
Venue: ICML 2025 poster (2025)
Keywords: Latent reasoning
Cache augmentation
LLM
Primary area: deep_learning->large_language_models
TL;DR: This paper introduces a method to improve LLM reasoning by using a separate "coprocessor" to add learned latent information to the LLM's KV-cache without changing the original LLM, leading to better performance on complex tasks.

Abstract: 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 improves performance across a range of reasoning-intensive tasks.

Deliberation in Latent Space via Differentiable Cache Augmentation
Luyang Liu'
Jonas Pfeiffer
Jiaxing Wu
Jun Xie! Arthur Szlam
Abstract
sequences of steps before producing a final answer can sig- nificantly improve performance reasoning tasks. More broadly, providing LLMs with the ability to allocate com- pute adaptively during generation can lead to more effective generation within a fixed compute budget (Schuster et al., 2022). However, at a high level, many these "extra think- ing" approaches are similar in that their sequences of inter- mediate outputs are discrete, making them difficult train in end-to-end fashion, and in that their extra "thinking" (i.e. computation) is performed just-in-time, as part of the output generating process.
Techniques enabling large language models (LLMs) to "think more" by generating and attend- ing 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 optimize. In this work, demonstrate that frozen L.M can be augmented with an offline coprocessor that operates on the model's s key-value (kv) cache. This coprocessor augments the cache with a set latent embed- dings designed to improve the fidelity of subse- quent decoding. We train this coprocessor using the language modeling loss from the decoder on standard pretraining data, while keeping the de- coder itself frozen. This approach enables the model to lear, in an end-to-end differentiable fashion, how to distill additional computation into kv-cache. Because the decoder remains un- changed, the coprocessor can operate offline and asynchronously, and the language model can func- tion normally if the coprocessor is unavailable if 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. Fur- thermore, even without any task-specific training. our experiments demonstrate that cache augmen- tation consistently improves performance across range of reasoning-intensiv tasks.
In this work, we introduce a fundamentally different ap- proach, inspired by the literature on kv-cache compres- sion (Mu al., 2024; Ge al., 2024). Our approach takes step towards LLMs that can deliberate on their memories (encoded in the kv-cache), and distill these deliberations into form usable for subsequent tasks. Specifically, our method processes the transformer's ' cache and augments it with a set of soft tokens produced in a single forward pass- not sequentially. This extra processing is performed by a separate model, which we refer as "'oprocessor", while the base transformer remains frozen. Once the kv-cache is augmented with the coprocessor's 's output (which term "latent embeddings"), decoding proceeds as normal until the coprocessor is called again. This approach offers the following key advantages:
End-to-end Differentiability: Our framework enables end- to-end backpropagation during coprocessor training, fa- cilitating efficient optimization without the need for rein- forcement learning techniques. We leverage the standard ionguage-modeling loss on pre-training data, making the method scalable.
Asynchronous Operation: Because cache augmentation improves results many tokens beyond the augmentation point, and because the base transformer remains frozen during coprocessor training, asynchronous coprocessor op- eration becomes feasible. This contrasts with existing meth- ods where additional computation occurs sequentially, and online. Our approach opens the door to models that can strategically bank computation by deliberating and refining their internal memory, independent of composing a response to query.
1. Introduction
Recent research (Wei et al., 2022; Kojima al., 2022; Wu et al., 2024) has shown that enabling large language mod- els (LLMMs) to generate, or even search over, intermediate
'Google DeepMind. Correspondence to: Luyang Liu <luyang. @ dirranggiuu@gool...comm.. gomm, Arthur Szlam <as- zlam@ maim@ggoole.com..
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
We evaluate our method using Gemma-2 (Team- et al., 2024) models pretrained on diverse dataset mixture.

---

Deliberation in Latent Space via Differentiable Cache Augmentation
Our experiments demonstrate that without any fine-tuning specific downstream tasks, our approach consistently improves performance across range of reasoning-intensinn tasks. We observed that increasing the number of injected la- tent embeddings generally leads to better performance. For example, we observe a 10.05% improvement on GSM8K and a 4.70% improvement on MMLU, when augmenting the Gemma-2 2B model with 64 latent embeddings. These results highlight the potential of enhancing LLMs with kv- cache coprocessing for augmenting model capabilities.
Frozen LM
KV-cache to coprocessor
Augmented cache to LM
2. Methodology
Coprocessor
We enhance a frozen LLM by training a coprocessor that inputs a key-value (kv) cache, and augments it with set of soft tokens. This section details the architecture and training process of our approach.
2.1. Problem statement
Given an input x and a desired target output y, and pre- trained, frozen LLM parameterized by 0, we seek to learn coprocessor, denoted by f. This coprocessor takes the kv-cache (kg... v generated by the frozen LLM when processing the input æ as input, and outputs a sequence of latent representations 2:
Figure I. Overview of the proposed architecture. The input se- quence processed by frozen LLM, generating kv-cache. This cache is then passed to a coprocessor, along with trainable soft tokens. The coprocessor outputs latent embeddings which are used to augment the original kv- before being fed back into the LLM for output generation.
f (kig,es 09 2
(1)
The objective of learning f is to produce latent embeddings that, when combined with the input x. improve the frozen LLM's ability to generate the correct target y. Specifically, we aim to maximize the expected log- ligeeiodl of the target y given the input x and the learned latent embeddings 2, as predicted by the frozen LLM:
words or sub-words but serve as abstract prompts for the co- processor. The coprocessor ingests the kv-cache and these tokens to produce a sequence of latent embeddings, z.
LLM Generation with Augmented Context: appended the original kv-cache. This augmented cache is then fed back into the frozen LLM, providing it with en- riched contextual information derived from the coprocessor. The LLM then proceeds to generate the output sequence, y. conditioned on both the original input & and the copro- cessor's output z. This allows the LLM to leverage the coprocessor's 's latent inferences without requiring to ex- plicitly verbalize intermediate steps.
Finally, 2
max F.g [log po polylo, 2]]
(2)
2.2. Model architecture
Our proposed architecture enhances frozen, pretrained LLM with a dedicated coprocessor module operating on the kv-cache. As illustrated in Figure 1, the interaction between these components unfolds in three stages:
Training focuses solely optimizing the coprocessor and trainable embeddings' weights. The coprocessor shares the same model architecture as the pretrained LLM, and its weights are initialized with the pretrained weights of the LLM. The loss is calculated on the final output y, and backpropagation is used to update only the coprocessor's 's parameters. This targeted training approach allows for ef- ficient fine-tuning without altering the pretrained LLM. In practice, the coprocessor's 's augmentation can potentially be performed offline and asynchronously, in parallel with the LLM's decoding process. This could enable continuous refinement of the LLM's contextual memory, leading to improved efficiency and faster response times.
K V-cache Generation: The input sequence, x. is first pro- cessed by the frozen LLM to generate its corresponding kv-cache (ko.zx U. This cache encapsulates the LLM''s internal representations of the input. Crucially, the LLM's weights remain frozen throughout the entire process.
The ky-cache is then passed to the copro- cessor module, which adopts the same model architecture as the pretrained LLM and therefore has an identical num- ber of parameters as the frozen LLM. The coprocessor also receives a sequence of distinct extra soft tokens with train- able embeddings. These tokens do not correspond to actual
Augmentation:
2

---

Deliberation in Latent Space via Differentiable Cache Augmentation
Text
a
b
C
Latent embeddlnng
b'
d'
Latent
embedding 2
b" ''
loss
"
loss
6'
Ahead taken
d'
Ahead token
f
Allowed Attention
Masked Attention
(8)
Figure 2. Our coprocessor training framework. (a) Illustration of
augmentation and ahead token prediction. For each selected augmentation position, latent embeddings generated by the coprocessor and inserted after the corresponding token't embedding. The target tokens for prediction ("ahead tokens") ") are then appended. A causal mask is applied to all sequences following these insertion points. (b) Structure of the modified input and attention mask for model training. We show example of 1 latent embedding and I ahead token here for simplicity.
2.3. Pretraining setup
bility while keeping the LLM frozen. Furthermore, because the base LLM remains frozen, the coprocessor can be called asynchronously and its computations can be performed in parallel with the LLM's decoding operation, and there no need to insert between every pair of consecutive tokens.
We employ a pretraining strategy designed to encourage the coprocessor to learn augmentations that will be useful for predicting larger segments of text beyond the next token after the augmentation. This strategy involves generating augmentations for multiple randomly selected positions and training the coprocessor predict many tokens ahead.
We implement efficient training framework by modify- ing the input, attention mask, position index, and target, to enable training everything together in one forward pass. Figure 2(b) illustrates how we modify the input and atten- tion mask to enable training on multiple positions. Instead of sequentially processing each augmentation position in multiple decoder forward calls, construct single, tended input sequence and corresponding attention mask. This allows us to compute the loss from all selected augmen- tation points in parallel. For each selected augmentation position, we insert the generated latent embeddings after the corresponding token's 's embedding in the input sequence. The ahead tokens for that position are then appended after the latent embeddings. This creates a concatenated input sequence containing the original text followed by multiple groups of latent embeddings and their corresponding ahead tokens. The attention mask is constructed to ensure correct causal attention. The original tokens attend to all preceding original tokens, as usual. Crucially, the latent embeddings and their corresponding ahead tokens only attend to the original tokens preceding their insertion point. They masked from attending any subsequent original tokens, other latent embeddings, other ahead tokens as shown in Figure 2(b). The resulting output is then used to calculate the loss and train the coprocessor.
Instead of training on single split of sequence into input a and target (which limits scalability), we augment at multiple points within each sequence. As shown in Figure 2 (a), given an input text sequence (c.g., "a c d e f"), we randomly select subset of positions (e-g.. "b" and "d"). For each selected position, the coprocessor generates con- figurable number of latent embeddings (e.g., b', b" and d', d" in the figure, where the number of latent embeddings is hyperparameter NL) in one transformer forward call. The training objective is to predict number of tokens (another hyperparameter A) beyond the placement of the augmenta- tion, in a feacher-forcing way. For instance, if "b" is chosen, and we are predicting two tokens ahead, the coprocessor uses the generated latent embeddings (b', b"") and the kv- cache of the preceding text "a" and "b" to predict "c" and "d". This process can be viewed as a form of latent space interpolation: the coprocessor learns to bridge the gap be- tween the known preceding context ("a", "b") and the future context ("c", "d") by generating meaningful latent represen- tations (b', b²)) Similarly, if dd" chosen, the targets would be "e" and "f", based on d', d" and the preceding context "a b c d". This approach is similar to the parallel decoding introduced in Quiet-Star (Zelikman et al., 2024), but since we generate latent embeddings instead of thoughts in token space, our approach has the advantages of fully differentia-
Rather than capturing different aspects of a single token's 's context, these embeddings are trained to generate informa-
3

---

Deliberation in Latent Space via Differentiable Cache Augmentation
11.2
10.90
Baseline
latents
Baseline latents
latents 16 latents
latents 64 latents
11.1
4 latents
16 latents
64 latents
10.85
11.0
Perplexity
Perplexity
10.9
10.80
10.8
10.7
10.75
10.6
10.5
10.70
20000
40000
60000
80000
100000
20000
40000
60000
80000
100000
Training steps
Training steps
(a) Predicting the 1st token after latents
(b) Predicting the 32nd token after latents
Figure 3. Validation perplexity of the baseline frozen Gemma-2 2B model and augmented models with varying numbers of latents (8, 16, 32, 64), when predicting the stt and 32nd tokens following latent augmentation. Lower perplexity indicates hetter performance.
tion useful for predicting future tokens, effectively enabling the model to perform a form of "latent thinking" before making predictions. This contrasts with standard next-token prediction and allows the coprocessor to learn more gen- eralizable representations. By learning to anticipate future tokens in this manner, the coprocessor develops a stronger understanding of sequential dependencies within text, which proves valuable downstream tasks.
3.1. Perplexity Evaluation
Our augmented Gemma model is able to achieve lower per- plexity on the validation dataset compared to the pretrained Gemma model on many tokens ahead, even beyond the ahead token NA we defined during training. We evaluate this using a proprictary validation dataset (Same as the one used in Gemma (Team-Gemme et al., 2024)) and evaluate the effect of augmenting the frozen Gemma-2 2B LLM's kv-cache on future token prediction. For each sequence in the validation set, we generate N2 latent embeddings after each token using our coprocessor. These embeddings are then used augment the cache each token position. We then measure the model's ability to predict the n-th future token. Specifically, the "1st token" perplexity measures the model's performance predicting the token immediately following the inserted latent embeddings. The "32nd token" perplexity measures the model's performance predicting the token 32 positions ahead, given the context preceding the latent embeddings, the embeddings themselves, and the following 31 tokens. This demonstrates that even though we train with NA 16. the benefits of cache augmenta- tion extend beyond this range, improving predictions even at position 32. Figure 3 presents perplexity curves during training for the baseline frozen Gemma-2 2B model and our augmented models using N, 16, 32, and 64 latent embeddings. Across all latent sizes, our approach consis- tently reduces perplexity, with the improvement scaling with the number of latent embeddings. This demonstrates that augmenting the cache with the coprocessor improves both short-range and longer-range prediction accuracy.
3. Experiments
We validate our approach using the frozen Gemma-2 2B model. Our augmented Gemma-2 models, with only the coprocessor being trained and the decoder-only LLM kept frozen, are trained on the same 2 trillion token, primarily- English dataset used for Gemma-2 pretraining (Team- Gemma et al., 2024), following the setup described in Sec- tion 2.2. This dataset includes variety of sources, such web documents, code, and scientific articles. We trained the model for 100,000 steps using batch size of 1024, packed sequences of length 2048, 16 ahead tokens (). and 128 randomly sampled augmentation positions ("traces"") for all training experiments. Importantly, no task-specific training performed for any of the experiments; all training is done on the pretraining dataset.
The frozen Gemma-2 model serves our primary base- line. We also experimentally verified that continued train- ing of the base Gemma-2 model on this same pretraining data subset results in negligible differences in perplexity or downstream performance, given its comprehensive initial pretraining on this data. Consequently, this was not pursued formal baseline for all reported comparisons.
Table quantifies the perplexity reduction achieved by our cache augmentation approach. The consistent reduction, generally correlating with the number of latents, confirms that the benefit of our method extends multiple subsequent token predictions, suggesting improved internal represen-

---

Deliberation in Latent Space via Differentiable Cache Augmentation
Position
Latents
16 Latents
32 Latents
64 Latents
3.3. Comparison with other baselines and variations
1
-1.53%
-2.48%
-3.28%
-3.94%
3.3.1. PAUSE TOKEN
2
-1.67%
-2.41%
-3.15%
-3.70%
4
-1.39%
-1.98%
-2.66%
-3.17%
We compare our approach with a closely related baseline: the Pause Token method (Goyal et al., 2023). Pause To- ken introduces trainable embeddings inserted between the input (x) and output (y) sequences, encouraging the LLM to perform latent "thinking" before generating the output. The crucial distinction between our approach and Pause Token lies in how these latent embeddings are generated. While Pause Token utilizes fixed and pretrained embeddings that do not condition on the input x. our method employs coprocessor that generates context-dependent, dynamic embeddings based on the input. This allows our approach to tailor the latent representations to the specific input, po- tentially leading to more effective reasoning.
8 -1.22%
-1.56%
-2.11%
-2.61%
16 -0.85%
-1.08%
-1.50%
-1.88%
32 -0.55%
-0.64%
-0.88%
-1.20%
Table Relative perplexity reduction (in %) achieved by aug- mented models compared the baseline, for various numbers of latent embeddings and positions following latent augmentation. "Position" indicates the token position relative the augmentation point. (c.g., Position is the immediately following token).
tations within the decoder, leading to more accurate and coherent generation.
Table directly compares the performance of the baseline Gemma-2 2B model against both Pause Token and our ap- proach, with the latter two using 32 embeddings. Notably, the Pause Token model was trained using the same training data and under the same experimental setup as our method. On the validation set, our method achieves a perplexity of 10.60 on the first token prediction, significantly lower than both the baseline (10.96) and Pause Token (11.63). Furthermore, our method achieves an accuracy of 26. 76% the GSM8K dataset, outperforming both the baseline (21.38%) and Pause Token (22.37%). These improvements underscore the effectiveness of our dynamic, contextually- informed embeddings, which provide a richer representation compared to the fixed embeddings in Pause Token, leading to better next token prediction and improved performance on reasoning tasks.
3.2. Public Benchmark Evaluation
We evaluated cache augmentation on a range of public benchmarks spanning natural language understanding and reasoning tasks (Table 2). In this setting, we only call the coprocessor once, at the end of the prompt. Our method con- sistently improves performance compared to the baseline frozen Gemma-2 2B model, with particularly substantial gains on reasoning-intensin benchmarks. Several tasks, including MMLU, GSM8K, TriviaQA, NQ, and MATH, exhibit a strong correlation between the number of latent embeddings and performance improvement. For example, on GSM8K, accuracy steadily climbs from a +1.29% gain with 4 latent embeddings to a notable +10.05% with 64. Similarly, MATH improves from -0.12% with 4 to +2.06% with 64, and MMLU shows jump from +0.45% with 4 to +4.70% with 64. This trend suggests that for certain challenging reasoning tasks, providing more latent embed- dings allows the model to perform more extensive "thinking" in the latent space, significantly enhancing its reasoning capabilities.
3.3.2. ZERO-SHIT CoT
Our technique be viewed as form of latent Chain- of- ofouhhtt (CoT) prompting. Therefore, we compare our approach to standard zero-shot CoT (Kojima et al., 2022), which involves appending "Lett' think step by step" to the input prompt. While zero-shot CoT can be effective, it relies on the LLM to generate intermediate reasoning steps token by token, which can be computationally expensive during inference. Our method, the other hand, generates latent embeddings in single forward pass, potentially offering more efficient approach to guiding reasoning.
Other reasoning tasks, including ARC-e/c, Winogrande, and Boolq, also show improvements with increasing latent counts. While some tasks, such as AGIEval, BBH, and HumanEval, show less pronounced improvements or oc- casional performance dips with higher latent embedding counts, our method still frequently provides a benefit. This broad improvement across diverse benchmarks underscores the effectiveness and general applicability of cache augmen- tation for enhancing frozen language models.
Table 4 presents the accuracy on GSM8K for the baseline Gemma-2 2B model, zero-shot CoT, and our approach with 16 and 32 latent embeddings. Our method shows clear im- provements. With 16 latent embeddings, accuracy of 24.72%, surpassing both the baseline (21.38%) and zero-shot CoT (23.20%) Performance further improves to 26.76% with 32 embeddings. This suggests that our learned, context-dependent latent embeddings provide a
achieve an
We further provide analysis in the appendix, showing that our method's 's performance scales with increasing training data (Section A. and that effectively adapts to down-
stream tasks (Section A.2).
5

---

Deliberation in Latent Space via Differentiable Cache Augmentation
Benchmark
Metric
Baseline
4 Latents
Latents
16 Latents
32 Latents
64 Latents
MMLU
5-shot
52.00
52.45 (+0.45)
52.24 (+0.24)
52.34 (+0.34)
54.61 (+2.61)
56.70 (+4.70)
GSM8K
8-shot
21.38
22.67 (+1.29)
23.12 (+1.74)
24.72 (+3.34)
26.76 (+5.38)
31.43 (+10.05)
DROP
3-shot, FI
$3.69
54.64 (+0.95)
54.91 (+1.23)
56.23 (+2.55)
57.37 (+3.68)
57.77 (+4.08)
ARC-
O-shot
80.56
81,52 (+0.97)
81.57 (+1.01)
83.12 (+2.57)
83.04 (+2.48)
83.67 (+3,11)
ARC-c
O-shot
50.26
51.28 (+1.02)
52.39 (+2.13)
53.24 (+2.99)
54.44 (+4.18)
54,44 (+4.18)
MATII
4-shot
16.50
16.38 (-0.12)
16.78 (+0.28)
17.00 (+0.50)
17.18 (+0.68)
18.56 (+2.06)
Winogrande
O-shot
64.01
65.35 (+1.34)
65.35
66.30 (+2.29)
66.30 (+2.29)
66.61 (+2.60)
PIQA
O-shot
78.18
78.62 (+0.44)
78.67 (+0.49)
78.94 (+0.76)
78.94 (+0.76)
79.00 (+0.82)
SIQA
O-shot
51.79
51.59 (-0.20)
(-0.5)
51.74 (-0.05)
$2.30 (+0.51)
52.00 (+0.20)
HellaSwag
O-shot
73.77
74.41 (+0.64)
74,41 (+0.64)
74.82 (+1.05)
75.04 (+1.27)
75.31 (+1.54)
Boolq
O-shot
75.41
75.29 (-0.12)
77.22 (+1.80)
78.17 (+2.75)
77.03 (+1.62)
76.91 (+1.50)
MBPP
3-shot
30.40
29.00 (-1.40)
(+1.00))
31.20 (+0.80)
31.40 (+1.00)
31.80 (+1.40)
AGlEval
3-5-shot
31.71
32.18 (+0.47)
30.04
31.32 (-0.38)
32.78 (+1.07)
33.85 (+2.14)
TriviaQA
5-shot
60.29
60.30 (+0.01)
60.83 (+0.54)
61.43 (+1.14)
62.05 (+1.76)
62.23 (+1.94)
NQ
5-shot
17.14
17.35 (+0.21)
17.89 (+0.75)
18.16 (+1.02)
18.91 (+1.77)
19.20 (+2.06)
HumanFval
pass@
19.51
18.29 (-1.22)
19.51 (+0.00)
20.73 (+1.22)
20.73 (+1.22)
22.56 (+3.05)
BBH
3-shot
42.22
42.36 (+0.14)
42.37 (+0.15)
42.53 (+0.31)
42.48 (+0.26)
42.64 (+0.41)
Table Performance of baseline and augmented models across various benchmarks. Results shown for the baseline (frozen Gemma-2 2B pretrained model) and the model augmented with a learned coprocessor using 4, 8, 16, 32, and 64 latent embeddings, respectively. Results reported for zero/few-shot settings as indicated in the "Metric" column. Results accuracy (in %) if not specified otherwise. Improvements over the baseline are shown in parentheses. In this setting, the coprocessor is called once, at the end of the prompt.
Method
Validation set perplexity (L)
GSM8K 8-shot accuracy (1)
32
Baseline
From Scratch Anetune
GSMBK Accuracy
30
Baseline Gemma 2B
10.96
21.38
Pause Token
11.63
22.37
28
Cache Augmentation
10.60
26.76
26-
Table 3. Comparison between the baseline Gemma-2 2B model,
the Pause Token method (Goyal et al., 2023) (using 32 embed- dings), and our approach (also using 32 embeddings). Lower perplexity indicates better next token prediction. Higher accuracy indicates better performance on GSM8K.
24
22-
4
8
16
32
64
Number of latent embeddings
Baseline
0-shot CoT
16 Latents
32 Latents
21.38
23.20
24.72
26.76
Figure 4. Finetuning the coprocessor from Gemma-2 2B pretrained weights significantly improves GSM8K accuracy compared to training from scratch. Lines represent the mean and shaded areas represent the 95% confidence interval, both estimated from the last 5 checkpoints.
Table 4. Accuracy on GSM8K 8-shot for the baseline Gemma-2 2B model, zero-shot Chain-of-Thoughe (CoT) prompting, and our approach with 16 and 32 latent embeddings.
more efficient and effective mechanism for guiding reason- ing compared to the generic prompt and sequential token generation of zero-shot COT.
the primary comparison point for the following experiments.
Training the Coprocessor from scratch: We also inves- tigated training the coprocessor from scratch-randome initializing its weights rather than finetuning from the pretrained weights of Gemma-2 2B. While training from scratch improves performance on all downstream tasks com- pared the baseline. finetuning from pretrained weights yields even better results. This suggests that the coproces- sor benefits from the foundational knowledge encoded in the pretrained LLM. Figure 4 illustrates this improvement in GSM8K accuracy as the number of latent embeddings increases, with the finetuned model consistently outperform- ing the model trained from scratch. Results of other bench- marks can be found in Table in the Appendix Section A.3.
It is worth noting that for some other benchmarks, such as MATH and HumanEval, the performance of zero-shot CoT was closer to the bascline Gemma-2 2B 2B model. GSM8K was highlighted here as it demonstrated a notable improvement with our method over both the baseline and zero-shot CoT.
3.3.3. ALTERNATIVE COPROCESSOR CONFIGURATIONS
We explored alternative configurations for our coprocessor to assess the importance of design choices. Our default setup involves finetuning the pretrained LLM to serve as the coprocessor (i.e., the coprocessor's weights are initialized with the pretrained weights of the LLM), which serves as
6

---

Deliberation in Latent Space via Differentiable Cache Augmentation
LoRA Finetuning the pretrained LLM the Coproces- sor: In addition to full finetuning the pretrained LLM and from-scratch training, we explored the efficacy of Low-Rank Adaptation (LoRA) (Hu et al., 2021) for tuning the copro- cessor from the pretrained L.M'ss weights. LoRA freezes the pretrained model weights and introduces trainable rank- decomposition matrices, significantly reducing the number of trainable parameters. For instance, LoRA with a rank of 64 for the Gemma-2 2B model adds approximately 2% extra trainable parameters compared to the base LLM. This approach offers substantial memory benefits, as only the rel- atively small LoRA weights need to be stored in addition to the base model. We experimented with LoRA using ranks of 64 and 128, comparing their performance on GSM8K to the baseline Gemma-2 2B model, the from-scratch training ap- proach discussed above, and our fully finetuned coprocessor. As shown in Table 5, which presents results using 32 latent embeddings for all methods, LoRA finetuning achieves rea- sonable improvements over the baseline, demonstrating that even a parameter-efficient approach can effectively train the coprocessor for improved reasoning. Specifically, LoRA with rank 64 achieved an accuracy of 23.35% while LoRA with rank 128 reached 24.03%. These results fall between the baseline performance (21.38%) and the performance achieved by from-scratch training (25.78%), indicating that while the LoRA-tunee coprocessor benefits from the pre- trained weights, generating high-quality latent embeddings for effective reasoning appears to require more substan- tial parameter updates than those provided by parameter- efficient methods like LoRA. While these results are not as strong full finetuning (26.76%), they represent notable improvement over the baseline and highlight the potential of LoRA for efficient training/inference of our coprocessor, especially in memory-constrained environments.
Method
GSM8K Accuracy
Baseline
21.38
LoRA (Rank 64)
23.35
LoRA (Rank 128)
24.03
From Scratch Training
25.78
Full Finetuning
26.76
Table S. GSM8K accuracy comparison of different finetuning methods for the coprocessor, all using 32 latent embeddings. LoRA offers a memory mmemory-efficient alternative to full finetuning, achieving reasonable performance gains.
3.4. Impact of the number of ahead token in training
We investigated the impact of varying the number of ahead tokens the number of future tokens the model is trained to predict-during coprocessor training. While larger looka- head improves perplexity on later tokens, it often leads to higher perplexity on earlier tokens. Though learning rate scaling might mitigate this, we empirically chose 16 ahead tokens for most experiments in this paper, given its strong performance on GSM8K, as shown in the Table 6.
4. Related Work
4.1. Chain-of-Thoughe Reasoning in LLMs
Limitations in eliciting complex reasoning from LLMs through standard prompting have motivated research into prompting strategies that encourage intermediate reasoning steps. Chain-of -hoouhhh--httttttthh (CoT) prompting (Wei et al., 2022) significantly improved reasoning performance by prompting L.LMs to "think step by step". Subsequent work explored zero-shot CoT (Kojima et al., 2022; Zhou et al., 2023),
Augmentation using Last Layer's Activations: Instead of using the kv-cache as input to the coprocessor, we ex- perimented with providing the last layer's activations from the frozen LLM, concatenated with the soft token embed- dings. This approach, using 32 latent embeddings, yielded perplexity of 10.81 1 on the validation set and an accuracy of 23.20% on the GSM8K benchmark under the same training setup. Both metrics are notably worse than those achieved with kv-cache augmentation, which resulted in a perplexity of 10.69 and GSMSK accuracy of 26. 76% (also with 32 latent embeddings). We hypothesize that the last layer's 's activations alone do not provide as rich representation for the coprocessor as the information aggregated across multiple layers in the kv-cache, hindering both next-token prediction (reflected in the higher perplexity) and reasoning ability (reflected in the lower GSM8K accuracy).
aggregating multiple reasoning paths (Wang et al., 2022), internalizing the intermediate reasoning steps (Deng et al., 2024), verifying generation steps (Lightman et al., 2023),
and broader search spaces for reasoning trajectories, such as Tree-of- of-Thought (Yao al., 2024: Wang and Zhou, 2024). Other approaches leverage reinforcement learning to opti- mize the reasoning process based on final answer accuracy or target text likelihood (e.g., StaR (Zelikman et al., 2022), TRICE (Hofffman et al., 2024), Quiet-STaR (Zelikman et al., 2024)). While effective, these methods are often constrained by the expressiveness of natural language and can be com- putationally expensive due to the sequential generation of reasoning steps, both during training and inference. On the other hand, our differentiable cache augmentation approach could be viewed as complementary to these RL techniques. For instance, RL could potentially be used to further opti- mize the generation of latent embeddings by the coprocessor, or our method could provide an improved base model upon which RL-based reasoning strategies are applied.
7

---

Deliberation in Latent Space via Differentiable Cache Augmentation
Baseline
4 Ahead
Ahead
16 Ahead
32 Ahead
efficiency of LLMs by identifying and removing redundant information from the input, making it more compact and easier to process.
24.03
24.11
24.72
23.73
21.38
(+2.65)
(+2.73)
(+3.34)
(+2.35)
While our approach also leverages the kv-cache, our moti- vation is fundamentally different. Rather than focusing on compression, we aim to augment the kv-cache with latent embeddings produced by an offline coprocessor, thereby enhancing the transformer's reasoning capabilities with- out modifying its architecture. This allows us to improve the fidelity of further decoding and boost performance on reasoning intensivve tasks. Our work is inspired by the idea of deliberation, where the coprocessor can "think" in the la- tent space by processing the kv-cache and generating mean- ingful embeddings that guide the transformer's subsequent generation.
Table 6. GSM8K accuracy for varying numbers of ahead tokens during coprocessor training. 16 ahead tokens achieves the highest accuracy (24.72%, +3.34% over the baseline of 21 21.38%). 16 latent embeddings are used for all these experiments.
4.2. Latent Space Reasoning
Previous research has investigated the role of latent trans- former computations in LLMs' reasoning abilities. (Cheng and Durme, 2024) elicits reasoning with a short sequence of continuous embeddings, aiming for greater inference effi- ciency. As demonstrated in (Biran et al., 2024), a sequential latent reasoning pathway in LLMS is identified for multi-hop reasoning problems. (Shalev et al., 2024) revealed that the middle layers of LLMs produce highly interpretable embed- dings, representing a set of potential intermediate answers for multi-hop queries. To improve LLMs' latent reason- ing ability, researchers have proposed to augment LLMs with meta tokens. The Pause Token method (Goyal et al., 2023), closely related to our work, introduces trainable em- beddings inserted between input and output sequences to encourage latent "thinking". Similarly, a recent work (Pfau et al., 2024) also studied the cireumstances under which causal transformers are able learn to utilize intermediate dummy tokens (e.g. the filler token used in this work). Un- like these studies which employ pretrained embeddings, our work generates latent tokens dynamically based on the input. More recently, COCONUT (Hao et al., 2024) introduced a new reasoning paradigm by utilizing LLMs' hidden states as input embeddings in latent space. While both our approach and COCONUT explore reasoning in latent space, key dis- tinctions exist. Our method maintains strictly frozen base LLM and generates multiple latent embeddings in paral- lel in a single forward pass of the coprocessor, trained on general pretraining data. In contrast, COCONUT focuses on generating sequential continuous CoT like embeddings and often involves task-specific fine-tuning and multi-stage training to internalize reasoning, without the constraint of frozen base model.
4.4. Augmenting LLMs with External Modules
Extensive research has focused on augmenting pretrained LLMs with external modules for improved efficiency and performance (Pfeiffer et al., 2023). Parameter-efficien fine- tuning methods like prompt tuning (Lester et al.. 2021), prefix tuning (Li and Liang, 2021), and adapters have also been explored. Adapters, first introduced for computer vi- sion by Rebuffi et al. (2017: 2018) and later popularized in NLP by Houlsby et al. (2019), insert small, trainable modules into the LLM. LoRA (Hu et al., 2021) further im- proves adapter efficiency by decomposing weight updates into low-rank matrices. Furthermore, multimodal models like Flamingo (Alayrac et al., 2022), CoCa (Yu et al., 2022), PaLI (Chen et al., 2022), and PaLM-E (Driess et al., 2023) leverage gross-attention or soft prompts to incorporate in- formation from other modalitics. Building upon these tech- niques, recent work has explored augmenting LLMs with modules specifically designed for reasoning. For exam- plc, CALM (Bansal et al., 2024) employs cross-attention between specialized and general models to enhance the gen- eral model's 's capabilities.
4.5. Hypernetworks for Parameter Generation
Our work shares conceptual similarities with the concept of hypernetworks, where a separate network (the hyper- network) generates the parameters of another network. In the context of LLM augmentation, rather than learning a fixed set of parameters for a module, a hypernetwork could generate these parameters conditioned on embeddings rep- resenting different tasks, inputs, or contexts (Ha et al., 2017; Platanios et al., 2018). This allows for form of parame- ter sharing and ""etanglement" between modules that are otherwise disjoint in their parameters (Goyal et al., 2021). Hypernetworks have been used to condition parameter gen- eration on inputs as well. Examples include conditional batch normalization (de Vries et al., 2017). feature-wist
4.3. KV-Cache Compression
KV-cache compression is a technique used to reduce the size of the transformer's kv-cache, the memory storing past activations, for efficient storage and faster computation. (Ge et al., 2024) propose an in-context autoencoder (ICAE) to compress the context into concise embeddings for the kv- cache. (Mu et al., 2024) introduce the concept of "gist to- kens" to learn compressed representations of the kv-cache. Alternatively, prior work (Li et al., 2023) also improves the
8

---

Deliberation in Latent Space via Differentiable Cache Augmentation
linear modulation (FiLM) for text-and-vision tasks (Perez et al., 2018), and self-modulation in GANs (Chen et al., 2019). Bertinetto et al. (2016) even conditioned parameter generation on individual examples for one-shot learning. In the context of L.Mss, hypernetworks have generated diverse module parameters, such as classifier heads (Ponti et al., 2021), continuous prompts (He et al., 2022), and adapter layers (Üstün al.. 2020; Ansell et al., 2021; Mahabadi et al., 2021), conditioned task language embeddings. These embeddings can be learned or fixed, incorporating side information about task or language relationships.
ates all latent embeddings for given augmentation point in parallel in single forward pass for efficiency, investigating the sequential generation of these latent embeddings is an important avenue for future research. Such an approach, where latent embeddings are generated iteratively and can attend to previously generated ones within the same deliber- ation step, could enable more complex and nuanced internal 'thuugh'' sequences and facilitate iterative refinement within the latent space, though it may introduce latency considera- tions. Finally, for efficient implementation of the structured sparse attention masks employed during training, partic- ularly with more complex latent interactions, techniques like FlashAttention (Dao et al., 2022) variants could be explored.
Importantly, our approach can be viewed through the lens of hypernetworks, with the coprocessor itself acting as a hypernetwork that is conditioned on the ky-cache of the frozen LLM. Instead of generating parameters for separate module, the coprocessor generates latent embeddings that augment the kv-cache. However, the core principle remains similar: a network dynamically generating outputs based on a rich contextual input. In our case, the kv-cache serves as a highly informative representation of the input sequence, allowing the coprocessor (hypernetwork) to generate aug- mentations tailored to the specific context.
Acknowledgement
We would like to acknowledge several individuals who con- tributed to this work. We thank Diane Wan, Matt Barnes, and Harrison Lee for engaging in early discussions that helped shape the initial ideas. We are also sincerely grate- ful to Shawn Obanion, Abhinav Rastogi, Bradly Green, Jeremiah Harmsen, and aurclao Ranzato for their in- valuable guidance and support throughout the completion of this work.
5. Conclusion
This paper introduces differentiable cache augmentation, a novel method for enhancing frozen decoder-only language models by incorporating a learned coprocessor that operates on the model's kv-cache. This coprocessor generates latent embeddings that enrich the context provided the L.M,, improving its ability to reason and predict future tokens without requiring any modifications to the original model ar- chitecture. Our experiments demonstrate that this approach consistently reduces perplexity and significantly improves performance on a variety of reasoning- -nntinni-n tasks, even in zero/few-shot settings. These improvements are particu- larly notable on tasks requiring complex reasoning, such as MMLU and GSM8K.
Impact Statement
Improved reasoning abilities in LLMs could lead to various societal consequences, both positive and negative. On the one hand, such models could be beneficial in applications requiring complex groblem-solving, such as education, sci- entific research, and automated decision-making. On the other hand, there are concerns about potential misuse, in- cluding generating harmful content, amplifying biases, and spreading misinformation.
Furthermore, it important to highlight that our approach performs reasoning in latent space, meaning the reasoning trace is not readily interpretable. This lack of transparency may create additional challenges in understanding and de- bugging incorrect outputs generated by the LLM.
Importantly, because the coprocessor operates offline and asynchronously, it opens up exciting possibilities for future research into models that can perform more deliberate and computationally intensive reasoning processes, including deliberation not necessarily conditioned on responding to particular prompt. Future work will explore scaling the coprocessor larger models using many modular processors, investigating different coprocessor architectures, and applying this method to more diverse downstream tasks. This includes investigating configurations where the copro- cessor might be significantly smaller (or larger) than the base LLM, which would require exploring techniques for adapting KV-cache information between models of hetero- geneous sizes, such as projection layers or sub-sampling strategies. Additionally, while our current approach gener-
References
J.-B. Alayrac, J. Donahue, P. Luc, A. Miech, I. Barr, Y. Has- son, K. Lenc, A. Mensch, K. Millican, M. Reynolds, et al. Flamingo: a visual language model for few-shot learning. Advances in Neural Information Processing Systems, 35: 23716-23785 2022.
A. Ansell, E. M. Ponti, J. Pfeiffer, S. Ruder, G. Glavaš, 1. Vulié, and A. Korhonen. MAD-G: Multilingual adapter generation for efficient cross- lingual transfer. In Find- ings of the Association for Computational Linguisties:

---

Deliberation in Latent Space via Differentiable Cache Augmentation
EMNLP 2021, pages 4762-4781, Punta Cana, Domini- can Republic, Nov. 2021. Association for Computational Linguistics. doi: 1w..8663/w12222..nndnnn----- 410. URL https: htt///aaaaatotooooeooooaa org/2021. Findings-emn 41.00.
4-9, 2017, Long Beach, CA, USA, pages 6594
6604, 2017. URL htttps:/pprceeddin
neurips. peurr/ps.cc/pppee//////aaaall Gfad63aaat22eeeeeaaaeeeeeeeeae Abbstaacc.....tt........... html.
R. Bansal, B. Samanta, S. Dalmia, N. Gupta, S. Vashishth, S. Ganapathy, A. Bapna, P. Jain, and P. Talukdar. Lim augmented Ims: Expanding capabilities through compo- sition. arXiv preprint arXiv:2401.0244 102412, 2024.
Y. Deng, Y. Choi, and S. Shicber. From explicit cot to implicit cot: Learning to internalize cot step by step. arXiv preprint arXiv:2405. 14838, 2024.
D. Driess, F. Xia, M. S. Sajjadi, C. Lynch, A. Chowdhery, Ichter, A. Wahid, J. Tompson, Q. Vuong, T. Yu, et al. Palm-e: An embodied multimodal language model. arXiv preprint arXiv:2303.03.777 2023.
L. Bertinetto, J. Henriques, J. Valmadre, P. H. S. Torr, and A. Vedaldi. Learning feed-forwart one-shot learners. In Advances in Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5- 10, 2016, Barcelona, Spain, pages 523-531, 2016. https:::::: http:/tpss//pccceeediiiie /pppceedinggs neurips /c/ htt.ppps....aaaaeee///////sessss 839ab4 09954468220662444460502222888888 html.
T. Ge, H. Jing, L. Wang. X. Wang. S.-Q. Chen, and F. Wei. In-context autoencoder for context compression in a large language model. In The Twelfth International Conference rasttract. on Learning Representations((ILI. 2024.
A. Goyal, A. Lamb, J. Hoffmann, S. Sodhani, S. Levine, Y. Bengio, and B. Schölkopf. Recurrent independent mechanisms. 9th International Conference on Learn- ing Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https: //apenreview.cem
E. Biran, D. Gottesman, S. Yang. M. Geva, and A. Glober- son. Hopping too late: Exploring the limitations of large language models multi-hop queries. In Y. Al-Onaizan, M. Bansal, and Y.-N. Chen, editors, Proceedings of the 2024 Conference on Empirical Methods in Natural Lan- guage Processing, pages 18113-141300 Miami, Florida, USA, Nov. 2024. Association for Computational Linguis- tics. doi: 10. 10.1866331122024emmlpp-.o 7811
S. Goyal, Z. Ji, A. S. Rawat, A. K. Menon, S. Kumar, and V. Nagarajan. Think before you speak: Training language models with pause tokens. arXiv preprint arXiv:2310.02226, 2023.
T. Chen, M. Lučié, N. Houlsby, and S. Gelly. On self modu- lation for generative adversarial networks. In 7th Inter- national Conference on Learning Representations, ICLR 2019, New Orleans, LA,, USA, May 6-9, 2019. OpenRe- view.net, 2019, URL https: https::ttpppeeeeee http:////ppeeeeeeeoooome Forum?id=Hkil 5aaoH5k.....mm.
D. Ha, A. M. Dai, and Q. V. Le. Hypernetworks. In 5th Inter- national Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26. 2017, Conference Track Proceedings. OpenReview.net, 2017. URL https:
/ wwopneeeieeiio net wtt.rrummmiddkppeeeeee1l
X. Chen, X. Wang, S. Changpinyo, A. Piergiovanni, P. Padlewski, D. Salz, S. Goodman, A. Grycner, B. Mustafa, L. Beyer, et al. Pali: A jointly-scaled multilingual language-image model. arXiv preprint arXiv:2209.06777 2022.
S. Hao, S. Sukhbaatar, D. Su, X. Li, Z. Hu, J. Weston, and Y. Tian. Training large language models to reason in continuous latent space, 2024.
Y. He, H. S. Zheng, Y. Tay, J. P. Gupta, Y. Du, V. Aribandi, Z. Zhao, Y. Li, Z. Chen, D. Metzler, H. Cheng, and H. Chi. Hyperprompt: Prompt-based cask-conditioning transformers. International Conference on Ma- chine Learning. ICML 2022, 17-23 July 2022. Balti- more, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pages 8678-8690. PMLR, 2022. URL wwt.httss//prrrcceeinne mll .prees/ vl 2162/heed hwt.:.22eff..e. html.
J. Cheng and B. V. Durme. Compressed chain of thought: Efficient reasoning through dense representations. 2024.
URL httpt:t/t////c... httttt//////xxii/ /ttp:///axvvvvviii org/ htrg//bb//412/ 13171.
T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré. FlashAt- tention:
Fast and memory-efficient exact attention with -w-awrrensss. In Advances in Neural Information Pro- cessing Systems (NeurIPS), 2022.
H. de Vries, F. Strub, J. Mary, H. Larochelle, O. Pietquin, and A. C. Courville. Modulating early visual processing by language. In Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December
D. Hendrycks, C. Burns, S. Kadavath, A. Arora, S. Basart, E. Tang, D. Song, and J. Steinhardt. Measuring mathe- matical problem solving with the math dataset. NeurIPS, 2021.
10

---

Deliberation in Latent Space via Differentiable Cache Augmentation
M. D. Hoffman, D. Phan, D. Dohan, S. Douglas, T. A. Le, A. Parisi, P. Sountsov, C. Sutton, S. Vikram, and R. A Saurous. Training chain-of-thought via latent- variable inference. Advances in Neural Information Pro- cessing Systems, 36, 2024.
J. Mu, X. Li, and N. Goodman. Learning compress prompts with gist tokens. Advances in Neural Information Processing Systems, 36, 2024.
E. Perez, F. Strub, H. de Vries, V. Dumoulin, and A. C. Courville. Film: Visual reasoning with a general conditioning layer. In Proceedings of the Thirty-Second AAAI Conference on Artificial Intelligence, (AAAI- 18), the 30th Innovative Applications of Artificial Intelligence (IAAI-88), and the 8th AAAI Symposium Educational Advances in Artificial Intelligence (EAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018, pages 3942-3951. AAAI Press, 2018. URL https: /twww//wwww... awwi. wttp:/aaaaa...occ/////aooooo//// /wnde////////xexeooo//ea//x///xe php/ TEAA///AAA1//pppee///eee
N. Houlsby, A. Giurgiu, S. Jastrzebski, B. Morrone, Q. de Laroussilhe, A. Gesmundo, M. Attariyan, and S. Gelly. Parameter-efficient transfer learning for NLP. In K. Chaudhuri and R. Salakhutdinov, editors, Pro- ceedings of the 36th International Conference on Ma- chine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 2790-2799. PMLR, 2019. URL ht http:://t/// http://tttoo pttpp/eecceeeececeepeepepcec.sse mal.prress v97 httphhoolsllllly 9a. html.
J. Pfau, W. Merrill, and S. R. Bowman. Let's think dot by dot: Hidden computation in transformer language models. In First Conference Language Modeling, 2024.
E. J. Hu, Y. Shen, Wallis, 7., Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen. Lora: Low-rank adaptation of large language models. arXiv preprint S.Ki:2206600688. 2021.
J. Pfeiffer, S. Ruder, I. Vulic, and E. M. Ponti. Modular deep learning. Transactions of Machine Learning Re- search, 2023, 2023. URL https: /apes///oppeeeeee. net/forunt 1d=222xXExxttt......
T. Kojima, S. S. Gu, M. Reid, Y. Matsuo, and Y. Iwa- sawa. Large language models are zero-shot reasoners. Advances in neural information processing systems, 35: 22199-22213, 2022.
E. A. Platanios, M. Sachan, G. Neubig, and T. Mitchell. Contextual parameter generation for universal neural ma- chine translation. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 425-435, Brussels, Belgium, Oct.-Nov. 2018. Asso- ciation for Computational Linguistics. doi: 10.18653/v1/ D18- 103-. URL, https::/aaeee https://aclannhoologyy.o org/ D18 1039.
B. Lester, R. Al-Rfou, and N. Constant. The power of scale for parameter-efffiient prompt tuning. arXiv preprint wrXiv:2/04.0861,, 2021.
X. Li and Liang. Prefix-tuning: Optimizing continuous prompts for generation. arXiv preprint 2rX1v:2101.00190, 2021.
Y. Li, B. Dong, F. Guerin, and C. Lin. Compressing context to enhance inference efficiency of large language models. In H. Bouamor, J. Pino, and K. Bali, editors, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 6342-6353, Singapore, Dec. 2023. Association for Computational Linguistics. doi:
E. M. Ponti, I. Vulié, R. Cotterell, M. Parovic, R. Reichart, and A. Korhonen. Parameter space factorization for zero- shot learning across tasks and languages. Transactions of the Association for Computational Linguistics, 9:410- 428, 2021. doi: 10.1162/tuc.a.000774 URL https: / //aclanthology wttp//2211.0000 tac1-1.25.
10. 10.18665/v1/2023.emmm 20tp5s///122233mmpmmma.m.33333 wmt.:3mmmmmmmmmmi.391111
H. Lightman, V. Kosaraju, Y. Burda, H. Edwards, B. Baker, T. Lee, J. Leike, J. Schulman, I. Sutskever, and K. Cobbe. Let's verify step by step. arXiv preprint arXiv:2305.20055 2023.
S. Rebuffi, H. Bilen, and A. Vedaldi. Learning multiple visual domains with residual adapters. In Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 506-516, 2017. URL https: http:::://p.ooceiing..
R. K. Mahabadi, S. Ruder, M. Dehghani, and J. Hen- derson. Parameter-efficient multi-task fine-tuning for transformers via shared hypernetworks. In Proceed- ings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th Interna- tional Joint Conference Natural Language Process- ing (Volume 1: Long Papers), pages 565-576, Online, Aug. 2021. Association for Computational Linguistics. URL https://
neurips cc cttccc/ppppppeeeccccccccccccecc 2atppper/2011/hhaee e7b2 9bbe999c - Abstract html.
S. Rebuffi, H. Bilen, and A. Vedaldi. Efficient parametriza- tion of multi-domain deep neural networks. In 2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018, pages 8119-8127.
doi: 10.116533/1/2021..aa acl- acl bclanthology www.gyy1.gg.00..0 wttp22211.1111110 all.lllo.....11...11.1m

---

Deliberation in Latent Space via Differentiable Cache Augmentation
Computer Vision Foundation / IEEE Computer So- ciety, 2018. doi: 10.1109/CVPR.22088.004.1 URL
E. Zelikman, Y. Wu, Mu, and N. Goodman. Star: Boot- strapping reasoning with reasoning. Advances Neural Information Processing Systems, 35:15476-1548, 2022.
http::///peeaccesss.heeeeeeo.oo cot.t//teocccccccccccc/c.ac content_
evpr_2018 / http:///i///ttee/eeuuff__BffCfCC
E. Zelikman, G. Harik, Y. Shao, Jayasiri, N. Haber, and N. D. Goodman. Quiet-star: Language models can teach themselves think before speaking. arXiv preprint arXiv:2403.0963 099.99999 2024.
2018_paperr.
html.
T. Schuster, A. Fisch, J. Gupta, M. Dehghani, D. Bahri, V. Tran, Y. Tay, and Metzler. Confident adaptive language modeling. Advances in Neural Information Processing Systems, 35:17456-1744 17472, 2022.
Zhou, N. Schärli, L.. Hou, J. Wei, Scales, X. Wang, D. Schuurmans, C. Cui, O. Bousquet, Q. V. Le, and E. H. Chi. Least-to-most prompting enables complex reasoning in large language models. In The Eleventh International Conference on Learning Representations, 2023.
Y. Shalev, A. Feder, and A. Goldstein. Distributional rea- soning in llms: Parallel reasoning processes in multi-hop reasoning. arXiv preprint arXiv:2406.
2024.
Team-Gemma, M. Riviere, S. Pathak, P. G. Sessa, C. Hardin, S. Bhupatiraju, L. Hussenot, T. Mesnard, B. Shahriari, A. Ramé, al. Gemma 2: Improving open language mod- els a practical size. arXiv preprint arXiv:2408.000 00118, 2024.
A. Üstün, A. Bisazza, G. Bouma, and G. Noord. UDapter: Language adaptation for truly Universal De- pendency parsing. In Proceedings of the 2020 Con- ference on Empirical Methods in Natural Language Processing (EMNLP), pages 2302-2315, Online, Nov. 2020. Association for Computational Linguistics. doi: 10. 10.116633/1122020.mmmmmmmmm main. 180. URL https:// aclanthology org/2020. wttp:///2222e...mmmmammemm.mema. emnlp-main 180.
X. Wang and D. Zhou. ohain-of-thought reasoning without prompting. arXiv preprint www.arrivvvv...... 10200, 2024.
X. Wang, J. Wei, D. Schuurmans, Q. Le. E. Chi, S. Narang, A. Chowdhery, and D. Zhou. Self-consistency improves chain of thought reasoning in language models. arXiv preprint arXiv:2203.11 1/171, 2022.
Wei, X. Wang, D. Schuurmans, M. Bosma, Xia, E. Chi, Q. V. Le, D. Zhou, al. Chain-of-thought prompting clicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24837, 2022.
Y. Wu, Z. Sun, S. Li, S. Welleck, and Y. Yang. Inference scal- ing laws: An empirical analysis compute-optima in- ference problem-solvine with language models, arXiv preprint 4rXiv:2408 00724,, 2024.
S. Yao, D. Yu, J. Zhao, I. Shafran, T. Griffiths, Y. Cao, and K. Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. Advances in Neural Information Processing Systems, 36, 2024.
J. Yu, Z. Wang, V. Vasudevan, L. Yeung, M. Seyedhos- seini, and Y. Wu. Coca: Contrastive captioners are image-text foundation models. arxiv 2022. arXiv preprint arXiv:2205.019/7. 2022.
12

---

Deliberation in Latent Space via Differentiable Cache Augmentation
A. Appendix
A.1. Scaling with Training Data
The scaling of performance with increasing training data is a crucial aspect of evaluating the effectiveness of our approach. Figure demonstrates the impact of training duration on both GSM8K accuracy and validation perplexity for our method. The x-axis represents the total number of training steps for the coprocessor. The baseline performance, representing the frozen Gemma-2 2B model, is shown for reference corresponding intervals along this axis. As shown, we observe clear trend of improved performance for our method with increased training data (i.e., more tokens seen during training of the coprocessor). Specifically, our method ("Ours") demonstrates a clear benefit from increased training exposure, with GSM8K accuracy exhibiting a consistent upward trend and validation perplexity showing a decreasing trend. This indicates that the coprocessor learns to generate more useful latent embeddings and better integrate with the frozen LLM as it is exposed to more data, improving next token prediction. This trend highlights the importance of scaling with training data for our approach.
27
Baseline
10.85
Ours
GSMBK Accuracy
26
Validation Perplexity
10.80
25
10.75
24
10.70
23
10.65
22-
10.60
20000
40000
60000
80000
100000
20000
40000
60000
80000
100000
Percentage of Total Training Steps
Training steps
(a) GSMBK
(b) Validation Set
Figure 5. Scaling of GSM8K accuracy and validation perplexity with increasing training steps for the coprocessor (using 32 latent embeddings). The baseline performance of the frozen Gemma-2 2B model is shown for reference.
A.2. Adaptation to Downstream Tasks
All experiments described thus far have focused training the processor using the pretraining dataset. To assess the adaptability of our approach to downstream tasks, we conducted experiments us- ing data mixture containing the training sets of the GSM8K and MATH (Hendrycks et al., 2021) datasets. We employed LoRA finetun- ing (with rank of 128) on both the baseline model and our augmented model. For the baseline, LoRA was applied directly the base LLM,, while for augmented model, LoRA applied specifically to the coprocessor, leaving the base LLM frozen.
Baseline
35
Ours
GSM8K Accuracy
Figure 6 presents the results of this downstream adaptation. We ob- serve a substantial improvement in performance for our augmented model compared to the baseline after LoRA finetuning. This improve- ment is likely attributable to the strong regularization imposed by keeping the base LLM frozen during coprocessor training. This freez- ing prevents overfitting to the relatively small downstream datasets, allowing the coprocessor to effectively learn task-specific reasoning patterns without disrupting the general knowledge encoded in the pre- trained LLM. The baseline model, with LoRA applied directly to the LLM, likely suffers from overfitting to the downstream data, limiting its performance gains. These results demonstrate the effectiveness of our approach in adapting to downstream tasks while maintaining the robustness of the pretrained L.M..
151
1000
2000
3000
4000
5000
Training steps
Figure 6. Accuracy on GSM8K's s test set after LoRA finetuning. Our augmented model shows significant improvement compared to the baseline.
13

---

Deliberation in Latent Space via Differentiable Cache Augmentation
A.3. Training Coprocessor from Scratch
We observed performance gains across most benchmarks when training the coprocessor from scratch (with randomly initialized weights), but finetuning from the pretrained LLM consistently yielded better results.
Benchmark
Metric
Baseline
4 Latents
8 Latents
16 Latents
32 Latents
64 Latents
MMLU
5-shot
52.00
$2.03 (+0.03)
52.21 (+0.21)
52.75 (+0.75)
53.55 (+1.S5)
$6.63 (+4.63)
GSM8K
8-shot
21.38
22.52 (+1.14)
22.59 (+1.21)
24.41 (+3.03)
25.78 (+4.40)
(+8.422
ARC-c
O-shot
80.56
81.69 (+1.13)
(+1.30)
82.79 (+2.23)
83.12 (+2.56)
83.21 (+2.65)
ARC-c
0-shot
50.26
51.71 (+1.45)
(+1.96)
52.47 (+2.21)
54.27 (+4.01)
53.24 (+2,98)
MATH
4-shot
16.50
16.22 (-0.28)
(-0.04)
16.92 (+0.42)
17.18 (+0.68)
18.34 (+1.84)
Winogrande
0-shot
64.01
65.19 (+1.18)
65.98 (+1.97)
66.54 (+2.53)
66.69 (+2.68)
67.25 (+3.24)
PIQA
O-shot
78. 18
78.13 (-0.05)
79.00 (+0.82)
79. (+0.98)
79.27 (+1.09)
79.22 (+1.04)
SIQA
O-shot
51.79
51,94 (+0.15)
(-0.15)
51.84 (+0.05)
51.94 (+0.15)
51.89 (+0.10)
HellaSwag
O-shot
73.77
74.37 (+0.60)
74.68 (+0.91)
74.82 (+1.05)
74,89 (+1.12)
(+1.44)
Boolg
0-shot
75.41
75.66 (+0.25)
(+1.53)
76.97 (+1.56)
77,80 (+2.39)
77.46 (+2.05)
MBPP
3-shot
30.40
30.40 (0.00)
30.60 (+0.20)
(+0.40)
32,00 (+1.60)
32.60 (+2.20)
AGIEval
3-5-shot
31.71
32.52 (+0.81)
32.22 (+0.51)
31.92 (+0.21)
32.78 (+1.07)
32.35 (+0.64)
TriviaQA
5-shot
60.29
60.53 (+0.24)
60.95 (40.66)
61.45 (+1.16)
61.93 (+1.64)
62.62 (42.33)
NQ
5-shot
17.14
17.26 (+0.12)
17.89 (40.75)
18.47 (+1.33)
18.68 (+1.54)
19.00 (+1.8)
Fvaa
pass @ 19.51
18.29 (-1.22)
18.90 (-0.61)
20.73 (+1.22)
19.51 (0.00)
19.51 (0.00)
BBII
3-shot
42.22
42.16 (-0.06)
42.24 (+0.02)
42.42 (+0.20)
43.19 (+0.97)
42.93 (+0.71)
Table Performance of baseline and augmented models across various benchmarks with coprocessor training from scratch. Check Table 2 for more detailed description.
A.4. 9B Model Experiments
We also experimented with scaling up our approach to larger LM, Gemma-2 9B. Due to limited compute resources, trained the 9B model with the same batch size and training steps as the 2B model. Despite this constraint, we observed meaningful improvements on many downstream tasks. For example, on GSM8K 8-shot, the accuracy improved from 64.7% to 69.4 4.4, and on MMLU, it increased from 71.3% to 74.5%. These results suggest that our approach can effectively enhance the reasoning capabilities of larger LMs, even with limited compute resources for training the coprocessor.
A.5. Asynchronous Coprocessor
To assess the practicality of deploying our coprocessor in real-world setting, we evaluated its performance when operating asynchronously alongside the base LLM. This configuration can hide the extra latency introduced from cache augmentation, as response sampling won't be blocked by the coprocessor's operation.
We simulated the case where the LLM performs N extra sampling steps (ranging from to 3) before receiving the soft embedding from the asynchronous coprocessor. Table 8 presents the results GSMSK, showing slight degradation in accuracy as N increases.
This performance drop is likely because our training procedure did not explicitly account for this asynchronous behavior. However, it's important to note that even with this degradation, the performance remains significantly better than the baseline frozen Gemma 2B model. This experiment demonstrates the potential for hiding the coprocessor's latency through asynchronous operation, paving the way for more efficient and practical deployment.
Baseline
No Asyne
N=1
N=2
N=3
GSM8K Accuracy
21.38
31.43
27.2
27.1
26.9
Table 8. Performance of asynchronous coprocessor on GSM8K.
14


Reviews:
[{"review_id": "9Glf09FQST", "reviewer": "Reviewer_n6p1", "summary": "In this work a co-processor is trained to get as input the generated KV-cache of a frozen model - after given an input x  and a set of soft tokens and produce a set of latent embeddings z. These embeddings are appended to the KV-cache (augmentation) and the original frozen model decodes towards output y (generation with augmented context). \n\nThe co-processor is trained by allowing it to first generate a number of latent representation and only then produce a predetermined number of output tokens for a training loss to be computed (starting from predetermined positions). Thus augmented models exhibit lower perplexity and increased performance across a broad selection of tasks, which typically increases when allowing the model to generate more latent embeddings. \n\nThis method, which allows asynchronous operations and sports end-to-end differentiability, compares favorably to closely or remotely similar schemes (Pause Token, zero-shot COT). The co-processor is typically initialized as a twin model but can also be trained from scratch or through LoRA (and the latter although incurring a performance hit is a very attractive option).", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer n6p1,\n\nThank you very much for your positive evaluation and your insightful review of our work. We greatly appreciate your accurate summary, positive feedback on our claims and experiments, and for reviewing the supplementary material thoroughly. We are encouraged that you see this as a timely contribution.\n\nWe would like to address your specific comments and question:\n\n### **Regarding Inference Overhead (Time/Space):**\n\nYou suggested including comments on the overhead related to generating embeddings and using two models at inference time.\n\n**Our Response:**\n> Thank you for raising this important point about inference overhead. Regarding space, while our main setup uses a coprocessor similar in size to the base model, we also demonstrated strong results with a parameter-efficient LoRA variant adding only ~2% extra parameters (as discussed in the supplement). Regarding time, our method adds only one extra forward pass through the coprocessor during inference. Crucially, this pass can operate asynchronously, potentially hiding much of its latency, as discussed regarding deployment scenarios. Furthermore, KV cache reuse limits the computational overhead of this pass. We will add a note discussing these inference time and space considerations more explicitly in the revised manuscript.\n\n### **Regarding Comparison with Coconut:**\n\nYou suggested a more detailed comparison with the Coconut paper.\n\n**Our Response:**\n> Thank you for the suggestion to expand the comparison with Coconut. While both works explore continuous representations, key differences make direct comparison challenging. Coconut focuses on training an LLM to generate sequential 'continuous CoT' embeddings, often fine-tuning on specific downstream tasks. In contrast, our method uses a separate coprocessor to generate *parallel* latent embeddings that augment a *frozen* base LLM, and our coprocessor is trained on general pretraining data, aiming for broad applicability without task-specific tuning. Given these methodological distinctions and the near-simultaneous appearance of the works (Dec 2024), a direct empirical comparison is difficult. However, we will expand the discussion in the related work section to better highlight these conceptual similarities and differences as requested.\n\n### **Regarding Relation to RL Techniques for Reasoning:**\n\nYou asked how our results/method for math reasoning compare to recent RL techniques and if there's potential synergy.\n\n**Our Response:**\n> Thank you for this insightful question regarding RL techniques for reasoning. Recent RL methods, particularly those using process supervision or outcome rewards, have shown impressive results, often by fine-tuning models to generate better reasoning steps or verify solutions. Our approach is largely orthogonal, focusing on enhancing the base model's capabilities through parallel 'latent deliberation' via the coprocessor, trained on general pretraining data without task-specific reward signals. However, we see potential for synergy: RL techniques could potentially be applied *on top* of our augmented model to further refine the final reasoning output or even guide the generation process. Alternatively, RL could be explored as a method to train the coprocessor itself, though this adds significant complexity. We view these as complementary approaches rather than competing ones.\n\n---\n\nThank you again for your valuable feedback and positive assessment of our work. We believe addressing these points will further strengthen the paper.\n\nSincerely,\n\nThe Authors of Submission 403", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "Most relevant: further explorations of latent space; Coconut and Pause Token papers are close along this dimension.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "This is a very timely contribution because decoding towards tokens could be seen as a kind of quantization, so exploring the continuous latent space more could be very advantageous towards finally producing tokens of \"higher\" quality.\n\nTraining seems more complex, there are more and different hyper-parameters to explore (positions, how many latents, how many tokens to generate). This will translate to training overhead.", "other_comments_or_suggestions": "- It would be nice to include a comment on the overhead of generating more embeddings (time) or on the fact that there are two, rather than one models in use (space)  - at inference time.\n\n- Coconut is also a recent similar approach in the sense that the continuous latent space is explored more before the materialization of tokens. The reader would expect a more detailed comparison - basically as started towards the end of page 7.", "questions_for_authors": "Could you position the results obtained here in math reasoning tasks and the training method used to reach those as compared to recently re-popularized RL techniques that attain impressive results for this type of tasks? Is there room for synergy? Or these views are totally orthogonal?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tExb5u8oWR", "reviewer": "Reviewer_w2nr", "summary": "In this work, the authors train a hyper-network (termed “coprocessor”) which takes a KV-cache from a language model mid-generation and produces a set latent embeddings which are appended to the KV-cache before producing the final answer. Critically, the coprocessor produces the latent embeddings with *a single forward pass* (unlike CoT decoding which requires decoding tokens step-by-step).\n\nThe main claims of this work are:\n\n1. Cache augmentation enables lower perplexity on subsequent tokens.\n2. When the cache is augmented the decoder achieves improved performance compared to CoT reasoning on a range of reasoning-intensive tasks. \n3. Using the coprocessor is more efficient than using CoT decoding because CoT decoding is sequential and the coprocessor operates in parallel.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer w2nr,\n\nThank you for the detailed review of Submission 403 and the constructive feedback. We're glad you found the method novel and the related work discussion good. We address your points below:\n\n### **Regarding Claim 1 (Perplexity Reduction & Baseline Fairness):**\n\nYou questioned the baseline fairness due to training data and suggested a fine-tuning comparison.\n\n**Our Response:**\n> We thank the reviewer for their detailed feedback and insightful question regarding the training data and baseline comparisons. We recognize the importance of providing more clarity here.\n\n> The dataset used to train the coprocessor is *a subset of the data the model was pretrained on*. We did perform the suggested experiment of further finetuning the base model on this same dataset subset. However, as the base model was already extensively optimized on this data distribution during its initial pretraining, we observed negligible differences in perplexity or downstream performance. Standard fine-tuning paradigms predominantly yielded diminishing returns when applied directly to the original pretraining corpus. Furthermore, consistent with typical large-scale pretraining data, the corpus we used consists of text for self-supervised next-token prediction and lacks curated instruction-following examples or specific question-answering formats as in downstream tasks.\n\n> Given that this fine-tuning baseline showed no significant change from the original frozen model, we initially omitted it to maintain focus on comparisons highlighting the architectural contribution of the coprocessor. However, we understand the reviewer's point about baseline completeness. We will update the manuscript (e.g., in the dataset description or appendix) to clarify the nature of this dataset (as part of the original pretraining corpus, focused on self-supervision) and explicitly state the result of the fine-tuning baseline experiment, explaining why it yielded minimal change.\n\n### **Regarding Claim 2 (Performance Improvement vs. CoT & Evaluation Scope):**\n\nYou noted the limited evaluation scope (GSM-8K only) for the CoT comparison.\n\n**Our Response:**\n> We appreciate this point. Our primary goal wasn't necessarily to universally outperform CoT, but rather to investigate parallel 'latent deliberation'. We compared against zero-shot CoT ('Let’s think step by step') because it enhances reasoning at inference time without task-specific demonstrations or fine-tuning, making it a comparable low-data approach similar to our use of general pretraining data. In our experiments, this prompt yielded near-baseline results on MATH/HumanEval; we highlighted GSM-8K because zero-shot CoT provided a notable improvement there, thus offering a meaningful comparison point against our method. We will add discussion clarifying this rationale, explicitly mention the MATH/HumanEval CoT results for completeness, and refine claims accordingly in the paper.\n\n### **Regarding Claim 3 (Efficiency Benefits vs. CoT & Hardware Considerations):**\n\nYou questioned the efficiency claim due to lack of empirical evidence considering hardware limitations.\n\n**Our Response:**\n> Thank you for raising these critical efficiency points regarding I/O/compute bounds. Your concerns are valid under the assumption of shared hardware. However, our proposed architecture does not impose this constraint. The asynchronous coprocessor performs its single forward pass (akin to prefill) and can be deployed on **separate hardware** (e.g., a different accelerator or node). This fundamentally changes the efficiency analysis by decoupling its execution from the base model's subsequent token-by-token generation. This mitigates resource contention on the decoder's device and enables potential efficiency gains not possible if constrained to the same hardware.\n\n### **Regarding your Question (Limitations/Tradeoffs of Parallel vs. Sequential Computation):**\n\nYou asked about the limitations of the parallel approach vs. sequential CoT.\n\n**Our Response:**\n> Thank you for this relevant point on CoT's sequential nature and the cited work. Our work takes a first step in exploring parallel latent deliberation. We agree it's plausible that tasks heavily reliant on iterative refinement, long causal dependencies, or where intermediate steps explicitly build upon each other might benefit more from sequential CoT. Our current parallel coprocessor doesn't explicitly model such sequential dependencies within the latent injection. This represents a potential limitation compared to sequential methods, which we will acknowledge and discuss in the revised paper. Investigating sequential latent deliberation is an important avenue for future work.\n\nWe hope this addresses your concerns and thank you again for the valuable feedback to improve the paper.\n\nSincerely,\n\nThe Authors of Submission 403", "claims_and_evidence": "Are the claims made in the submission supported by clear and convincing evidence? If not, which claims are problematic and why?\n\n1. If the claim is about the comparison with a frozen model, this claim is well-supported by evidence (see Figure 3). However, I don’t think this is a fair baseline. The coprocessor saw over 200 billion tokens of training data from a “proprietary pretraining dataset” before evaluation on that same proprietary dataset. In contrast, the baseline did not see any of that data. Without any further information on that pretraining dataset, it is not clear to me whether the reduction in perplexity is due to the proposed methodology (i.e. the coprocessor), or just the fact that we’re training on from the same distribution as the test data. To be convinced that the benefits come from the specifics of having a coprocessor, I’d like to see how perplexity improves when we simply fine-tune the base model on the same pretraining data. \n2. This claim is not very well-supported by evidence. Comparisons with CoT are only made on GSM-8K, a single reasoning task. To support the claim that cache augmentation actually improves performance w.r.t. CoT I would expect evaluations on more datasets. \n3. This idea is not supported by strong empirical or theoretical evidence in the paper. The authors argue in the introduction and abstract that cache augmentation has efficiency benefits over CoT because the coprocessor can be run asynchronously while the model continues to generate. However, it’s not obvious to me that this could actual yield latency or throughput improvements in practice. At small batch sizes, language model generation is typically I/O-bound (the bottleneck is loading weights from HBM to the cores, not performing the compute). Since the coprocessor has its own weights, loading those weights will contend for bandwidth with the generation. As a result, the coprocessor and generator cannot fully overlap. At high batch sizes, we are typically compute bound (the cores are “full”), so adding more parallelism won’t help significantly. To convince me that the method provides speedup, I would expect to see empirical speed benchmarks or at least a theoretical cost model. Otherwise, I think this claim should be downplayed in the abstract and introduction.", "methods_and_evaluation_criteria": "Please see discussion of claims and evidence above.", "theoretical_claims": "Not applicable.", "experimental_designs_or_analyses": "Yes I did. Please see the discussion of claims and evidence above where I discuss the experimental design and analyses in the context of the paper’s claims.", "relation_to_broader_scientific_literature": "In my view, the paper is most closely related to the hypernetwork literature. The authors do a great job situating the work in the context of the broader literature in the related work.", "essential_references_not_discussed": "None that I’m aware of. The authors provide a nice review of the relevant literature.", "other_strengths_and_weaknesses": "The method proposed is interesting and to my knowledge novel. However, due to some weaknesses in the evidence discussed above, it’s not clear to me how much of the benefit comes from the proposed method vs. the choice of specific training data.", "other_comments_or_suggestions": "None.", "questions_for_authors": "Several works (e.g. https://arxiv.org/abs/2310.07923) have shown that the power of chain of thought actually comes specifically from the sequential nature of the computation. I suspect that on some tasks which require sequential computation (e.g. S5), CoT could outperform the parallel coprocessor method. Have you found any limitations or tradeoffs when it comes to the parallel nature of the coprocessor?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "g4lPIiFPIi", "reviewer": "Reviewer_9maT", "summary": "This paper proposes a novel method that augment the memories (kv-caches) with a set of latent embeddings from auxiliary compressor modules. This offers two main advantages, end-to-end differentiability by using soft (continuous) tokens and asynchronous operation by using compressor in offline while freezing the base model. Without fine-tuning on downstream tasks, they showed good performance improvement across diverse tasks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer 9maT,\n\nThank you for the positive evaluation, thoughtful review, and examination of the supplement. We appreciate your constructive questions and address them below:\n\n### **Regarding Training Speed/Process:**\n\nYou asked about the training speed and the forward passes involved.\n\n**Our Response:**\n> We thank the reviewer for the question regarding training speed. The training involves three main forward passes as outlined in Figures 1 and 2: (1) LLM processes input text to generate KV caches. (2) The coprocessor takes these KV caches to generate latent embeddings in parallel for all insertion points using an attention mask (avoiding extra forward steps per sample). (3) Latent embeddings are combined with the original text KV caches in the LLM to predict ahead tokens for loss calculation. Due to KV cache reuse, the effective sequence length involves the original length plus terms related to the number of latent and ahead tokens per insertion point. We will ensure this process is clearly described.\n\n### **Regarding Asynchronous Inference Elaboration:**\n\nYou asked for elaboration on the asynchronous inference process and the placement of latent embeddings when extra sampling steps occur.\n\n**Our Response:**\n> Thank you for this important question about asynchronous inference. To clarify the process: If, for example, the base model performs N=2 asynchronous sampling steps immediately after the prompt while the coprocessor computes, the coprocessor uses the KV cache of the *original prompt only* (ignoring the N async steps) to generate latent embeddings in a single pass. These generated latents are then inserted *between* the original prompt's KV cache and the KV cache of the N=2 asynchronously sampled tokens before resuming further generation. We will add a detailed elaboration of this process to Appendix A.5 as requested.\n\n### **Regarding Sequential Latent Embedding Generation:**\n\nYou asked if we tried outputting latent embeddings sequentially.\n\n**Our Response:**\n> We thank the reviewer for asking about sequential latent generation. We did not pursue this approach in the current work primarily due to the significant latency costs associated with sequential generation during both training and inference compared to our parallel method. However, we agree that exploring sequential latent deliberation is an interesting direction for future research, potentially offering different trade-offs.\n\n### **Regarding Scalability and Heterogeneous Model Sizes:**\n\nYou raised a valid point about the potential disadvantage of the coprocessor needing to be the same size as the base model and asked for discussion on heterogeneous sizes.\n\n**Our Response:**\n> Adapting to different coprocessor/base model sizes (requiring KV cache adaptation like sub-sampling) needs further research. We acknowledge this limitation and will add discussion on scalability considerations for heterogeneous sizes.\n\n### **Regarding Multi-step Latent Generation:**\n\nYou asked about generating and using latent embeddings multiple times during generation.\n\n**Our Response:**\n> Our training uses single-step parallel latent generation for efficiency, as multi-step dependencies increase training compute. Multi-step generation (where latents attend to prior latents) is compelling future work, especially for dialogue, though evaluating it requires infra changes beyond this rebuttal. We will discuss this possibility as future work.\n\n### **Regarding results of 1-ahead + Cache Augmentation:**\n\nYou asked about the results when training with 1-ahead token prediction combined with cache augmentation.\n\n**Our Response:**\n> Thank you for this specific question regarding Table 6. When training with only 1 ahead token combined with cache augmentation, we observed only negligible improvements over the baseline model. \n\n### **Regarding Suggested References:**\n\nYou suggested citing concurrent works like CCOT and SoftCoT.\n\n**Our Response:**\n> Thank you for suggesting CCOT and SoftCoT. We agree they are relevant and will cite/discuss them appropriately in the revised related work section.\n\n### **Regarding Efficient Parallel Decoding Implementation:**\n\nYou had a further question on implementing the parallel decoding efficiently, specifically regarding FlashAttention.\n\n**Our Response:**\n\n> Thank you for the question about efficient implementation. While we use standard attention optimizations, specialized techniques like FlashMask offer a path for future efficiency gains. FlashMask is suited for our method's sparse non-causal masks, potentially improving speed/reducing overhead by skipping computation for masked connections (e.g., latent-to-latent). \n---\n\nWe hope these responses adequately address your questions and comments. Thank you again for your insightful feedback and positive evaluation, which will help us improve the paper significantly.\n\nSincerely,\n\nThe Authors of Submission 403", "claims_and_evidence": "The claims and their motivations make sense a lot. Especially, using extra modules, which can be used asynchronously in potential, can efficiently improve the performance through ‘latent thinking’.", "methods_and_evaluation_criteria": "Methods are well aligned to their motivations. Yet, I have some following concerns and questions:\n\n- How about the training speed for these methods? KV-caches from frozen base models are reused in coprocessor, while coprocessor is being updated. So, it seems like we have to run additional forward step with frozen LLM for every samples. Though there could be advantages like no need to update some parameters related to KV caches, or reduced memory footprint for running coprocessor.\n- The authors suggested these methods can be utilized asynchronously, but it would be good to elaborate inference process for it in Section A.5. If there are extra sampling steps in base model, which positions do we put latent embeddings?\n- Have you tried to output latent embeddings sequentially, not in parallel?\n- One disadvantage could be coprocessor should be the same size with base model due to reusing KV caches and augmenting soft embeddings. As this point can make these methods not to be scalable, further discussions for heterogeneous model sizes would be great.\n- Is there any result that generate and use latent embeddings for multiple times during generation, not just a single forward call like in main experiments. I guess we have to adjust attention masks if we generate additional latent embeddings as latent embeddings do not attend each other in training settings. This could induce additional overhead during inference.\n- What were the results when you train the models with 1 ahead token and cache augmentation techniques (regarding to Table 6 results)?\n\nI cannot see any problems in evaluation criteria.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "- Analysis on computational costs would make these methods more validate.", "relation_to_broader_scientific_literature": "I believe this work has been related to recent latent reasoning research. Especially, this work can be classified to latent reasoning methods that use extra modules to augment thought tokens. CCOT and SoftCoT (see below references) seem well aligned with this paper’s research direction.", "essential_references_not_discussed": "Like above explanations, there are some concurrent works that use extra modules (LoRA or assistant LLM) to generate thought tokens. I’d just suggest the authors to cite these references in their paper.\n\n[1] Compressed Chain of Thought: Efficient Reasoning Through Dense Representations\n\n[2] SoftCoT: Soft Chain-of-Thought for Efficient Reasoning with LLMs", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "I’d be glad to discuss with the authors to raise my evaluation score. Please refer to some questions above.\n\nOne further question is how the authors implement parallel decoding parts efficiently. Is it possible to use FlashAttention mechanism for this non-causal attention masking?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BE7QsYOhSx", "reviewer": "Reviewer_vZqn", "summary": "This work proposes a novel framework to augment an existing pretrained LLM with a differentiable cache module that can be finetuned on a set of data to improve performance of the model when its combined with aforementioned cache extender. Authors claim that the novel design of the module and look-ahead training objective allows the final model to perform reasoning in the latent space introduced by the cache processor. The cache processor operate by taking in the kv cache of a given prompt (x) input, and producing a sequence of latent embeddings that are appended to the existing representations computed by the frozen LLM model.\nExperiments and analysis done with proprietary data and public benchmarks show that the combined model improves the performance compared to the frozen LLM itself, as well a few other methods known from previous work.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer vZqn,\n\nWe sincerely thank you for your detailed review and insightful comments on our work (Submission 403). We appreciate you acknowledging the merit of our core idea, the quality of the related work section, and the potential for future work stemming from our method.\n\nWe understand your primary concerns revolve around the proprietary dataset used and the need for specific baselines and parameter details. We appreciate the opportunity to clarify these points.\n\n### **Regarding Proprietary Data and Fine-tuning Baseline:**\n\nWe recognize your significant concerns regarding the clarity of the \"proprietary training and validation data,\" its potential influence, and the corresponding lack of a baseline showing the frozen model fine-tuned on this same data. These points were raised multiple times and in your third question.\n\n**Our Response:**\n> We sincerely thank the reviewer for highlighting these concerns regarding our dataset and baselines. We understand that clarity here is crucial for evaluating our method's contribution.\n\n> The dataset used to train the coprocessor is *a subset of the data the model was pretrained on*. We did perform the suggested experiment of further finetuning the base model on this same dataset subset. However, as the base model was already extensively optimized on this data distribution during its initial pretraining, we observed negligible differences in perplexity or downstream performance. Standard fine-tuning paradigms predominantly yielded diminishing returns when applied directly to the original pretraining corpus. Furthermore, consistent with typical large-scale pretraining data, the corpus we used consists of text for self-supervised next-token prediction and lacks curated instruction-following examples or specific question-answering formats as in downstream tasks.\n\n> Given that this fine-tuning baseline showed no significant change from the original frozen model, we initially omitted it to maintain focus on comparisons highlighting the architectural contribution of the coprocessor. However, we understand the reviewer's point about baseline completeness. We will update the manuscript (e.g., in the dataset description or appendix) to clarify the nature of this dataset (as part of the original pretraining corpus, focused on self-supervision) and explicitly state the result of the fine-tuning baseline experiment, explaining why it yielded minimal change.\n\n### **Regarding Parameter Counts and Scaled Baseline:**\n\nRegarding your request for more details on parameter capacity and the suggestion of a scaled-up baseline, as asked in your second question:\n\n**Our Response:**\n> We thank the reviewer for the important question on parameter counts. In our main experiments, the coprocessor has the same number of parameters as the base LLM.\n> While a baseline scaling the base model to ≈2x size is a relevant comparison point, our primary goal was specifically enhancing and augmenting an existing, frozen model, rather than the distinct and computationally intensive task of training a larger model from scratch, placing the scaled baseline outside this study's scope.\n> Importantly, we also explored parameter efficiency using a LoRA variant, which added only ≈2% of the base model's parameters while still providing substantial performance gains.\n> We will revise the paper to clearly state all relevant parameter counts (full coprocessor and LoRA variant) and briefly discuss the scaled baseline context.\n\n### **Regarding Data Decontamination / Use of Open Data:**\n\nAddressing your first question regarding data decontamination, the potential use of open-source datasets, and the condition for increasing your score:\n\n**Our Response:**\n> We appreciate the reviewer's specific condition for potential score improvement. While re-running our main experiments on large open-source datasets is infeasible within the rebuttal period (though valuable future work), we confirm that we performed extensive decontamination of our proprietary pretraining data specifically against all benchmark tasks reported in Table 2.\n\n---\n\nWe hope these clarifications, the results from experiments already performed (such as the fine-tuning baseline), and our commitments to update the manuscript with further details (data description, parameter counts, decontamination confirmation) adequately address the main concerns raised. We thank you again for your constructive feedback and willingness to reconsider your score based on these clarifications.\n\nSincerely,\n\nThe Authors of Submission 403", "claims_and_evidence": "* The idea of adding more trainable parameters to allow the model to gain extra performance improvements has clear reasoning behind it. For instance, deep seek v3 showed how multi token prediction with little parameter overhead gives significant improvements to the base model. Here we see how extra latent representation conditioned on the input is useful for predicting the output sequence.\n\n* A big concern is lack of clarify about the \"proprietary training and validation data\". This weakens all experimental results as its unclear how that data is aligned with public benchmarks used in the experiments. Given that there exists lots of open source pretraining data (such as dclm), I don't see any benefit of using such data setup in a research paper aimed at sharing publicly in the conference.\n\n* A big concern is lack of baselines showing how would a frozen model finetuned on exactly the same data perform on the benchmarks. Given that cache processor model shares the arch (more on this in the next point) with the main LLM, such baseline is required.\n\n* Experiment descriptions do not mention much of details about the parameter capacity of cache processor w.r.t. the main model as well as competing methods such as LoRA. Its crucial to be shared as its one of the main drivers behind the model improvements: adding cache processor adds more parameters, so another baseline would be scaling up the frozen model with more params to match cache processor. While authors claim the efficiency of async cache processor abilities, its still a valid baseline that can serve as an upper bound of the performance.", "methods_and_evaluation_criteria": "Public benchmarks make sense and allows readers to understand the ballpark of the performance, and reproduce evaluations if needed.", "theoretical_claims": "There are no theoretical claims with corresponding proofs in the manuscript.", "experimental_designs_or_analyses": "The choice of training and validation data make it hard for readers to understand the influence of that data on the analysis of the method. This is one of the biggest concerns I have regarding this work.", "relation_to_broader_scientific_literature": "Related method section is well written and address different alternative methods such as reasoning in discrete space with CoT, kv cache compression for model efficiency; those method address features of the proposed method from different angles.", "essential_references_not_discussed": "The presented discussion provides enough context for this work.", "other_strengths_and_weaknesses": "This method might open room for future work that'd train a number of such \"expert\" cache processors that can be trained separately on specific data domains to contribute towards the final mixture model.", "other_comments_or_suggestions": "none", "questions_for_authors": "* I'd like to increase my score as long as authors would consider designing experiments using open sourced training datasets that are known to not be contaminated with any benchmark data. If thats not possible, would authors be able to confirm that they did extensive de-contamination of their data w.r.t. all benchmark tasks from table 2? \n\n* Could you please include more information about the capacity (number of parameters) of cache processors that were trained in the experimental models? \n\n* Could you please provide any baseline result showing the performance of the baseline/frozen LLM if it was finetuned using same data and same hyper-parameters?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes a novel and efficient method for augmenting LLMs using a differentiable coprocessor that injects latent embeddings into the KV-cache, enabling improved reasoning without modifying the base model. Despite some concerns about baseline comparisons and proprietary data, the reviewers converged on the view that the method is timely, well-motivated, and shows strong empirical promise across multiple tasks.
Decision: Accept (poster)