跳到正文
原文
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微调居然会弄坏模型的长距离记忆,这篇论文不仅把原因扒清楚了,还给出了零成本修复方案,值得放进参考列表。

正文

View PDF HTML (experimental)

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