Attention Amnesia in Hybrid LLMs: When CoT Fine-Tuning Breaks Long-Range Recall, and How to Fix It

📄 arXiv: 2606.11052v1 📥 PDF

作者: Xinyu Zhou, Boyu Zhu, Yi Xu, Zhiwei Li, Yingfa Chen, Huiming Wang, Zhijiang Guo

分类: cs.CL

发布日期: 2026-06-09

备注: 28 pages

💡 一句话要点

提出QK-Restore以解决混合LLM中的长范围记忆衰退问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长范围记忆 链式思维 混合线性注意力 无训练方法 信息检索 推理能力提升

📋 核心要点

  1. 现有的链式思维微调方法在提升推理能力的同时，导致混合线性注意力模型的长范围记忆能力显著下降。
  2. 论文提出QK-Restore方法，通过恢复查询-键投影参数，解决了长上下文记忆衰退的问题，且无需额外训练。
  3. 实验结果显示，QK-Restore在多个架构上有效恢复了长上下文能力，例如HypeNet-5B的性能从65.4%提升至76.4%。

📝 摘要（中文）

链式思维（CoT）监督微调（SFT）被广泛应用于提升推理能力，但我们发现其在混合线性注意力模型中系统性地降低了长上下文的记忆能力。在包括HypeNet和Jet-Nemotron等架构中，经过CoT-SFT后，针对于“针在干草堆中”（NIAH）的检索性能显著下降，尤其在更困难的检索设置和更长的上下文窗口下更为明显。我们将此归因于CoT-SFT使注意力梯度偏向短期模式，干扰了负责长范围路由的查询-键投影（$W_Q, W_K$）。基于此观察，我们提出了QK-Restore，这是一种无训练的方法，仅恢复来自预-SFT检查点的$W_Q$和$W_K$，同时保留所有其他后SFT参数。实验表明，QK-Restore在零训练成本下有效恢复了长上下文能力，同时保持了推理性能。

🔬 方法详解

问题定义：论文要解决的具体问题是链式思维微调（CoT-SFT）导致混合线性注意力模型在长上下文检索中的性能显著下降，尤其是在更复杂的检索任务中。现有方法在提升推理能力的同时，未能有效维持长范围记忆能力。

核心思路：论文的核心解决思路是通过提出QK-Restore方法，专注于恢复查询-键投影参数（$W_Q, W_K$），以此来修复长范围路由的能力，而不需要对其他参数进行训练。这样设计的原因在于，注意力梯度的偏向性影响了长范围信息的处理。

技术框架：整体架构包括两个主要阶段：首先是进行链式思维的监督微调，然后在此基础上应用QK-Restore方法，仅恢复特定的查询-键投影参数。该方法不涉及额外的训练过程，直接利用预训练模型的参数。

关键创新：最重要的技术创新点在于提出了QK-Restore这一无训练方法，能够有效恢复长上下文能力，同时保持推理性能。这与现有方法的本质区别在于，传统方法往往需要重新训练或调整多个参数。

关键设计：在QK-Restore中，关键设计包括对查询-键投影参数的选择和恢复策略，确保在恢复过程中不影响其他参数的性能。此外，论文还引入了Procrustes变体，以平衡路由保持和推理适应之间的关系。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，QK-Restore在HypeNet-5B上将S3@256K的性能从65.4%提升至76.4%，同时保持了强大的推理性能。这一提升在多个架构上均得到了验证，展示了该方法的有效性和广泛适用性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、信息检索和对话系统等。通过恢复长上下文记忆能力，QK-Restore能够提升模型在复杂任务中的表现，具有重要的实际价值和未来影响，尤其是在需要处理大量上下文信息的场景中。

📄 摘要（原文）

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.