Have you ever crammed vocabulary all night, aced the test the next day, and then forgotten everything you learned the week before?
You're not alone. New research shows large language models suffer from "attention amnesia": when you fine-tune a hybrid attention model on Chain-of-Thought (CoT) data to boost reasoning, its long-context retrieval ability collapses. HypeNet-9B falls from 67.2% to 9.4% accuracy on 256K-context retrieval tasks.
Hybrid attention models: a trade-off between efficiency and memory
Background: standard Transformers use softmax attention, whose compute grows quadratically with sequence length—256K context costs 4096× more than 4K. Linear attention (e.g., Mamba, GLA) reduces compute to linear, but at the price of compression: like compressing a book into a mind map, details are inevitably lost.
Hybrid models are the compromise: most layers use linear attention for efficiency, while a few softmax attention layers handle "long-range memory." These softmax layers act like the hippocampus in the brain—routing distant information into current decisions.
CoT fine-tuning: medicine for reasoning, poison for memory
The problem arises during fine-tuning. When you perform supervised fine-tuning (SFT) on CoT data (e.g., detailed math reasoning steps), gradient updates systematically favor "short-range patterns"—because each step of CoT reasoning depends mainly on the previous few steps, not on information tens of thousands of tokens back.
The paper proposes a "gradient locality theorem" to explain this: CoT data has Markov structure (each step mainly depends on the previous one), so attention gradients concentrate on nearby tokens, and distant query-key (QK) projections are gradually "forgotten."
An analogy: CoT fine-tuning is like making a student repeatedly do math proofs. Each problem only requires looking at the previous two or three steps, so over time their ability to look up distant references atrophies—the training simply never uses it.
QK-Restore: a zero-cost memory recovery technique
Once the cause is identified, the fix is remarkably simple: swap back only the pre-fine-tuning QK weights, keeping every other parameter at its fine-tuned version.
It's like a student who has completed reasoning boot camp keeping all their reasoning skills, but putting back their pre-training "reference-lookup glasses." Specifically:
1. QK-Restore: directly replace the fine-tuned W_Q and W_K with the pre-fine-tuning versions—zero training cost. 2. QK-Pro (Procrustes variant): when direct replacement hurts reasoning performance, use Procrustes alignment to balance "preserving routing" and "adapting to reasoning."
Results? HypeNet-5B's 256K-context retrieval accuracy recovers from 65.4% to 76.4%, while reasoning performance remains essentially unchanged. Consistent improvements appear on other architectures such as Jet-Nemotron.
Deeper implications
This work is more than "fixing a bug":
- A win for mechanistic interpretability: the researchers didn't stumble onto the problem by trial and error—they located the QK projections as the "victim" through rigorous gradient analysis, offering a methodology for understanding fine-tuning side effects.
- Capabilities aren't free: there's a deep tension between reasoning and long-range memory. Optimizing one can silently degrade the other—a warning for everyone fine-tuning models.
- Simple solutions are often the most effective: no retraining, no extra data—just swap back two weight matrices. Good solutions aren't necessarily complex.
Limitations
The paper mainly validates on hybrid linear-attention models; pure softmax models are less affected (softmax naturally maintains global attention). Also, QK-Restore is a "rollback" strategy—a more fundamental solution would be fine-tuning methods that never damage long-range memory in the first place.
---
Paper: Attention Amnesia in Hybrid LLMs: When CoT Fine-Tuning Breaks Long-Range Recall, and How to Fix It Authors: Xinyu Zhou, Boyu Zhu, Yi Xu, Zhiwei Li, Yingfa Chen, Huiming Wang, Zhijiang Guo Link: https://arxiv.org/abs/2606.11052 Code: https://github.com/LARK-AI-Lab/QK-Restore