Bottlenecked Transformers: Periodic KV Cache Consolidation for Generalised Reasoning

This paper introduces the Bottlenecked Transformer, a novel architecture that enhances general reasoning by applying Information Bottleneck theory to periodically consolidate and reconsolidate KV cache entries via an auxiliary processor, achieving significant performance gains on math benchmarks compared to standard and pause-token baselines.

The Gist

Imagine you are trying to solve a very difficult math problem. You start writing down your thoughts, step by step. As you write, your brain holds onto every single word you've ever written, every number you've calculated, and every rule you've recalled.

Eventually, your "mental scratchpad" gets so crowded with details that it becomes hard to see the big picture. You remember everything, but you can't easily find the specific piece of information you need right now to solve the next step.

This is exactly the problem the paper "Bottlenecked Transformers" tries to solve for AI models.

Here is the story of their solution, explained simply.

1. The Problem: The AI's "Overloaded Brain"

Modern AI models (like the ones that write essays or solve math) work by predicting the next word. To do this, they keep a running memory of everything they've said so far, called a KV Cache.

Think of the KV Cache like a growing pile of sticky notes.

• Every time the AI thinks of a new step, it adds a new sticky note to the pile.
• The problem is that the AI never throws anything away. It keeps every detail, even the boring ones or the ones that don't matter anymore.
• As the pile gets huge, the AI gets "distracted" by irrelevant details. It struggles to generalize (apply what it learned to new problems) because it's too focused on memorizing the exact history rather than understanding the logic.

2. The Inspiration: How Human Brains Work

The authors looked at how human brains handle memory. We have two special processes:

• Consolidation: When you learn something new, your brain stabilizes it so it sticks.
• Reconsolidation: When you remember something old, your brain briefly makes that memory "plastic" (malleable) again. It updates that old memory with new context before locking it back down.

The Analogy: Imagine you are writing a diary.

• Standard AI: You write every single thought down and never edit. Your diary becomes a 1,000-page mess of rambling.
• Human Brain: Every night, you review your diary. You rewrite the messy parts to make them clearer, and you update old entries with new insights you gained today. You keep the essence but throw away the clutter.

3. The Solution: The "Bottlenecked Transformer"

The authors built a new type of AI that does this "diary review" automatically. They call it the Bottlenecked Transformer.

Here is how it works, step-by-step:

The "Pause" Button

The AI doesn't just keep typing forever. Every time it finishes a logical step (like finishing a sentence or a math equation), it hits a "Pause."

The "Cache Processor" (The Editor)

At this pause, a special, smaller AI module (called the Cache Processor) wakes up. It doesn't write new text; instead, it acts as an Editor for the AI's memory.

• Consolidation: It looks at the most recent thoughts the AI just had and rewrites them to make them clearer and more stable.
• Reconsolidation: It looks back at the most important old thoughts (the ones it needs to remember) and updates them with the new context it just learned.

The "Bottleneck"

Why call it a "Bottleneck"?
Imagine a funnel. If you pour a huge bucket of water (all the raw data) into a narrow neck, the water has to squeeze through. This forces the water to organize itself.

• The AI is forced to squeeze its massive, messy memory through this "Editor."
• It keeps the predictive information (the logic needed to solve the problem) but discards the redundant noise (the unnecessary details).
• This makes the AI's memory more efficient and smarter, not just bigger.

4. The Results: Smarter Math Solvers

The researchers tested this on hard math problems (like those found in high school competitions).

• The Old Way: The AI just kept generating text, getting confused by its own long history.
• The New Way: The AI paused, cleaned up its memory, updated its understanding, and then continued.

The Outcome: The new AI solved significantly more problems correctly. It didn't just get better at memorizing; it got better at reasoning. It was able to take what it learned in one problem and apply it to a slightly different one, just like a human student who understands the concept rather than just memorizing the steps.

Summary

Think of the Bottlenecked Transformer as an AI that has learned the art of reflection.

Instead of mindlessly churning out words and hoarding every detail, it stops periodically to say: "Wait, let me clean up my notes. What actually matters here? Let me update my old memories with this new insight."

By doing this "mental housekeeping," the AI becomes less cluttered, more focused, and surprisingly better at solving complex puzzles.

Technical Summary

1. Problem Statement

Large Language Models (LLMs) based on Transformers have demonstrated strong reasoning capabilities, particularly when inference-time compute is increased via "thinking" chains of thought (token-space generation). However, a growing body of research explores Auxiliary Latent-Space Computation (ALSC), where models perform extra computation in their internal continuous states (latent space) rather than emitting intermediate tokens.

Existing ALSC methods generally fall into three categories:

1. Token-mediated latent rollouts: Generating internal "latent tokens" that extend the sequence.
2. Residual/activation steering: Modifying hidden states to influence output.
3. Memory (KV) compression: Pruning or merging Key-Value (KV) cache entries to save memory.

The Gap: Current methods largely focus on compression (reducing memory footprint) or extension (adding more tokens). There is an underexplored direction: Memory Consolidation and Reconsolidation. In neuroscience, consolidation stabilizes new memories, while reconsolidation allows recalled memories to be updated with new context before restabilizing. The authors argue that standard autoregressive Transformers lack a mechanism to periodically reprocess and "rewrite" their working memory (KV cache) to discard irrelevant details and integrate new context, which hinders generalization in complex reasoning tasks.

2. Theoretical Motivation: Information Bottleneck (IB) Theory

The authors provide a theoretical justification for memory rewriting using Information Bottleneck (IB) theory.

