混合LLM中的注意力失忆：CoT微调破坏长距离召回及修复方法
阅读原文· arxiv.org做长上下文推理的同学注意了，CoT微调居然会弄坏模型的长距离记忆，这篇论文不仅把原因扒清楚了，还给出了零成本修复方案，值得放进参考列表。
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）监督微调（SFT）被广泛用于提升推理能力，但我们发现它会系统性地降低混合线性注意力模型的长上下文召回能力。在包括HypeNet和Jet-Nemotron在内的多种架构上，经过CoT-SFT后，在"大海捞针"（NIAH）测试中的检索性能大幅恶化，并且在更困难的检索设置和更长的上下文窗口下，这种退化变得更加严重。例如，HypeNet-9B在NIAH-S2@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%，同时保持了强大的推理性能。