Deliberation in Latent Space via Differentiable Cache Augmentation

Dear Reviewer vZqn,

We 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.

We 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.

Regarding Proprietary Data and Fine-tuning Baseline:

We 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.

Our Response:

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.

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.

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.

Regarding Parameter Counts and Scaled Baseline:

Regarding your request for more details on parameter capacity and the suggestion of a scaled-up baseline, as asked in your second question:

Our Response:

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. 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. 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. We will revise the paper to clearly state all relevant parameter counts (full coprocessor and LoRA variant) and briefly discuss the scaled baseline context.

Regarding Data Decontamination / Use of Open Data:

Addressing your first question regarding data decontamination, the potential use of open-source datasets, and the condition for increasing your score:

Our Response:

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.

---

We 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.

Sincerely,

The Authors of Submission 403

Dear Reviewer 9maT,

Thank you for the positive evaluation, thoughtful review, and examination of the supplement. We appreciate your constructive questions and address them below:

Regarding Training Speed/Process:

You asked about the training speed and the forward passes involved.

Our Response:

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.

Regarding Asynchronous Inference Elaboration:

You asked for elaboration on the asynchronous inference process and the placement of latent embeddings when extra sampling steps occur.

Our Response:

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.

Regarding Sequential Latent Embedding Generation:

You asked if we tried outputting latent embeddings sequentially.

Our Response:

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.

Regarding Scalability and Heterogeneous Model Sizes:

You 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.

Our Response:

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.

Regarding Multi-step Latent Generation:

You asked about generating and using latent embeddings multiple times during generation.

Our Response:

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.

Regarding results of 1-ahead + Cache Augmentation:

You asked about the results when training with 1-ahead token prediction combined with cache augmentation.

Our Response:

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.

Regarding Suggested References:

You suggested citing concurrent works like CCOT and SoftCoT.

Our Response:

Thank you for suggesting CCOT and SoftCoT. We agree they are relevant and will cite/discuss them appropriately in the revised related work section.

Regarding Efficient Parallel Decoding Implementation:

You had a further question on implementing the parallel decoding efficiently, specifically regarding FlashAttention.

Our Response:

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).

---

We 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.

Sincerely,

The Authors of Submission 403

Dear Reviewer n6p1,

Thank 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.

We would like to address your specific comments and question:

Regarding Inference Overhead (Time/Space):

You suggested including comments on the overhead related to generating embeddings and using two models at inference time.

Our Response:

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.

Regarding Comparison with Coconut:

You suggested a more detailed comparison with the Coconut paper.

Our Response:

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.

Regarding Relation to RL Techniques for Reasoning:

You asked how our results/method for math reasoning compare to recent RL techniques and if there's potential synergy.

Our Response:

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.

---

Thank you again for your valuable feedback and positive assessment of our work. We believe addressing these points will further strengthen the paper.

Sincerely,

The Authors of Submission 403

Dear Reviewer w2nr,

Thank 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:

Regarding Claim 1 (Perplexity Reduction & Baseline Fairness):

You questioned the baseline fairness due to training data and suggested a fine-tuning comparison.

Our Response:

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.

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.

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.

Regarding Claim 2 (Performance Improvement vs. CoT & Evaluation Scope):

You noted the limited evaluation scope (GSM-8K only) for the CoT comparison.

Our Response:

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.

Regarding Claim 3 (Efficiency Benefits vs. CoT & Hardware Considerations):

You questioned the efficiency claim due to lack of empirical evidence considering hardware limitations.

Our Response:

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.

Regarding your Question (Limitations/Tradeoffs of Parallel vs. Sequential Computation):

You asked about the limitations of the parallel approach vs. sequential CoT.

Our Response:

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.

We hope this addresses your concerns and thank you again for the valuable feedback to improve the paper.

Sincerely,

The Authors of Submission 403