• The Goal: Optimize a latent representation $Z$ to be maximally informative about the output $Y$ (next reasoning step) while being minimally informative about the input $X$ (raw history). This balance ($I(X;Z)$ vs. $I(Z;Y)$) drives generalization.
• The Problem with Vanilla Transformers: In standard autoregressive training, the KV cache ($C_{0:n}$) acts as a "terminal bottleneck." The training objective (maximizing next-token likelihood) incentivizes the model to preserve all historical information to predict the next token perfectly. Consequently, the cache accumulates extraneous, non-predictive details ($I(X;Z)$ remains high), which can hinder the model's ability to generalize to new reasoning patterns.
• The Proposed Solution: Periodically rewriting the KV cache can act as a reprocessing step. By selectively compressing $I(X;Z)$ (removing noise/irrelevant details) while preserving or enhancing $I(Z;Y)$ (predictive utility), the model can achieve a better information bottleneck, leading to improved generalization.

3. Methodology: The Bottlenecked Transformer

The authors introduce the Bottlenecked Transformer, an architecture that augments a pre-trained backbone LLM with a Cache Processor.

Architecture

• Backbone: A standard decoder-only Transformer (frozen during the second stage of training).
• Cache Processor: A smaller, auxiliary Transformer that operates in parallel with the backbone. It is invoked at specific intervals (e.g., at newline tokens marking the end of a reasoning step).
• Mechanism:
  1. Selection: When invoked, the processor selects two sets of KV entries:
     • Recent Tokens: The KV entries from the most recent reasoning step (Consolidation).
     • Recalled Tokens: A top-$k$ set of prior KV entries selected based on attention mass relative to the recent step (Reconsolidation).
  2. Processing: These selected entries are converted into "KV-tokens" and fed into the Cache Processor. The processor performs non-causal, in-place rewrites on these entries.
  3. Rewrite: The processor outputs updates ($\Delta K, \Delta V$) which are added to the original KV entries via a gated residual connection. Crucially, no dimensionality reduction is applied; the goal is content restructuring, not compression.

Training Strategy

The training occurs in two stages:

1. Stage 1 (SFT): The backbone is fine-tuned on reasoning data using standard next-token prediction.
2. Stage 2 (Processor Training): The backbone is frozen. The Cache Processor is trained to minimize the cross-entropy loss of the next reasoning step, conditioned on the rewritten cache.
   • The processor learns to reorganize the memory to improve future predictions without an explicit IB loss function; the compression emerges implicitly via the optimization of predictive efficiency.

4. Key Contributions

1. Theoretical Framework: A novel application of Information Bottleneck theory to decoder-only Transformers, arguing that periodic KV cache rewrites are necessary to break the "high-fidelity but non-compressive" nature of standard autoregressive caches, thereby improving generalization.
2. Novel Architecture: The Bottlenecked Transformer, which implements neuroscientific concepts of consolidation/reconsolidation via a dedicated Cache Processor that performs in-place, non-causal rewrites of KV entries.
3. Empirical Validation: Extensive evaluation showing that this architecture outperforms vanilla Transformers and other ALSC baselines (pause tokens, latent rollouts) on mathematical reasoning tasks.

5. Experimental Results

The model was evaluated on seven benchmarks (GSM8K, MATH, SVAMP, TheoremQA, LogiQA, Gaokao-MathQA, GSM-Hard) across four backbone models (Llama 3.2 1B/3B, Llama 3.1 8B, Qwen 3 0.6B).

• Performance Gains: The Bottlenecked Transformer consistently outperformed baselines.
  • Llama 3.2 1B: Improved SVAMP accuracy by +6.6% (38.0% $\to$ 44.6%).
  • Llama 3.2 3B: Improved GSM8K by +4.6%.
  • Llama 3.1 8B: Improved LogiQA by +3.1%.
• Comparison to Baselines:
  • Pause Tokens: Showed variable performance, often underperforming standard SFT unless paired with continued pretraining.
  • Latent Rollouts: Generally underperformed, sometimes causing model destabilization (especially in larger models like Llama 3.1 8B).
• Ablation Studies:
  • Top-$k$ Reconsolidation: Moderate values ($k \approx 32-64$) were optimal for most tasks, while harder tasks (MATH) benefited from larger $k$.
  • Recent Step Window ($R$): Performance was robust across a range of window sizes, suggesting the mechanism benefits from medium-horizon context reshaping rather than token-by-token updates.
  • Rewrite Magnitudes: Analysis showed the processor primarily edits Value vectors (memory content) rather than Key vectors (memory addressing), with the most significant changes occurring in the earliest layers of the network.

6. Significance and Future Work

• Significance: This work bridges the gap between neuroscience (memory consolidation) and LLM architecture. It demonstrates that reprocessing internal states (rather than just generating more tokens or compressing memory) is a viable path to improving reasoning generalization. It challenges the assumption that larger caches are always better, suggesting that quality of memory (via reorganization) matters more than quantity.
• Limitations & Future Directions:
  • Credit Assignment: Training solely on next-step loss can lead to high variance in credit assignment for cache rewrites.
  • Biological Fidelity: The current model collapses consolidation and reconsolidation into a single online step. Future work could separate these into offline (replay-based) and online (prediction-error gated) processes.
  • Explicit IB Loss: Future iterations could explore explicit noise injection or denoising objectives to directly optimize the Information Bottleneck trade-off.

In summary, the Bottlenecked Transformer offers a compelling alternative to token-heavy reasoning, proving that periodic, intelligent rewrites of the model's internal memory can significantly boost reasoning capabilities without increasing inference latency via token generation.