混合大语言模型中的注意力遗忘症：思维链微调如何破坏长程记忆及其修复方法

摘要

链式思维（CoT）监督微调（SFT）被广泛用于提升推理能力，但我们发现它会系统性降低混合线性注意力模型的长上下文召回能力。在HypeNet和Jet-Nemotron等架构中，经过CoT-SFT后，NIAH（大海捞针）检索任务的性能显著下降，且检索设置越困难、上下文窗口越长，退化越严重。例如，HypeNet-9B在NIAH-S2@256K上的性能从67.2%降至9.4%。我们将此归因于CoT-SFT使注意力梯度偏向短程模式，破坏了负责长程路由的查询-键投影（W_Q, W_K）。基于这一发现，我们提出QK-Restore——一种无需训练的方法，仅从SFT前的检查点恢复W_Q和W_K，同时保留其他所有SFT后的参数。我们进一步引入一个Procrustes变体，以平衡路由保留和推理适应。跨架构实验表明，QK-Restore在零训练成本下持续恢复长上下文能力，同时保持推理性能；例如，在HypeNet-5B上，它使S3@256K从65.4%提升至76.4%，同时保持强大的推理能力。

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.