HuggingFace Daily Papers(社区热门论文)· HuggingFace Daily Papers(社区热门论文)·· 2026-06-09精选AI 评分75
混合LLM中的注意力失忆:CoT微调破坏长距离召回及修复方法
AI 导读
CoT监督微调系统性地降低混合线性注意力模型(如HypeNet、Jet-Nemotron)的长上下文召回能力。在NIAH任务上,HypeNet-9B的S2@256K从67.2%降至9.4%,原因是CoT-SFT使注意力梯度偏向短程模式,破坏长程路由的W_Q和W_K投影。QK-Restore方法无需训练,从微调前检查点恢复W_Q和W_K,保留其余参数;Procrustes变体平衡路由保留与推理适应。在HypeNet-5B上,QK-Restore将S3@256K从65.4%提升至76.4%,推理性能不变。
推荐理由
做长上下文推理的同学注意了,CoT微调居然会弄坏模型的长距离记忆,这篇论文不仅把原因扒清楚了,还给出了零成本修复方案,值得放进参考列表。
正文
Abstract:Chain-of-thought (CoT) supervised fine-tuning (SFT) is widely adopted to improve reasoning ability, yet we find that it systematically degrades long-context recall in hybrid linear-attention models. Across architectures including HypeNet and Jet-Nemotron, retrieval performance on Needle-In-A-Haystack (NIAH) deteriorates substantially after CoT-SFT, and the degradation becomes more severe under harder retrieval settings and longer context windows. For example, HypeNet-9B on NIAH-S2@256K decreases from $67.2\%$ to $9.4\%$. We attribute this to CoT-SFT biasing attention gradients toward short-range patterns, disrupting query-key projections ($W_Q, W_K$) that are responsible for long-range routing. Motivated by this observation, we propose QK-Restore, a training-free method that restores only $W_Q$ and $W_K$ from the pre-SFT checkpoint while preserving all other post-SFT parameters. We further introduce a Procrustes variant to balance routing preservation and reasoning adaptation. Across architectures, QK-Restore consistently restores long-context capability at zero training cost while preserving reasoning performance; for instance, on HypeNet-5B it improves S3@256K from $65.4\%$ to $76.4\%$ while maintaining strong reasoning performance.
| Comments: | Accepted to EMNLP 2026 (Main) |
| Subjects: | Computation and Language (cs.CL) |
| Cite as: | arXiv:2606.11052 [cs.CL] |
| (or arXiv:2606.11052v2 [cs.CL] for this version) | |
| https://doi.org/10.48550/arXiv.2606.11052 arXiv-issued DOI via DataCite |
Submission history
From: Xinyu Zhou [view email]
[v1]
Tue, 9 Jun 2026 16:17:19 UTC (3,923 KB)
[v2]
Sun, 30 Aug 2026 15:47:03 UTC (3,923 KB)
来源:HuggingFace Daily Papers(社区热门论文) · arxiv.org