CoT微调会损害长文本召回，提出无需训练的修复方法
Attention Amnesia in Hybrid LLMs: When CoT Fine-Tuning Breaks Long-Range Recall, and How to Fix It
- 通过恢复预微调时的查询键投影矩阵修复长程记忆
- 在256K上下文中，召回率从9.4%提升至76.4%
- 零训练成本，适合需长上下文推理的部署场景
链式思维（CoT）监督微调虽广泛用于提升模型推理能力，但会系统性破坏混合线性注意力模型的长上下文召回性能。在针堆草场（NIAH）任务中，如HypeNet-9B在256K上下文下的召回率从67.2%降至9.4%，且在更难设置和更长上下文时恶化更严重。我们发现其根源在于CoT-SFT使注意力梯度偏向短程模式，干扰负责长程路由的查询-键投影矩阵（W_Q, W_K）。为此提出QK-Restore：仅从预微调检查点恢复W_Q、W_K，保留其余后微调参数。进一步引入Procrustes变体平衡路由保持与推理适应。在多架构上，QK-Restore以零训练成本稳定恢复长上下文能力，如HypeNet-5B在S3@256K下从65.4%提升至76.4%，同时维持强推理表现。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。