混合 LLM 中的注意力遗忘症：当 CoT 微调 破坏远程回忆时，以及如何修复它

思想链（CoT）监督的 微调（SFT）被广泛采用来提高推理能力，但我们发现它系统地降低了混合线性注意力模型中的长上下文回忆。在包括 HypeNet 和 Jet-Nemotron 在内的架构中，Needle-In-A-Haystack (NIAH) 上的检索性能在 CoT-SFT 后大幅恶化，并且在更难的检索设置和更长的上下文窗口下，退化变得更加严重。例如，NIAH-S2@256K 上的 HypeNet-9B 从 $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\%$，同时保持强大的推理性能。