Research map / Large language models / Reasoning and chain of thought

Deliberation in Latent Space via Differentiable Cache Augmentation

ICML 2025 · Large language models › Reasoning and chain of thought · approach: Efficient LLM serving · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Trains a coprocessor (initialized from Gemma-2 2B) that reads the frozen LLM's kv-cache and emits latent embeddings inserted into the cache, using only the language-modeling loss on pretraining data, yielding gains on reasoning and knowledge benchmarks.

Problem

Standard test-time reasoning methods such as chain-of-thought generate discrete intermediate tokens just-in-time, which is slow and hard to optimize end-to-end; the paper asks whether a frozen LLM can instead be improved by latent computation over its kv-cache.

Approach

A coprocessor, a copy of the frozen LLM, takes the kv-cache plus trainable soft tokens and produces latent embeddings in a single forward pass; these are appended to the cache so the frozen LLM decodes with augmented context. Training uses multi-position augmentation with ahead-token prediction in one forward pass via a custom attention mask, so it is differentiable and needs no RL.

Key results

• With 64 latent embeddings on Gemma-2 2B, GSM8K 8-shot accuracy rises from 21.38% to 31.43% (+10.05) and MMLU 5-shot from 52.00% to 56.70% (+4.70).
• Validation perplexity at the first token after augmentation drops by up to 3.94% with 64 latents, and gains persist to the 32nd token.
• With 32 latents, the method reaches 26.76% GSM8K vs 22.37% for Pause Token and 21.38% baseline, with validation perplexity 10.60 vs 11.63 for Pause Token.

Measured effects reported in the paper

Surprising findings

• Training the coprocessor from random initialization still beats the baseline on most benchmarks, though fine-tuning from pretrained weights is better.
• LoRA-tuned coprocessors give smaller gains (23.35% and 24.03%) than from-scratch (25.78%) and full fine-tuning (26.76%) on GSM8K.

Limitations

• Gains on several benchmarks are small or negative (e.g., HumanEval, AGIEval, SIQA, MATH at 4 latents), and the method is evaluated on a single 2B model size.
• Evaluation relies on a proprietary validation set and results are single-run with no reported variance for most benchmarks.
• The coprocessor doubles the model parameter count and compute at training time, and inference cost of the coprocessor is not fully analyzed.

Experimental setup

Datasets: Gemma-2 pretraining mixture (2T tokens) GSM8K training set MATH training set

Benchmarks: MMLU GSM8K DROP ARC-e ARC-c MATH Winogrande PIQA SIQA HellaSwag BoolQ MBPP

Metrics: accuracy F1 pass@1 perplexity

Baselines: Gemma-2 2B baseline Pause Token zero-shot CoT COCONUT last-layer activation augmentation

Builds on: Pause Token Quiet-STaR chain-of-thought prompting LoRA

Closest papers by meaning

• Reasoning Cache: Continual Improvement Over Long Horizons via Short-Horizon RL (ICML 2026 · similarity 0.90)
• Cache-to-Cache: Direct Semantic Communication Between Large Language Models (ICLR 2026 · similarity 0.89)
• Get More with LESS: Synthesizing Recurrence with KV Cache Compression for Efficient LLM Inference (ICML 2024 · similarity 0.89)
• RelayCaching: Accelerating LLM Collaboration via Decoding KV Cache Reuse (ICML 2026 · similarity 0.89)
• CateKV: On Sequential Consistency for Long-Context LLM Inference Acceleration (ICML 2025 · similarity 0.89)
• Hogwild! Inference: Parallel LLM Generation via Concurrent Attention (NeurIPS 2025 · similarity 0.88)