周新宇
LARK，香港科技大学（广州）
朱博宇
伦敦大学学院
徐毅
Mistral AI
李志伟
LARK，香港科技大学（广州）
陈英发
王慧明
新加坡科技设计大学
郭志江
LARK，香港科技大学（广州）
摘要
思维链监督微调被广泛用于提升推理能力，然而我们发现，在混合线性注意力模型中，它会系统性地损害长上下文召回能力。在包括 HypeNet 和 Jet-Nemotron 在内的多种架构上，经过思维链监督微调后，模型在"大海捞针"测试中的检索性能显著下降，并且在更困难的检索设置和更长的上下文窗口下，退化程度更为严重。例如，HypeNet-9B 在 NIAH-S2@256K 上的性能从某个值下降到了另一个值。我们将此归因于思维链监督微调使注意力梯度偏向短距离模式，从而破坏了负责长距离路由的查询-键投影。基于这一观察，我们提出了 QK-Restore，一种无需训练的方法，该方法仅从监督微调前的检查点恢复查询投影和键投影，同时保留所有其他监督微调后的参数。我们进一步引入了一种 Procrustes 变体，以平衡路由保留与推理适应。在各种架构上，QK-Restore 在零训练成本下持续恢复长上下文能力，同时保持推理性能；例如，在 HypeNet-5B 上，它将 S3@256K 从某个值提升到另一个值，同时保持了强大的推理性能。
摘要
思维链（CoT）监督微调（SFT）被广泛用于提升推理能力，然而我们发现它会系统性地降低混合线性注意力模型的长上下文召回能力。在包括 HypeNet 和 Jet-Nemotron 在内的多种架构上，经过 CoT-SFT 后，大海捞针（NIAH）测试中的检索性能显著下降，并且在更困难的检索设置和更长的上下文窗口下，这种退化变得更加严重。例如，HypeNet-9B 在 NIAH-S2@256K 上的表现从 [原文此处数值缺失] 下降至 [原文此处数值缺失]。我们将此归因于 CoT-SFT 使注意力梯度偏向短程模式，从而破坏了负责长程路由的查询-键投影（query-key projections）。基于这一发现，我们提出了 QK-Restore，一种无需训练的方法，该方法仅从预 SFT 检查点恢复 [原文此处变量缺失] 和 [原文此处变量缺失]，同时保留所有其他后 SFT 参数。我们还引入了一种 Procrustes 变体，以平衡路由保留与推理适应。在多种架构上，QK-Restore 在零训练成本下一致地恢复了长上下文能力，同时保持了推理性能；例如，在 HypeNet-5B 上，它将 S3@256K 从 [原文此处数值缺失] 提升至 [原文此处数值缺失]，同时保持了强大的推理性能。
1 引言
高效序列模型，例如线性注意力（katharopoulos2020transformersrnnsfastautoregressive; qin2024variouslengthsconstantspeed; yang2024gatedlinearattentiontransformers）和状态空间模型（gu2024mambalineartimesequencemodeling; dao2024transformersssmsgeneralizedmodels），已成为长上下文处理中 softmax 注意力（vaswani2023attentionneed）的极具吸引力的替代方案，通过紧凑的循环或结构化状态将二次复杂度降低为线性复杂度。然而，这种压缩在长上下文理解（zhang2024hedgehogporcupineexpressive）所依赖的强召回任务（例如“大海捞针”NIAH; hsieh2024ruler）上引入了信息瓶颈。混合线性注意力模型通过保留少量 softmax 注意力层用于全局召回，同时将其余层转换为高效的线性注意力层（nvidia2025nvidianemotron3efficient），来缓解这种权衡。尽管 Qwen3.5（qwen35blog）和 Kimi Linear（kimiteam2025kimilinearexpressiveefficient）等原生混合模型实现了强大的性能和效率，但从头开始预训练它们仍然非常耗费资源。这一成本推动了近期的工作，即通过知识蒸馏，仅使用一小部分预训练 token（chen2026hybridlinearattentionright; li2025distilling; hoshino2025radredundancyawaredistillationhybrid），将预训练的 softmax 注意力模型转换为混合线性注意力模型。现有的转换工作主要关注在架构转换过程中保留预训练能力，而下游推理后训练对长上下文检索的影响仍未得到充分探索。
虽然蒸馏混合模型主要作为高效的基础模型使用，但现代大语言模型通常依赖后训练来获得更强的指令遵循和推理能力（Guo_2025；qwen2025qwen25technicalreport；ouyang2022traininglanguagemodelsfollow）。特别是，链式思维监督微调被广泛用于增强数学和多步推理（wei2023chainofthoughtpromptingelicitsreasoning；li2025tl）。我们观察到，在数学推理数据上应用链式思维监督微调可以提升推理性能，但会显著降低蒸馏混合模型在预训练或架构转换过程中获得的长上下文召回能力，尤其是在更具挑战性的 NIAH 设置和扩展上下文长度中（图 1）。与一般的灾难性遗忘不同，这种退化具有高度结构性：它主要影响由保留的 softmax 注意力层介导的长距离检索行为，而预期的推理增益则基本保持不变。这揭示了混合模型中的一个根本性矛盾：链式思维监督微调在强化局部推理的同时，可能会破坏长距离检索所需的路由机制。
为了研究这一现象，我们将链式思维推理轨迹建模为中间推理状态上的潜在马尔可夫过程，以反映数学推导中局部的逐步结构（wang2026doeschainofthoughthelpmarkovian；prystawski2023thinkstepstepreasoning）。在此链式思维-马尔可夫假设下，我们推导出注意力 logits 上的期望梯度幅度随 token 距离呈指数衰减。我们通过测量 token 自相关性和注意力 logits 的梯度，从经验上验证了这一预测。这些结果表明，链式思维监督微调可以改善局部多步推理，同时会侵蚀检索所需的长距离路由行为。
为了更好地隔离召回能力下降的根源，我们通过路由-提取分解来分析保留的 softmax 注意力层。我们发现，CoT-SFT 会在查询-键投影中引发局部性偏差漂移，而查询-键投影决定了信息检索的来源；相比之下，值侧提取仍能从 SFT 后的适应中受益。由于蒸馏混合模型中的长上下文召回高度依赖于少量保留的 softmax 注意力层，这种查询-键漂移会不成比例地破坏检索能力。受此不对称性的启发，我们提出了 QK-Restore，这是一种无需训练的方法，仅从 SFT 前的检查点恢复这些层中的查询和键投影，同时保留所有其他 SFT 后的参数，从而在恢复长距离路由能力的同时，保留 CoT-SFT 带来的大部分推理改进。我们的贡献可总结如下：
-
我们识别出 CoT-SFT 引发的召回能力下降是蒸馏混合模型的一种结构化后训练失效模式。
-
我们提供了理论和实证分析，表明 CoT-SFT 将训练信号集中在局部 token 交互上，并在保留的 softmax 注意力层的查询-键路由几何结构中引发局部性偏差漂移。
-
我们引入了 QK-Restore，这是一种无需训练的方法，仅恢复查询和键投影，从而在基本保持推理性能的同时，恢复长上下文召回能力。
2 相关工作
高效长上下文模型与混合注意力机制。近期混合注意力模型将 softmax 注意力层与高效循环或线性注意力层交错结合，在常识推理和需要高召回率的长上下文任务上均取得了有竞争力的性能，同时提升了推理效率（qwen35blog; kimiteam2025kimilinearexpressiveefficient; nvidia2025nvidianemotron3efficient）。从零开始大规模预训练强大的混合模型仍然成本过高，这促使近期研究转向通过 Transformer 到混合模型的转换或知识蒸馏来获得混合模型（chen2026hybridlinearattentionright; li2025distilling）。然而，现有工作主要聚焦于获取高性能的转换后混合模型检查点，而它们在后续后训练阶段的行为仍未被充分探索。
将 Transformer 蒸馏为混合模型。Transformer 到混合模型的蒸馏将选定的 softmax 注意力层转换为线性或循环混合器，其中层选择对长上下文检索性能至关重要（goldstein2026radladsrapidattentiondistillation; chen2026hybridlinearattentionright; li2025distilling; gu2026jet）。然而，转换后强大的召回性能并不一定意味着在面向推理的后训练后仍能保持稳定；具体而言，在多个具有强长上下文检索能力的蒸馏混合模型中，我们观察到 CoT-SFT 会显著降低召回能力，这揭示了一个现有蒸馏方法尚未解决的关键失效模式。
思维链动态与注意力路由。近期研究将思维链推理视为一种在中间推理状态上的结构化生成过程（wang2026doeschainofthoughthelpmarkovian; prystawski2023thinkstepstepreasoning）。马尔可夫视角通过将推理建模为潜在状态之间的局部转移来形式化这一结构，表明思维链监督主要强化了短程的步骤间依赖关系（wang2026doeschainofthoughthelpmarkovian）。与先前利用这种局部性来解释思维链在推理中有效性的工作不同，我们将其与高效长上下文模型中的一种失败模式联系起来：偏向局部性的思维链监督微调会侵蚀召回所需的长程注意力路由。这一观点与将注意力视为信息路由机制的分析相一致，其中查询-键交互定义了用于选择位置的路由逻辑值，而值投影则决定了从这些位置提取的内容（vaswani2023attentionneed）。
3 背景与问题设定
3.1 混合模型
在本工作中，我们主要关注混合模型。为了在提升长上下文效率的同时保留 softmax 注意力的表达能力，混合架构将少量 softmax 注意力层与高效的线性注意力层交错排列（yang2025zebra; chen2026hybridlinearattentionright）。
Softmax 注意力层 对于 Transformer 的第 层、第 个注意力头，每个头的投影矩阵（）对隐藏状态 进行变换：
| (1) | ||||
其中标量逻辑值 和注意力权重 为
| (2) |
每个头的输出向量 和该层的输出为
| (3) |
线性注意力层 对于第 层，一个循环状态矩阵 通过以下方式演化：
| (4) |
其中 、 是特定于层的更新和读取算子。以 Lightning Attention（qin2024lightningattention2freelunch）为例，状态更新和读取过程为：
| (5) | ||||
| (6) |
我们将 softmax 注意力层的集合记为 。实证研究表明，混合模型中长距离召回能力不成比例地依赖于一组 softmax 注意力层，而其余大部分层可以用循环机制替代，且性能下降极小（wang2025systematicanalysishybridlinear; chen2026hybridlinearattentionright; jelassi2024repeat）。因此，保留 的路由行为对于长上下文召回至关重要。
3.2 注意力路由
我们通过 softmax 之前的 logit 来实现注意力路由的操作化，该 logit 决定了位置 对位置 的注意力强度。由于路由模式完全由 决定，对这些矩阵的改动会直接改变长距离检索行为。为了理解 CoT-SFT 训练如何影响检索行为，我们考察了关于 QK 指标的梯度。 的梯度更新为：
| (7) |
如果 在小距离时很大，在大距离时可忽略不计，那么梯度下降会系统地推动 趋向局部模式，无论上下文长度如何。此外，我们有：
| (8) |
其中 是梯度向量。因此，理解 的距离依赖性成为理解 CoT-SFT 如何影响检索的核心。
4 为什么 CoT-SFT 会破坏路由？
4.1 CoT 数据假设
我们首先通过一个潜在马尔可夫模型（wang2026doeschainofthoughthelpmarkovian; prystawski2023thinkstepstepreasoning）来刻画 CoT 数据结构，这为后续的梯度分析提供了动机。
假设 4.1（CoT-马尔可夫结构）。
存在潜在推理状态 ， ，使得：
1. 马尔可夫转移： 。
2. 观测模型： 。位置 处的 token 由当前推理状态决定。
3. 遍历性：该链是不可约且非周期的，具有平稳分布 ， 。
4. 可逆性：细致平衡条件成立： 对所有 成立。
5. 谱间隙：转移矩阵的第二大特征值幅度为 。
我们在附录 A 中讨论该假设有效性的相关论述。
4.2 梯度局部性理论
Softmax 注意力层是混合模型中长距离召回的唯一场所。因此，我们分析 CoT-SFT 如何通过梯度局部性侵蚀这些层中的路由能力。
4.2.1 定义与假设
定义 4.2（距离条件梯度幅度）。
| (9) |
假设 4.3（范数边界）。
存在常数，使得在训练分布上几乎必然地，
| (10) |
假设 4.4（得分函数恒等式）。
设 表示损失函数关于位置 处每个注意力头输出的梯度。在交叉熵损失下的近似最优性条件下，模型满足：对于所有 以及 上几乎必然地，
| (11) |
上述假设的论证包含在附录 B 和附录 C 中。
4.2.2 梯度局部性定理
引理 4.5（谱相关性衰减）。
对于一条具有谱间隙 的平稳可逆遍历马尔可夫链，以及任意均值为零的 ：
| (12) |
定理 4.6（梯度局部性）。
设 为假设 4.3 中的常数， 为假设 4.1 中的谱半径。在假设 4.3–4.4 以及 CoT-马尔可夫模型（假设 4.1）下，对于所有 ：
| (13) |
| (14) |
常数 仅取决于模型的范数边界；衰减率完全由潜在转移矩阵的谱间隙决定。完整证明包含在附录 F 中。
4.3 实证验证
定理 4.6 表明，CoT 训练会优先强化邻近 token 之间的交互，导致远距离路由信号随距离增加而衰减。这一理论链条依赖于关于数据和模型的两个实证主张：
1. 引理 4.5 的假设：控制 token 生成的潜在马尔可夫链具有谱间隙 。然而，由于潜在状态不可观测，我们转而考虑观测到的 token 序列，并使用 token 自相关性作为 的代理指标。具体而言， 衡量的是某个 token 在距离 处重新出现的概率比随机情况高出多少，该值按所有 token 类型以其平稳频率 加权平均后得到。如附录 G 所证明，在隐马尔可夫模型下成立。
2. 定理 4.6 的预测：期望梯度幅度随 呈指数衰减，且 CoT 数据的有效窗口比通用文本更小。我们直接从模型中测量每个距离 处的 ：
| (15) |
3. 曲线拟合。两者均采用形式 ，其中 作为衰减的自然尺度。因此，我们通过对数线性回归来拟合模型。取对数可将模型线性化：
| (16) |
4. 解释。如图 3 所示，左子图展示了模型 token 自相关衰减参数 ，这是一个语料库级别的统计量，用于衡量文本保持结构化、非随机自相似性的特征距离。由于数学推理会反复引用相同的符号、变量和公式片段，思维链文本在更长距离上（ token）维持着比普通散文（ token）更高的自相似性。右子图展示了注意力梯度衰减参数 ，它通过衡量梯度信号 的有效作用范围来刻画训练动态。两条衰减曲线均能很好地用指数拟合来描述，这与定理的预测一致。
当 时，存在一个距离区间，在该区间内数据需要长程注意力，但训练梯度不再对其进行强化。关于预训练阶段也观察到的注意力梯度衰减的讨论，请参见附录 H。
5 方法
定理 4.6 已经确定，注意力 logit 上的梯度随距离呈几何级数衰减。现在我们研究这种衰减会传播到哪些权重指标中。
5.1 路由-提取梯度解耦
定理 5.1（路由-提取梯度解耦）
在思维链监督微调条件下，当 token 生成过程被建模为具有谱间隙 的平稳可逆遍历马尔可夫链时，对 softmax 注意力层参数的梯度更新表现出以下不对称性：
[路由] 来自距离为 的位置对的 的逐对贡献满足：
其中 为常数。 也满足相同界限。
[提取] 在任意位置 的值向量上的梯度满足：
其中 对所有 成立，且 对所有 成立。
详细证明见附录 I。因此，思维链监督微调对路由和知识提取的影响不同。 仅从短程位置对接收有意义的梯度信号，而到达 的梯度在所有位置上都有统一的下界，与上下文长度无关。
5.2 QK-恢复
定理 5.1 表明，CoT-SFT 以根本不同的方式影响 softmax 注意力层的两个功能组件。
路由参数（控制模型关注的位置）接收到的梯度信号随距离呈指数衰减：它们仅由短距离配对塑造，并逐渐丧失长距离路由能力。提取参数（控制检索内容）在每个位置接收到的梯度均具有一致的下界：它们会累积价值处理方面的改进。
换句话说，CoT-SFT 同时具有破坏性和改进性。这些效应被分离到不相交的参数集中。我们提出了 QK-Restore，该方法通过移植 SFT 前检查点的参数，同时保留 SFT 后的参数（算法 1），从而消除路由破坏并保留提取改进。
6 实验
6.1 实验设置
模型。HypeNet (chen2026hybridlinearattentionright) 近期展现了强大的长上下文召回性能，因此我们主要关注该模型在 2B 到 9B 规模上的表现。为确保分析的全面性，我们还包含了 Jet-Nemotron-2B111https://huggingface.co/collections/jet-ai/jet-nemotron。
训练。为研究 CoT-SFT 的影响，我们主要聚焦于数学领域，该领域研究广泛且拥有充足的高质量数据集 (ultradata-math; mitra2024orcamath; yu2025dapoopensourcellmreinforcement)。我们在预训练和 SFT 阶段自行训练了不同规模的 HypeNet，并基于 Jet-Nemotron 的预训练检查点对其应用 SFT。数据集和训练配置的更多细节见附录 J。
评估。我们评估模型在长上下文召回和数学推理任务上的表现。对于长上下文召回，我们报告在 NIAH 上的准确率。为衡量数学推理能力，我们在 GSM8K (cobbe2021gsm8k) 和 MATH500 (lightman2023letsverifystepstep) 上对模型进行测试。
| 模型 | 方法 | NIAH-Single-1 | NIAH-Single-2 | NIAH-Single-3 | MATH500 | GSM8K | AIME24 | ||||||||||||
| 32K | 64K | 128K | 256K | 32K | 64K | 128K | 256K | 32K | 64K | 128K | 256K | Avg@16 | Maj@16 | Avg@16 | Maj@16 | Avg@16 | Maj@16 | ||
| Jet-Nemotron-2B | 预训练 | 100.0 | 99.8 | 100.0 | 100.0 | 95.6 | 91.0 | 61.0 | 22.2 | 57.2 | 48.0 | 73.6 | 40.6 | 37.5 | 64.0 | 52.6 | 84.3 | 0.80 | 3.30 |
| +SFT | 100.0 | 100.0 | 100.0 | 98.6 | 57.8 | 28.4 | 21.0 | 11.6 | 34.4 | 33.6 | 67.2 | 27.8 | 49.4 | 70.6 | 71.9 | 88.0 | 1.50 | 3.30 | |
| +QK-Restore | 100.0 | 100.0 | 100.0 | 98.6 | 64.6 | 34.0 | 22.0 | 7.60 | 41.8(7.4) | 39.0(5.4) | 64.4(2.8) | 34.0(6.2) | 49.0 | 69.4 | 71.2 | 88.6 | 1.90 | 6.70 | |
| HypeNet-2B | 预训练 | 99.2 | 97.6 | 97.2 | 98.4 | 75.8 | 94.0 | 88.0 | 56.2 | 48.8 | 34.8 | 30.4 | 23.4 | 4.70 | 25.4 | 2.20 | 20.0 | 0.00 | 0.00 |
| +SFT | 99.8 | 99.8 | 99.2 | 99.6 | 82.2 | 81.8 | 48.4 | 41.8 | 47.4 | 32.6 | 7.60 | 10.6 | 34.4 | 55.2 | 40.8 | 68.1 | 0.00 | 0.00 | |
| +QK-Restore | 99.4 | 99.6 | 99.4 | 99.8 | 97.6 | 91.2 | 83.2 | 40.8 | 56.2(8.8) | 52.8(20.2) | 30.8(23.2) | 30.2(19.6) | 33.7 | 56.6 | 41.2 | 69.0 | 0.20 | 3.30 | |
| HypeNet-5B | 预训练 | 100.0 | 100.0 | 100.0 | 100.0 | 99.4 | 99.8 | 97.0 | 93.4 | 95.0 | 94.2 | 87.2 | 75.2 | 4.70 | 33.8 | 4.10 | 39.5 | 0.00 | 0.00 |
| +SFT | 100.0 | 100.0 | 100.0 | 100.0 | 98.8 | 100.0 | 91.4 | 83.6 | 96.6 | 96.6 | 90.6 | 65.4 | 49.3 | 67.0 | 61.3 | 81.4 | 0.60 | 0.00 | |
| +QK-Restore | 100.0 | 100.0 | 100.0 | 100.0 | 99.8 | 100.0 | 96.8 | 85.4 | 97.0(0.4) | 96.4(0.2) | 90.2(0.4) | 76.4(11.0) | 48.7 | 68.4 | 60.9 | 81.8 | 0.40 | 3.30 | |
| HypeNet-9B | 预训练 | 100.0 | 100.0 | 100.0 | 100.0 | 64.2 | 95.8 | 72.4 | 67.2 | 53.0 | 95.0 | 83.8 | 52.0 | 17.6 | 57.6 | 15.3 | 62.4 | 0.80 | 6.70 |
| +SFT | 100.0 | 100.0 | 100.0 | 100.0 | 89.0 | 64.2 | 17.4 | 9.40 | 98.2 | 90.0 | 47.8 | 22.8 | 39.8 | 69.5 | 62.3 | 87.8 | 1.20 | 10.0 | |
| +QK-Restore | 100.0 | 100.0 | 100.0 | 100.0 | 92.6 | 100.0 | 44.0 | 19.6 | 98.6(0.4) | 97.0(7.0) | 66.4(18.6) | 42.6(19.8) | 37.6 | 68.4 | 59.3 | 88.4 | 0.60 | 3.30 | |
6.2 主要结果
性能对比结果见表 1。
CoT-SFT 导致长上下文性能下降。一个关键发现是，CoT-SFT 通常会削弱模型在预训练阶段获得的长上下文检索能力，尤其是在更具挑战性的检索设置和更长的上下文窗口下。虽然像 NIAH-Single-1 这样的简单任务几乎仍能保持饱和性能，但在 NIAH-Single-2 和 NIAH-Single-3 上出现了显著的性能下降。
例如在表 1 中，HypeNet-2B 在 128K 上下文下的 NIAH-Single-3 任务上，经过 SFT 后从 [原文此处缺失数值] 下降至 [原文此处缺失数值]；而 HypeNet-9B 在 256K 上下文下的 NIAH-Single-2 任务上，从 [原文此处缺失数值] 下降至 [原文此处缺失数值]。此外，随着上下文长度增加，性能下降持续加剧，这表明 SFT 对负责长距离检索的机制造成了不成比例的破坏。
QK-Restore 恢复长上下文能力。QK-Restore 能够持续缓解监督微调（SFT）过程中产生的长上下文性能损失，在退化最严重的配置中观察到的提升幅度最大。
具体在表 1 中，对于 HypeNet-2B，在 256K 长度下的 NIAH-Single-3 任务上，性能从提升至（+19.6），而 HypeNet-9B 则从提升至（+19.8）。在 OpenThoughts-3 的链式思维监督微调（CoT-SFT）下，HypeNet-9B 在 256K 长度下的 NIAH-Single-2 任务上从提升至，在 NIAH-Single-3 任务上从提升至。在若干配置中，QK-Restore 甚至超越了原始的预训练基线。例如，HypeNet-2B 在 64K 长度下的 NIAH-Single-3 任务上，从预训练结果提升至。这表明我们的方法并非简单地恢复预训练行为；相反，它能够在保留预训练阶段检索结构的同时，仍然受益于监督微调（SFT）过程中引入的任务特定表示。
QK-Restore 保留下游推理性能。除了检索能力之外，我们还评估了恢复路由参数对下游推理能力的影响。在各项基准测试中，QK-Restore 在显著提升长上下文检索能力的同时，很大程度上保留了从链式思维监督微调（CoT-SFT）中获得的性能提升。
在数学推理任务上，性能保持接近监督微调（SFT）后的模型水平。例如，对于 HypeNet-5B，QK-Restore 在 MATH500 和 GSM8K 上分别仅带来 和 个百分点的微小变化，同时恢复了长上下文检索性能。在其他模型规模上也观察到了类似的趋势。这些结果表明，可以选择性地恢复路由参数以恢复长上下文检索能力，而不会显著影响后训练阶段获得的任务特定知识和推理能力。
检索难度加剧长上下文退化。我们进一步观察到，检索的复杂性对鲁棒性有强烈影响。NIAH-Single-1 在各训练阶段几乎保持饱和状态，而在 Single-2 和 Single-3 任务中则出现了显著更大的差距。这一趋势表明，难度更高的检索设置对精确的长距离 token 交互提出了更高要求，因此对退化现象更为敏感。
| 模型 | 方法 | NIAH-Single-1 | NIAH-Single-2 | NIAH-Single-3 | IFEval | |||||||||
| 32K | 64K | 128K | 256K | 32K | 64K | 128K | 256K | 32K | 64K | 128K | 256K | |||
| HypeNet-5B | 预训练 | 100.0 | 100.0 | 100.0 | 100.0 | 99.4 | 99.8 | 97.0 | 93.4 | 95.0 | 94.2 | 87.2 | 75.2 | 12.0 |
| +SFT | 99.8 | 100.0 | 100.0 | 100.0 | 99.8 | 100.0 | 99.6 | 93.6 | 98.6 | 97.2 | 85.4 | 83.8 | 26.6 | |
| +QK-Restore | 100.0 | 100.0 | 100.0 | 100.0 | 99.6 | 100.0 | 99.8 | 93.4 | 97.4 | 95.2 | 85.4 | 85.2 | 26.6 | |
| HypeNet-9B | 预训练 | 100.0 | 100.0 | 100.0 | 100.0 | 64.2 | 95.8 | 72.4 | 67.2 | 53.0 | 95.0 | 83.8 | 52.0 | 14.0 |
| +SFT | 100.0 | 100.0 | 100.0 | 100.0 | 86.4 | 97.8 | 88.2 | 78.4 | 96.2 | 96.8 | 95.8 | 86.8 | 27.2 | |
| +QK-Restore | 100.0 | 99.8 | 100.0 | 100.0 | 88.2 | 97.4 | 92.8 | 78.6 | 96.0 | 97.4 | 95.2 | 88.2 | 29.0 | |
7 分析
7.1 非 CoT SFT 分析
我们的理论分析将长上下文性能下降归因于 CoT 监督所引发的局部优化动态。一个自然的问题是：这种现象是后训练阶段的普遍现象，还是 CoT 风格推理轨迹所特有的？为了回答这个问题，我们使用 Tulu-3 数据集，将 CoT-SFT 与传统的指令遵循 SFT 进行了比较。
如表 2 所示，标准的指令微调并未表现出 CoT-SFT 下观察到的相同性能下降模式。在 HypeNet-5B 和 HypeNet-9B 上，经过 Tulu-3 训练后，长上下文检索性能基本得以保持，并且经常得到提升。例如，HypeNet-5B 在 NIAH-Single-3 的 256K 长度上从 [原文此处缺少数值] 提升至 [原文此处缺少数值]，而 HypeNet-9B 则从 [原文此处缺少数值] 提升至 [原文此处缺少数值]。在多个上下文长度下的 NIAH-Single-2 上也观察到了类似的提升。
此外，根据 IFEval 的衡量，经过 Tulu-3 微调后，指令遵循能力显著提升。尽管有这些改进，长上下文检索能力仍然保持完好。这与 CoT-SFT 的行为形成了鲜明对比——在相同的架构和上下文长度下，CoT-SFT 的检索性能出现了下降。
这些发现表明，长上下文性能下降并非后训练阶段的固有结果。相反，它与 CoT 风格的监督密切相关，这支持了我们的假设：推理轨迹的顺序依赖结构将梯度更新集中在短程交互上，并对长程路由行为产生了不成比例的影响。
| 模型 | 方法 | NIAH-Single-1 | NIAH-Single-2 | NIAH-Single-3 | LCB-V5 | |||||||||
| 32K | 64K | 128K | 256K | 32K | 64K | 128K | 256K | 32K | 64K | 128K | 256K | Avg@8 | ||
| HypeNet-5B | 预训练 | 100.0 | 100.0 | 100.0 | 100.0 | 99.4 | 99.8 | 97.0 | 93.4 | 95.0 | 94.2 | 87.2 | 75.2 | 0.27 |
| +SFT | 100.0 | 100.0 | 99.8 | 100.0 | 99.8 | 100.0 | 99.0 | 89.6 | 98.8 | 97.8 | 94.2 | 91.4 | 6.37 | |
| +QK-Restore | 100.0 | 99.8 | 99.8 | 100.0 | 99.8 | 100.0 | 100.0 | 92.8 | 98.6 | 97.4 | 93.2 | 79.4 | 6.92 | |
| HypeNet-9B | 预训练 | 100.0 | 100.0 | 100.0 | 100.0 | 64.2 | 95.8 | 72.4 | 67.2 | 53.0 | 95.0 | 83.8 | 52.0 | 1.66 |
| +SFT | 100.0 | 100.0 | 100.0 | 100.0 | 91.8 | 100.0 | 41.8 | 22.4 | 96.2 | 98.6 | 76.6 | 40.6 | 13.11 | |
| +QK-Restore | 100.0 | 100.0 | 100.0 | 99.8 | 94.2 | 99.8 | 78.2 | 45.0 | 95.0 | 97.4 | 90.8 | 61.2 | 12.91 | |
7.2 超越数学推理的分析
我们的主要实验聚焦于数学推理，因为该领域容易获得高质量的 CoT 监督信号。为了评估观察到的检索-推理权衡是否也存在于数学领域之外，我们利用 OpenThoughts-3 进行了额外实验，并在 LiveCodeBench-V5 上评估了下游编码性能。
如表 3 所示，在编码领域也观察到了相同的定性行为。CoT-SFT 显著提升了编码能力，使 HypeNet-5B 在 LiveCodeBench-V5 上的得分从 [原文缺失数值] 提升至 [原文缺失数值]，HypeNet-9B 从 [原文缺失数值] 提升至 [原文缺失数值]。然而，这些提升伴随着长上下文检索能力的明显下降，尤其是在具有挑战性的 NIAH 设置和较长上下文长度上。
应用 QK-Restore 可以在保持下游编码性能的同时，恢复大部分丢失的检索能力。例如，在 HypeNet-5B 上，LiveCodeBench-V5 得分进一步从 [原文缺失数值] 提升至 [原文缺失数值]，而 NIAH-Single-3 在 256K 长度下的得分从 [原文缺失数值] 恢复至 [原文缺失数值]。类似地，HypeNet-9B 在 NIAH-Single-3 的 256K 长度下得分从 [原文缺失数值] 提升至 [原文缺失数值]，同时保持了有竞争力的编码性能（[原文缺失数值]）。
这些结果表明，CoT-SFT 导致的检索能力下降并非数学推理所特有。相反，它似乎源于与长推理链相关的一种更普遍的优化现象。QK-Restore 在数学和编码两个领域均有效，这进一步支持了我们的核心假设：长距离检索主要由路由参数控制，而任务特定的推理能力则主要编码在网络的其它部分。
7.3 注意力路由恢复分析
在本节中，我们对注意力图进行可视化，以描述 CoT 监督微调如何改变 HypeNet-5B 中 softmax 注意力层的长程路由行为。我们基于在 Miromind 数据集上训练的模型进行分析。
具体而言，我们比较了 HypeNet-5B 在 SFT 前和 CoT-SFT 后两个检查点之间所有 softmax 注意力层的注意力矩阵。然后，我们通过每个注意力头的平均注意力距离指标来识别受影响最大的层：
| (17) |
其中 索引注意力头， 是从查询 到键 的 softmax 后权重。根据 SFT 后 值下降的头数量对层进行排序，第 33 层成为受影响最大的层，32 个头中有 25 个转向更短的注意力距离。因此，我们将分析重点放在这一层。
我们在一个包含 8,192 个 token 的合成 NIAH 提示词上可视化召回率最高的注意力头，目标事实嵌入在上下文深度中。对于每个注意力头，我们展示一个四面板图：SFT 前注意力矩阵（蓝色）、SFT 后矩阵（红色）、QK-Restore 矩阵（绿色）以及逐元素对数比率（面板 4），其中正值表示 QK-Restore 比 SFT 后恢复了更多长程注意力的位置。
分析。图4展示了HypeNet-5B中第33层第8个注意力头的模式。从图4中我们观察到，非对角线的下三角区域主要呈红色，这证实了QK-Restore部分恢复了被CoT微调抑制的长距离路由：对于大多数查询位置，QK-Restore对远距离键的关注程度高于微调后的模型。其次，右下角区域（查询和键位置均位于"针"之后）沿主对角线呈现蓝色带状，这意味着微调后模型在该区域对邻近token赋予了更高权重，这与我们的假设相符。这些观察结果证明了QK-Restore在缓解路由崩溃和恢复长距离注意力行为特征方面的有效性。
7.4 QK-Restore消融实验
QK-Restore将预微调检查点的完整矩阵移植到微调后模型中，大幅恢复了长距离召回能力，但导致数学性能下降。微调后模型可分解为：
其中 编码了有害的局部性漂移， 编码了有益的数学能力适应。由于这两个分量在 中纠缠在一起，QK-Restore将它们一并丢弃，在恢复路由的同时牺牲了数学收益。如图5和图6所示， 中出现了不可忽略的漂移，其中 。
因此，我们将其建模为一个Procrustes问题：
| (18) |
我们将其称为QK-Pro变体（更多细节见附录N）。
此外，我们评估了注意力中各组件（即 ）的作用，并将CoT-SFT与冻结 的预防性方案（记为QK-Frozen）进行对比，以回答一个问题：是在训练过程中预防路由漂移更好，还是在事后进行修正？
| 方法 | NIAH-Single-2 | NIAH-Single-3 | MATH500 | ||
| 64K | 128K | 64K | 128K | Avg@16 | |
| QK-Restore | 91.2 | 83.2 | 52.8 | 30.8 | 33.7 |
| QK-Pro | 94.0 | 79.0 | 51.8 | 26.0 | 34.2 |
| Q-Restore | 78.4 | 79.0 | 50.8 | 24.4 | 34.2 |
| K-Restore | 90.4 | 71.6 | 46.2 | 13.2 | 34.2 |
| V-Restore | 57.2 | 36.6 | 24.0 | 15.2 | 32.1 |
| QK-Frozen | 90.6 | 76.4 | 49.4 | 23.6 | 33.1 |
分析。表 4 验证了三个关键主张。V-Restore 无法恢复 NIAH，同时导致 MATH500 性能下降，证实其无助于路由恢复。Q-Restore 和 K-Restore 各自实现了部分召回恢复，但在 128K 长度下显著不足（例如，K-Restore 在 NIAH-Single-3 上崩溃，得分 13.2 对比 30.8），表明孤立恢复引入了不匹配，且联合恢复对于连贯的路由几何结构是必要的。QK-Pro 在 MATH500 上取得了最高性能，但以牺牲长程路由恢复来换取数学能力保留。QK-Frozen 在 NIAH 和数学任务上始终不如 QK-Restore，表明无约束的 SFT 需要更好的适应性。
8 结论
在这项工作中，我们将 CoT-SFT 引发的长上下文召回退化识别为蒸馏混合模型中的一个关键后训练失败模式。我们表明，面向推理的 SFT 会使优化偏向短程注意力模式，扰乱保留的 softmax 注意力层的查询-键路由几何结构，并削弱长程检索能力。受此观察启发，我们提出了 QK-Restore，一种无需训练的方法，该方法仅从 SFT 前的检查点恢复查询和键投影，同时保留所有其他参数。跨多种架构的实验表明，QK-Restore 在基本保留 SFT 推理增益的同时，显著恢复了长上下文能力。总体而言，我们的发现强调了路由稳定性是构建同时支持检索与推理的高效长上下文模型的关键因素。
参考文献
附录 A 关于 CoT-马尔可夫结构假设的讨论
讨论 1。1 步马尔可夫条件是为了分析而做的理想化处理。任意步的马尔可夫链等价于在状态乘积空间上的 1 步链，产生相同的指数衰减，但推理范围更大。
讨论 2. 可逆性在引理 4.5 的证明中被用于其中一步：它意味着马尔可夫算子在 上是自伴的。我们承认，因果生成的序列（例如 CoT 推理链）通常不满足细致平衡条件：其潜在转移概率在构造上就是不对称的。然而，可逆性是一个充分条件而非必要条件。对于一般的遍历链（可能不可逆），相同的界仍然成立，只需将 替换为 的第二大奇异值，该值通过其奇异值分解而非特征分解获得。由于奇异值始终是实数且非负，证明过程无需修改即可进行，最终界的形式也保持不变。因此，我们在引理 4.5 中假设可逆性是为了表述清晰，同时理解所有后续结果通过将特征值替换为奇异值后，对一般遍历情况同样成立。
附录 B 范数界假设的论证
对 的界。值向量为 ，其中 表示层归一化。层归一化对每个坐标进行标准化，使得 ，其中 和 是 各分量的经验均值和标准差。归一化后，仿射变换前的向量满足 ，因此得到
| (19) | ||||
| (20) | ||||
| (21) |
其中 是一个仅依赖于学习到的缩放参数 和偏置 的常数。投影矩阵 是一个固定的训练权重；其算子范数是一个有限的、依赖于模型的常数。综合以上，我们得到：
| (22) | ||||
| (23) |
对 的界。梯度通过最后一层层归一化（其雅可比矩阵为 ）和语言模型头 反向传播，得到
| (24) |
其中 是预测分布， 是独热目标。根据算子范数的次可乘性：
| (25) |
由于 是概率向量的差，因此 。对于一个固定的训练模型， 和 都是有限常数。因此：
| (26) |
附录 C 得分函数恒等式假设的论证
在最优解 处，关于注意力输出的交叉熵损失对于任意固定上下文 满足：
| (27) | |||
| (28) | |||
| (29) |
由于交叉熵梯度在其最小值处消失。应用全期望公式可得。更精确地说，因为梯度是在给定上下文下计算的，然后对从条件分布中抽取的目标取平均，该等式有条件地成立：
| (30) |
附录 D 预备定义与引理
我们列出证明所需的以下定义和引理。
定义 D.1（空间）
对于具有有限状态空间和平稳分布的马尔可夫链，定义加权函数空间：
| (31) |
带有内积
| (32) | ||||
| (33) |
以及范数。
此外，均值为零意味着，即该函数在下的期望为零。
定义 D.2（作用于函数的算子）
我们定义作用于函数的马尔可夫算子：
| (34) | ||||
| (35) |
引理 D.3
对于可逆链，细致平衡条件成立：，这使得在中是对称的。
证明
根据 D.2 中马尔可夫算子的定义，我们有：
| (36) | ||||
| (37) | ||||
| (38) | ||||
| (39) | ||||
| (40) |
由于，自伴算子的有限维谱定理给出了一个实的标准正交特征基，其特征值为，谱间隙为，其中：
| (41) |
进一步，对于均值为零的：
| (42) |
这意味着沿方向没有分量。其展开仅使用：
| (43) |
附录 E 谱相关衰减引理的证明
证明
在平稳状态下（），塔性质表明：
| (44) |
注意内部期望是步预测：对于固定的起始状态，经过步后的期望值恰好是：
| (45) |
因此：
| (46) | ||||
| (47) | ||||
| (48) |
然后我们通过标准正交特征基展开：
| (49) |
利用标准正交性，我们进一步简化：
| (50) | ||||
| (51) |
然后对其进行约束：
| (52) | ||||
| (53) | ||||
| (54) |
应用柯西-施瓦茨不等式：
| (55) | ||||
| (56) |
现在将所有结果合并：
| (57) | ||||
| (58) |
附录 F 梯度局部性定理的证明
证明
为清晰起见，我们在以下推导中省略层/头的上标。首先，对梯度进行约束：
| (59) |
然后我们考虑如何约束。应用塔性质：
| (60) | ||||
| (61) |
其中我们定义：
| (62) |
根据假设 4.4：
| (63) | ||||
| (64) | ||||
| (65) | ||||
| (66) |
类似地，中心化不会改变协方差（因为）。和都是上的均值为零的函数。
然后应用引理 4.5：
| (67) | ||||
| (68) | ||||
| (69) |
对项应用相同的约束，该项满足作为凸组合：
| (70) |
其中且。∎
附录 G：模型 token 自相关衰减的证明
定义中心化指示特征以及滞后 处的 -加权自相关：
| (71) |
对于分子 ，我们将其展开为：
| (72) |
考虑 的期望：
| (73) | ||||
| (74) |
因此我们可以进一步简化：
| (75) |
在隐马尔可夫模型（HMM）中，模型 token 在给定其潜在状态的情况下是条件独立的： 。根据塔性质：
| (76) | |||
| (77) | |||
| (78) |
其中 是 的一个均值为零的函数：
| (79) | ||||
| (80) | ||||
| (81) |
由于 和 都是均值为零的，我们可以直接应用引理 4.5：
| (82) |
其中 是潜在转移矩阵 的第二大特征值。现在我们可以限定自相关的范围：
| (83) | ||||
| (84) |
结论。在语料库中观察到 且 证实了潜在链具有谱间隙 ，这恰好是引理 4.5 的假设条件。
附录 H：预训练阶段注意力梯度衰减的讨论
定理 4.6 基于引理 4.5，该引理适用于任何具有谱间隙的遍历马尔可夫链：这是数据生成过程的一个属性，并非 CoT 结构所特有，在一般的预训练数据语料库中也能观察到。
因此，梯度局部性界限适用于任何训练语料库，包括一般的预训练数据。路由坍缩的严重程度由失配程度 决定：即数据要求模型关注的距离与梯度实际强化该关注的距离之间的差异。对于一般的预训练语料库，这种失配程度小于 CoT 语料库，因此路由退化得更慢，但机制是相同的。
然而，对于混合架构来说，即使是预训练阶段较慢的侵蚀，也会在微调开始之前显著降低可用的召回容量，使模型更接近阈值。因此，混合架构的脆弱性源于预训练阶段，并会因 CoT 微调而进一步加剧。
附录 I：路由-提取梯度解耦定理的证明
证明。
根据公式 7，我们有：
| (85) |
其中 是对应的 KV 头。然后考虑其 Frobenius 范数：
| (86) | ||||
| (87) | ||||
| (88) |
通过定理 4.6 以及有界范数 、 。
对于提取参数，值向量通过 进入损失函数。累积所有下游依赖关系：
| (89) |
现在我们假设对于所有 和所有 ：
| (90) |
附录 J：训练数据与配置详情
对于 HypeNet，我们同时进行了预训练和 SFT。
预训练数据。由于我们主要关注长上下文召回和数学推理能力，在预训练阶段，我们混合使用了 FineWeb-Edu [lozhkov2024fineweb-edu] 和 UltraData-Math [ultradata-math]。根据近期预训练研究的观察结果——即保持网络语料库为主要来源，同时增加结构化推理数据 [allal2025smollm2smolgoesbig]——我们构建了一个包含通用网络数据和数学专注数据的混合数据集。先前的证据表明，网络规模的语料库能保留广泛的语言和事实能力，而专门的数学语料库则能提升推理和 STEM 表现。
预训练配置。我们对 HypeNet 采用三阶段训练，保持与 chen2026hybridlinearattentionright 中相同的设置，如表 5、表 6 和表 7 所示。
| 阶段 | 模型 token | 学习率 | 学习率调度器 | 上下文长度 | 批次大小 | 训练步数 |
| 1 | 3.2 亿 | 1e-3 至 1e-5 | 余弦 | 512 | 32 | 20,000 |
| 2 | 10 亿 | 1e-4 至 1e-5 | 余弦 | 512 | 96 | 20,000 |
| 3 | 10 亿 | 1e-5 | 常数 | 16,384 | 128 | 500 |
| 阶段 | 模型 token | 学习率 | 学习率调度器 | 上下文长度 | 批次大小 | 训练步数 |
| 1 | 3.2 亿 | 1e-3 至 1e-5 | 余弦 | 512 | 32 | 20,000 |
| 2 | 10 亿 | 5e-5 至 1e-5 | 余弦 | 512 | 96 | 20,000 |
| 3 | 10 亿 | 1e-5 | 常数 | 16,384 | 128 | 500 |
| 阶段 | 模型 token | 学习率 | 学习率调度器 | 上下文长度 | 批次大小 | 训练步数 |
| 1 | 3.2 亿 | 1e-3 至 1e-5 | 余弦 | 512 | 32 | 20,000 |
| 2 | 10 亿 | 3e-5 至 1e-5 | 余弦 | 512 | 96 | 20,000 |
| 3 | 10 亿 | 1e-5 | 常数 | 16,384 | 128 | 500 |
SFT 数据。为了进一步提升模型的数学推理性能，我们使用 MiroMind-M1 数据集 [li2025miromindm1opensourceadvancementmathematical] 进行 CoT-SFT，该数据集收集自 OpenR1 [openr1]、Open-thoughts [guha2025openthoughts]、Light-R1 [wen2025light] 和 Synthetic-1 [2025synthetic1]。SFT 配置。详细配置如表 8 和表 9 所示。
附录 K 评估细节
在评估 NIAH 时，我们使用 LM Evaluation Harness222https://github.com/EleutherAI/lm-evaluation-harness [eval-harness] 进行官方测试。在评估 MATH500 和 GSM8K 时，我们报告每个问题在 16 次独立生成上的平均 pass@1 以及 Maj@16（即多数投票作为最终预测）。具体来说，对于 MATH500，生成长度为 8,192；对于 GSM8K，我们将其设置为 2,048，因为该任务相对简单。对于 LiveCodeBench，我们报告在 16,384 响应长度下 8 次独立生成的平均 pass@1。为了衡量模型在 Tulu3 上微调后的性能，我们报告 IFEval 任务 [zhou2023instructionfollowingevaluationlargelanguage] 上的提示词级别严格准确率。
附录 L 纯 Softmax 注意力模型分析
我们的主要实验集中在混合模型上，在本节中，我们探讨 SFT 是否也会降低纯 Softmax 注意力模型的性能。我们选取了 Qwen2.5-3B、Qwen2.5-7B [qwen2025qwen25technicalreport] 和 Mistral-7B-Instruct-v0.3 [jiang2023mistral7b] 进行研究。关于如何识别用于检索的顶层，我们在附录 M 中提供了详细信息。
| 模型 | 方法 | NIAH-Single-2 | NIAH-Single-3 | ||
| 32K | 64K | 32K | 64K | ||
| Qwen2.5-7B | 预训练 | 100.0 | 95.8 | 99.8 | 98.6 |
| +SFT | 100.0 | 92.6 | 99.8 | 99.4 | |
| Qwen2.5-3B | 预训练 | 100.0 | 94.6 | 100.0 | 99.6 |
| +SFT | 100.0 | 93.8 | 100.0 | 90.8 | |
| +QK-Restore | 100.0 | 93.8 | 100.0 | 93.4 | |
| Mistral-7B | 预训练 | 99.8 | 90.2 | 99.8 | 81.0 |
| +SFT | 99.6 | 59.6 | 65.8 | 4.80 | |
| +QK-Restore | 99.8 | 62.0 | 70.0 | 5.40 | |
分析。表 10 显示，对于 Mistral-7B，在较长上下文中，它对 SFT 表现出强烈的敏感性。SFT 导致 NIAH-Single-3 在 64K 处严重退化，从 下降到 ，表明长距离检索能力崩溃。QK-Restore 对这种退化的恢复效果有限，表明预训练模型中的长上下文表示较弱。
对于 Qwen2.5-7B 和 Qwen2.5-3B，长上下文性能在 NIAH-Single-2 和 Single-3 任务上高度稳定。SFT 仅引起微小波动（例如，Qwen2.5-7B 在 64K 处从某值下降至某值），这表明这些模型在预训练期间已经学会了稳健的长上下文检索。因此，QK-Restore 仅带来边际收益，这与能力几乎未损失的接近饱和状态一致。为了探究其底层机制，我们测量了 Qwen2.5-3B 上的注意力梯度衰减，如图 7 所示。与 HypeNet-2B（其全注意力层呈现拟合良好的指数衰减）不同，该模型具有明显的近平坦尾部。我们假设，在纯 softmax 注意力模型中，长上下文召回是通过丰富的全注意力层在结构上得以维持的。
附录 M 纯 Softmax 注意力模型中检索层的识别方法
为了识别哪些 softmax 注意力层主要负责长距离检索，我们在集合上进行了留一法消融研究。对于每一层，我们通过设置来构建一个消融模型。每个消融模型在多个检索基准（Ruler-QA-SQuAD [hsieh2024ruler]、FDA [arora2023language]、SWDE [lockard2019openceres]）上进行评估。层的重要性得分定义为：
其中表示各基准上的平均检索准确率。较大的值表明该层对检索贡献不成比例地大，且其路由权重至关重要；接近零的得分则表示冗余。各层按降序排列，选择前个层作为权重恢复的目标集合。在我们的实验中，我们保持与 HypeNet 中相同。
附录 N QK-Pro 推导的详细过程
双线性约束没有封闭形式的联合解，因此我们通过将其中一个因子固定为其 SFT 前的值来进行线性化，使剩余问题成为可通过拉格朗日乘数法求解的标准约束最小二乘系统。
我们固定，求最小偏差：
| (91) | |||
| (92) |
然后应用拉格朗日乘数法：
| (93) |
求导，得到：
| (94) |
Xinyu Zhou
LARK, HKUST(GZ)
Boyu Zhu
UCL
Yi Xu
Mistral AI
Zhiwei Li
LARK, HKUST(GZ)
Yingfa Chen
Huiming Wang
SUTD
Zhijiang Guo
LARK, HKUST(GZ)
Abstract
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 to . We attribute this to CoT-SFT biasing attention gradients toward short-range patterns, disrupting query-key projections () that are responsible for long-range routing. Motivated by this observation, we propose QK-Restore, a training-free method that restores only and 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 to while maintaining strong reasoning performance.
Abstract
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 to . We attribute this to CoT-SFT biasing attention gradients toward short-range patterns, disrupting query-key projections () that are responsible for long-range routing. Motivated by this observation, we propose QK-Restore, a training-free method that restores only and 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 to while maintaining strong reasoning performance.
1 Introduction
Efficient sequence models, such as linear attention (katharopoulos2020transformersrnnsfastautoregressive; qin2024variouslengthsconstantspeed; yang2024gatedlinearattentiontransformers) and state-space models (gu2024mambalineartimesequencemodeling; dao2024transformersssmsgeneralizedmodels), have emerged as attractive alternatives to softmax-attention (vaswani2023attentionneed) for long-context processing, reducing the quadratic cost to linear complexity through compact recurrent or structured states. However, this compression introduces an information bottleneck on recall-intensive tasks that are crucial for long-context understanding (zhang2024hedgehogporcupineexpressive), such as Needle-in-a-Haystack (NIAH; hsieh2024ruler). Hybrid linear-attention models mitigate this tradeoff by retaining a small set of softmax-attention layers for global recall while converting the remaining layers into efficient linear-attention layers (nvidia2025nvidianemotron3efficient). Although native hybrid models such as Qwen3.5 (qwen35blog) and Kimi Linear (kimiteam2025kimilinearexpressiveefficient) achieve strong performance and efficiency, pretraining them from scratch remains highly resource-intensive. This cost motivates recent work that converts pretrained softmax-attention into hybrid linear-attention models through distillation, using only a small fraction of pretraining tokens (chen2026hybridlinearattentionright; li2025distilling; hoshino2025radredundancyawaredistillationhybrid). Existing conversion work primarily focuses on preserving pretraining capabilities during architectural conversion, while the impact of downstream reasoning post-training on long-context retrieval remains underexplored.
While distilled hybrid models primarily serve as efficient base models, modern LLMs typically rely on post-training to acquire stronger instruction-following and reasoning abilities (Guo_2025; qwen2025qwen25technicalreport; ouyang2022traininglanguagemodelsfollow). In particular, CoT-SFT is widely used to enhance mathematical and multi-step reasoning(wei2023chainofthoughtpromptingelicitsreasoning; li2025tl). We observe that applying CoT-SFT on mathematical reasoning data improves reasoning performance, but can substantially degrade the long-context recall ability that distilled hybrid models acquire during pretraining or architectural conversion, particularly in more challenging NIAH settings and extended contexts length (Figure 1). Unlike generic catastrophic forgetting, this degradation is highly structured: it primarily affects long-range retrieval behavior mediated by the retained softmax-attention layers, while leaving the intended reasoning gains largely intact. This reveals a fundamental tension in hybrid models: CoT-SFT strengthens local reasoning can simultaneously disrupt the routing mechanisms required for long-range retrieval.
To investigate this phenomenon, we model CoT reasoning traces as a latent Markov process over intermediate reasoning states, reflecting the local step-to-step structure of mathematical derivations (wang2026doeschainofthoughthelpmarkovian; prystawski2023thinkstepstepreasoning). Under this CoT-Markov assumption, we derive that the expected gradient magnitude on attention logits decays exponentially with token distance. We empirically validate this prediction by measuring token autocorrelation and gradients of attention-logit. These results suggest that CoT-SFT can improve local multi-step reasoning while eroding the long-range routing behavior needed for retrieval.
To better isolate the source of recall degradation, we analyze the retained softmax-attention layers through a routing-extraction decomposition. We find that CoT-SFT induces locality-biased drift in the query-key projections , which determine the source of information retrieval, whereas value-side extraction can still benefit from post-SFT adaptation. Since long-context recall in distilled hybrid models depends heavily on a small number of retained softmax-attention layers, such query-key drift can disproportionately disrupt retrieval. Motivated by this asymmetry, we propose QK-Restore, a training-free method that restores only and in these layers from the pre-SFT checkpoint while preserving all other post-SFT parameters, thereby recovering long-range routing while retaining most reasoning improvements from CoT-SFT. Our contributions can be summarized as:
-
We identify CoT-SFT-induced recall degradation as a structured post-training failure mode of distilled hybrid models.
-
We provide a theoretical and empirical analysis showing that CoT-SFT concentrates training signals on local token interactions and induces locality-biased drift in the query-key routing geometry of retained softmax-attention layers.
-
We introduce QK-Restore, a training-free method that restores only query and key projections, recovering long-context recall while largely preserving reasoning performance.
2 Related Work
Efficient Long-Context Models and Hybrid Attention. Recent hybrid attention models interleave softmax-attention layers with efficient recurrent or linear-attention layers, achieving competitive performance on both commonsense reasoning and recall-intensive long-context tasks while improving inference efficiency (qwen35blog; kimiteam2025kimilinearexpressiveefficient; nvidia2025nvidianemotron3efficient). Pretraining strong hybrid models from scratch at large scale remains prohibitively expensive, which motivates recent efforts to obtain hybrid models through Transformer-to-hybrid conversion or distillation (chen2026hybridlinearattentionright; li2025distilling). However, existing work primarily focus to obtaining high-performing converted hybrid checkpoints, while their behavior during subsequent post-training remains underexplored.
Distilling Transformers into Hybrid Models. Transformer-to-hybrid distillation converts selected softmax-attention layers into linear or recurrent mixers, where layer selection critically influences long-context retrieval performance (goldstein2026radladsrapidattentiondistillation; chen2026hybridlinearattentionright; li2025distilling; gu2026jet). However, strong recall performance after conversion does not necessarily imply stability after reasoning-oriented post-training; specifically, across multiple distilled hybrid models with strong long-context retrieval, we observe that CoT-SFT can substantially degrade recall capabilities, revealing a critical failure mode that remains unaddressed by existing distillation methods.
CoT Dynamics and Attention Routing. Recent work has studied CoT reasoning as a structured generation process over intermediate reasoning states (wang2026doeschainofthoughthelpmarkovian; prystawski2023thinkstepstepreasoning). A Markovian view formalizes this structure by modeling reasoning as local transitions between latent states, suggesting that CoT supervision primarily reinforces short-range step-to-step dependencies (wang2026doeschainofthoughthelpmarkovian). In contrast to prior work that utilizes this locality to explain the efficacy of CoT in reasoning, we connect it to a failure mode in efficient long-context models: locality-biased CoT-SFT can erode long-range attention routing required for recall. This perspective aligns with analyses of attention as an information-routing mechanism, where query-key interactions define routing logits that select positions, while value projections determine the content extracted from those positions (vaswani2023attentionneed).
3 Background and Problem Setup
3.1 Hybrid Model
In this work, we primarily focus on the hybrid model. To retain the expressive power of softmax-attention while improving long-context efficiency, hybrid architectures interleave a small set of softmax-attention layers with efficient linear-attention layers (yang2025zebra; chen2026hybridlinearattentionright).
Softmax-Attention Layers For Transformer’s layer , head , the per-head projection matrices () transform hidden state :
| (1) | ||||
with scalar logit and attention weight
| (2) |
The per-head output vector and layer output are
| (3) |
Linear-Attention Layers For layer , a recurrent state matrix evolves via
| (4) |
where , are layer-specific update and readout operators. Taking Lightning Attention (qin2024lightningattention2freelunch) as an example, the state update and readout are:
| (5) | ||||
| (6) |
We denote the set of softmax-attention layers as . Empirical studies suggest that long-range recall in hybrid models depends disproportionately on a set of softmax-attention layers, whereas most remaining layers can be replaced by recurrent mechanisms with minimal degradation (wang2025systematicanalysishybridlinear; chen2026hybridlinearattentionright; jelassi2024repeat). Thus, preserving the routing behavior of is critical for long-context recall.
3.2 Attention Routing
We operationalize attention-routing through the pre-softmax logit , which determines how strongly position attends to position . Since the routing pattern is determined solely by , changes to these matrices directly alter long-range retrieval behavior. To understand how CoT-SFT training affects retrieval behavior, we examine the gradients w.r.t QK metrics. The gradient update of is:
| (7) |
If is large for small distance and negligible for large , gradient descent systematically pushes toward local patterns regardless of context length. Furthermore, we have:
| (8) |
where is the gradient vector. So understanding the distance dependence of becomes central to understanding how CoT-SFT affects retrieval.
4 Why CoT-SFT Disrupts Routing?
4.1 CoT Data Assumption
We first characterize CoT data structure via a latent Markov model (wang2026doeschainofthoughthelpmarkovian; prystawski2023thinkstepstepreasoning), which motivates the gradient analysis that follows.
Assumption 4.1(CoT-Markov structure).
There exist latent reasoning states , , such that:
1. Markov transitions: .
2. Observation model: . The token at position is determined by the current reasoning state.
3. Ergodicity: The chain is irreducible and aperiodic with stationary distribution , .
4. Reversibility: Detailed balance holds: for all .
5. Spectral gap: The transition matrix has second-largest eigenvalue magnitude .
We include the discussion of the validity of this assumption in Appendix A.
4.2 Gradient Locality Theory
Softmax-attention layers are the sole locus of long-range recall in a hybrid model. Therefore, we analyze how CoT-SFT erodes their routing capacity in these layers through gradient locality.
4.2.1 Definition and Assumptions
Definition 4.2(Distance-conditioned gradient magnitude).
| (9) |
Assumption 4.3(Norm bounds).
There exist constants such that, almost surely over the training distribution,
| (10) |
Assumption 4.4(Score-function identity).
Let denote the gradient of the loss with respect to the per-head output at position . At approximate optimality under cross-entropy, the model satisfies, for all and almost surely over ,
| (11) |
The justifications for the above assumptions are included in Appendix B and Appendix C.
4.2.2 Gradient Locality Theorem
Lemma 4.5(Spectral Correlation Decay).
For a stationary reversible ergodic Markov chain with spectral gap , and any mean-zero :
| (12) |
Theorem 4.6(Gradient locality).
Let be the constants from Assumption 4.3, and let be the spectral radius from Assumption 4.1. Under Assumptions 4.3–4.4 and the CoT-Markov model (Assumption 4.1), for all :
| (13) |
| (14) |
The constant depends only on the model’s norm bounds; the decay rate is set entirely by the spectral gap of the latent transition matrix. The full proof is included in Appendix F.
4.3 Empirical Validation
Theorem 4.6 shows that CoT training preferentially reinforces nearby token interactions, causing long-range routing signals to decay with distance. The theoretical chain rests on two empirical claims about data and model:
1. Lemma 4.5’s hypothesis: the latent Markov chain governing token generation has a spectral gap . However, since the latent states are unobservable, we instead consider the observed token sequence and use the token autocorrelation as a proxy for . Concretely, measures how much more likely a token is to reappear at distance than by chance, averaged over all token types weighted by their stationary frequency . As proved in Appendix G, under the hidden markov model.
2. Theorem 4.6’s prediction: the expected gradient magnitude decays exponentially in , with a smaller effective window for CoT data than for general text. We measure at each distance directly from the model:
| (15) |
3. Curve Fitting. Both and take the form , where serves as a natural scale of the decay. Therefore, we fit the model via log-linear regression. Taking logarithms linearizes the model:
| (16) |
4. Interpretation. As shown in Figure 3, the left subfigure presents the token autocorrelation decay parameter , a corpus-level statistic that measures the characteristic distance over which text retains structured, non-random self-similarity. Because mathematical reasoning repeatedly references the same symbols, variables, and formula fragments, CoT text sustains high self-similarity over longer distances ( tok) than general prose ( tok). The right subfigure shows the attention gradient decay parameter , which characterises the training dynamics by measuring the effective reach of the gradient signal . Both decay curves are well described by exponential fits, consistent with the thoerem’s prediction.
When , there exists a band of distances where the data demands long-range attention but the training gradient no longer reinforces it. The discussions about the attention gradient decay also observed in the pre-training stage are included Appendix H.
5 Method
Theorem 4.6 has established that the gradient on the attention logit decays geometrically in distance . Now we investigate which weight metrics this decay propagates into.
5.1 Routing–Extraction Gradient Decoupling
Theorem 5.1(Routing-Extraction Gradient Decoupling).
Under CoT-SFT, where the token-generation process is modeled as a stationary reversible ergodic Markov chain with a spectral gap , gradient updates to the parameters of softmax-attention layer exhibit the following asymmetry:
[Routing] The per-pair contribution to from position pairs at distance satisfies:
for constant . The same bound holds for .
[Extraction] The gradient on the value vector at any position satisfies:
where satisfies for all and satisfies for all .
The detailed proof is in Appendix I. Therefore, CoT-SFT affects routing and knowledge extraction differently. receive meaningful gradient signal only from short-range pairs, while the gradient reaching is bounded below uniformly over all positions, independent of context length.
5.2 QK-Restore
Theorem 5.1 shows that CoT-SFT affects the two functional components of a softmax-attention layer in fundamentally different ways.
The routing parameters and , which control where the model attends, receive gradient signal that decays as with distance: they are shaped exclusively by short-range pairs and progressively lose long-range routing capacity. The extraction parameters and , which control what is retrieved, receive a gradient uniformly bounded below for every position: they accumulate value-processing improvements.
In other words, CoT-SFT simultaneously corrupts and improves . These effects are segregated into disjoint parameter sets. We propose QK-Restore , which eliminates routing corruption and preserves extraction improvement by transplanting from the pre-SFT checkpoint while retaining post-SFT and (Algorithm 1).
6 Experiments
6.1 Setup
Models. HypeNet (chen2026hybridlinearattentionright) has recently presented strong long-context recall performance, therefore we focus mainly on this model at scales from 2B to 9B. To ensure a comprehensive analysis, we also include Jet-Nemotron-2B111https://huggingface.co/collections/jet-ai/jet-nemotron.
Training. To investigate the influence of CoT-SFT, we primarily focus on the math domain, which is widely studied and has sufficient high-quality datasets (ultradata-math; mitra2024orcamath; yu2025dapoopensourcellmreinforcement). We train HypeNet in different scales on our own in both pre-training and SFT stage, and apply SFT to Jet-Nemotron from their pre-trained checkpoints. More details of datsets and training configuration are in Appendix J.
Evaluation. We evaluate the model’s performance on long-context recall and math reasoning tasks. For long-context recall, we report the accuracy on NIAH. To measure math reasoning, we test the model on GSM8K (cobbe2021gsm8k) and MATH500 (lightman2023letsverifystepstep).
| Model | Method | NIAH-Single-1 | NIAH-Single-2 | NIAH-Single-3 | MATH500 | GSM8K | AIME24 | ||||||||||||
| 32K | 64K | 128K | 256K | 32K | 64K | 128K | 256K | 32K | 64K | 128K | 256K | Avg@16 | Maj@16 | Avg@16 | Maj@16 | Avg@16 | Maj@16 | ||
| Jet-Nemotron-2B | Pre-train | 100.0 | 99.8 | 100.0 | 100.0 | 95.6 | 91.0 | 61.0 | 22.2 | 57.2 | 48.0 | 73.6 | 40.6 | 37.5 | 64.0 | 52.6 | 84.3 | 0.80 | 3.30 |
| +SFT | 100.0 | 100.0 | 100.0 | 98.6 | 57.8 | 28.4 | 21.0 | 11.6 | 34.4 | 33.6 | 67.2 | 27.8 | 49.4 | 70.6 | 71.9 | 88.0 | 1.50 | 3.30 | |
| +QK-Restore | 100.0 | 100.0 | 100.0 | 98.6 | 64.6 | 34.0 | 22.0 | 7.60 | 41.8(7.4) | 39.0(5.4) | 64.4(2.8) | 34.0(6.2) | 49.0 | 69.4 | 71.2 | 88.6 | 1.90 | 6.70 | |
| HypeNet-2B | Pre-train | 99.2 | 97.6 | 97.2 | 98.4 | 75.8 | 94.0 | 88.0 | 56.2 | 48.8 | 34.8 | 30.4 | 23.4 | 4.70 | 25.4 | 2.20 | 20.0 | 0.00 | 0.00 |
| +SFT | 99.8 | 99.8 | 99.2 | 99.6 | 82.2 | 81.8 | 48.4 | 41.8 | 47.4 | 32.6 | 7.60 | 10.6 | 34.4 | 55.2 | 40.8 | 68.1 | 0.00 | 0.00 | |
| +QK-Restore | 99.4 | 99.6 | 99.4 | 99.8 | 97.6 | 91.2 | 83.2 | 40.8 | 56.2(8.8) | 52.8(20.2) | 30.8(23.2) | 30.2(19.6) | 33.7 | 56.6 | 41.2 | 69.0 | 0.20 | 3.30 | |
| HypeNet-5B | Pre-train | 100.0 | 100.0 | 100.0 | 100.0 | 99.4 | 99.8 | 97.0 | 93.4 | 95.0 | 94.2 | 87.2 | 75.2 | 4.70 | 33.8 | 4.10 | 39.5 | 0.00 | 0.00 |
| +SFT | 100.0 | 100.0 | 100.0 | 100.0 | 98.8 | 100.0 | 91.4 | 83.6 | 96.6 | 96.6 | 90.6 | 65.4 | 49.3 | 67.0 | 61.3 | 81.4 | 0.60 | 0.00 | |
| +QK-Restore | 100.0 | 100.0 | 100.0 | 100.0 | 99.8 | 100.0 | 96.8 | 85.4 | 97.0(0.4) | 96.4(0.2) | 90.2(0.4) | 76.4(11.0) | 48.7 | 68.4 | 60.9 | 81.8 | 0.40 | 3.30 | |
| HypeNet-9B | Pre-train | 100.0 | 100.0 | 100.0 | 100.0 | 64.2 | 95.8 | 72.4 | 67.2 | 53.0 | 95.0 | 83.8 | 52.0 | 17.6 | 57.6 | 15.3 | 62.4 | 0.80 | 6.70 |
| +SFT | 100.0 | 100.0 | 100.0 | 100.0 | 89.0 | 64.2 | 17.4 | 9.40 | 98.2 | 90.0 | 47.8 | 22.8 | 39.8 | 69.5 | 62.3 | 87.8 | 1.20 | 10.0 | |
| +QK-Restore | 100.0 | 100.0 | 100.0 | 100.0 | 92.6 | 100.0 | 44.0 | 19.6 | 98.6(0.4) | 97.0(7.0) | 66.4(18.6) | 42.6(19.8) | 37.6 | 68.4 | 59.3 | 88.4 | 0.60 | 3.30 | |
6.2 Main Results
The performance comparisons are in Table 1.
CoT-SFT drives long-context degradation. A key observation is that CoT-SFT often weakens long-context retrieval abilities acquired during pre-training, particularly under more challenging retrieval settings and longer context windows. While simpler tasks such as NIAH-Single-1 remain nearly saturated, substantial degradation emerges on NIAH-Single-2 and NIAH-Single-3.
For example in Table 1, HypeNet-2B on NIAH-Single-3 at 128K drops from to after SFT, while HypeNet-9B on NIAH-Single-2 at 256K decreases from to . Moreover, the degradation consistently becomes more severe as context length increases, suggesting that SFT disproportionately disrupts the mechanisms responsible for long-range retrieval.
QK-Restore recovers long-context capability. QK-Restore consistently mitigates the loss of long-context performance incurred during SFT, with the largest gains observed in configurations where the degradation is most severe.
Specifically in Table 1, for HypeNet-2B, NIAH-Single-3 at 256K improves from to (+19.6), while HypeNet-9B increases from to (+19.8). Under OpenThoughts-3 CoT-SFT, HypeNet-9B improves from to on NIAH-Single-2 at 256K and from to on NIAH-Single-3 at 256K. In several configurations, QK-Restore even surpasses the original pre-training baseline. For example, HypeNet-2B on NIAH-Single-3 at 64K improves from (pre-train) to . This suggests that our approach does not simply recover pre-training behavior; rather, it can preserve retrieval structures from pre-training while still benefiting from task-specific representations introduced during SFT.
QK-Restore preserves downstream reasoning performance. Beyond retrieval, we evaluate whether restoring routing parameters affects downstream reasoning capabilities. Across benchmarks, QK-Restore largely preserves the gains obtained from CoT-SFT while substantially improving long-context retrieval.
On mathematical reasoning tasks, performance remains close to the post-SFT model. For example, on HypeNet-5B, QK-Restore incurs only minor changes of and points on MATH500 and GSM8K, respectively, while recovering long-context retrieval performance. Similar trends are observed across other model scales. These results suggest that routing parameters can be selectively restored to recover long-context retrieval without significantly affecting the task-specific knowledge and reasoning capabilities acquired during post-training.
Retrieval difficulty amplifies long-context degradation. We further observe that retrieval complexity strongly influences robustness. NIAH-Single-1 remains nearly saturated across training stages, whereas substantially larger gaps emerge in Single-2 and Single-3. This trend indicates that difficult retrieval settings place greater demands on precise long-range token interactions and are therefore more sensitive to the degradation.
| Model | Method | NIAH-Single-1 | NIAH-Single-2 | NIAH-Single-3 | IFEval | |||||||||
| 32K | 64K | 128K | 256K | 32K | 64K | 128K | 256K | 32K | 64K | 128K | 256K | |||
| HypeNet-5B | Pre-train | 100.0 | 100.0 | 100.0 | 100.0 | 99.4 | 99.8 | 97.0 | 93.4 | 95.0 | 94.2 | 87.2 | 75.2 | 12.0 |
| +SFT | 99.8 | 100.0 | 100.0 | 100.0 | 99.8 | 100.0 | 99.6 | 93.6 | 98.6 | 97.2 | 85.4 | 83.8 | 26.6 | |
| +QK-Restore | 100.0 | 100.0 | 100.0 | 100.0 | 99.6 | 100.0 | 99.8 | 93.4 | 97.4 | 95.2 | 85.4 | 85.2 | 26.6 | |
| HypeNet-9B | Pre-train | 100.0 | 100.0 | 100.0 | 100.0 | 64.2 | 95.8 | 72.4 | 67.2 | 53.0 | 95.0 | 83.8 | 52.0 | 14.0 |
| +SFT | 100.0 | 100.0 | 100.0 | 100.0 | 86.4 | 97.8 | 88.2 | 78.4 | 96.2 | 96.8 | 95.8 | 86.8 | 27.2 | |
| +QK-Restore | 100.0 | 99.8 | 100.0 | 100.0 | 88.2 | 97.4 | 92.8 | 78.6 | 96.0 | 97.4 | 95.2 | 88.2 | 29.0 | |
7 Analysis
7.1 Analysis on Non-CoT SFT
Our theoretical analysis attributes long-context degradation to the localized optimization dynamics induced by CoT supervision. A natural question is whether this phenomenon arises from post-training in general, or whether it is specific to CoT-style reasoning traces. To answer this question, we compare CoT-SFT with conventional instruction-following SFT using the Tulu-3 dataset.
As shown in Table 2, standard instruction tuning does not exhibit the same degradation pattern observed under CoT-SFT. Across both HypeNet-5B and HypeNet-9B, long-context retrieval performance is largely preserved and frequently improved after Tulu-3 training. For example, HypeNet-5B improves from to on NIAH-Single-3 at 256K, while HypeNet-9B improves from to . Similar gains are observed on NIAH-Single-2 across multiple context lengths.
Moreover, instruction-following ability improves substantially after Tulu-3 tuning, as measured by IFEval. Despite these improvements, long-context retrieval remains intact. This contrasts sharply with the behavior of CoT-SFT, where retrieval performance deteriorates under identical architectures and context lengths.
These findings suggest that long-context degradation is not an inherent consequence of post-training. Instead, it is closely associated with CoT-style supervision, supporting our hypothesis that the sequential dependency structure of reasoning traces concentrates gradient updates on short-range interactions and disproportionately affects long-range routing behavior.
| Model | Method | NIAH-Single-1 | NIAH-Single-2 | NIAH-Single-3 | LCB-V5 | |||||||||
| 32K | 64K | 128K | 256K | 32K | 64K | 128K | 256K | 32K | 64K | 128K | 256K | Avg@8 | ||
| HypeNet-5B | Pre-train | 100.0 | 100.0 | 100.0 | 100.0 | 99.4 | 99.8 | 97.0 | 93.4 | 95.0 | 94.2 | 87.2 | 75.2 | 0.27 |
| +SFT | 100.0 | 100.0 | 99.8 | 100.0 | 99.8 | 100.0 | 99.0 | 89.6 | 98.8 | 97.8 | 94.2 | 91.4 | 6.37 | |
| +QK-Restore | 100.0 | 99.8 | 99.8 | 100.0 | 99.8 | 100.0 | 100.0 | 92.8 | 98.6 | 97.4 | 93.2 | 79.4 | 6.92 | |
| HypeNet-9B | Pre-train | 100.0 | 100.0 | 100.0 | 100.0 | 64.2 | 95.8 | 72.4 | 67.2 | 53.0 | 95.0 | 83.8 | 52.0 | 1.66 |
| +SFT | 100.0 | 100.0 | 100.0 | 100.0 | 91.8 | 100.0 | 41.8 | 22.4 | 96.2 | 98.6 | 76.6 | 40.6 | 13.11 | |
| +QK-Restore | 100.0 | 100.0 | 100.0 | 99.8 | 94.2 | 99.8 | 78.2 | 45.0 | 95.0 | 97.4 | 90.8 | 61.2 | 12.91 | |
7.2 Analysis Beyond Mathetical Reasoning
Our primary experiments focus on mathematical reasoning, where high-quality CoT supervision is readily available. To evaluate whether the observed retrieval–reasoning trade-off extends beyond mathematics, we conduct additional experiments using OpenThoughts-3 and evaluate downstream coding performance on LiveCodeBench-V5.
As shown in Table 3, the same qualitative behavior persists in the coding domain. CoT-SFT substantially improves coding capability, increasing LiveCodeBench-V5 from to on HypeNet-5B and from to on HypeNet-9B. However, these gains are accompanied by noticeable degradation in long-context retrieval, particularly on challenging NIAH settings and long context lengths.
Applying QK-Restore recovers a large fraction of the lost retrieval capability while preserving downstream coding performance. For example, on HypeNet-5B, LiveCodeBench-V5 further improves from to , while NIAH-Single-3 at 256K recovers from to . Similarly, HypeNet-9B improves from to on NIAH-Single-3 at 256K while maintaining competitive coding performance ().
These results suggest that the retrieval degradation induced by CoT-SFT is not specific to mathematical reasoning. Instead, it appears to arise from a more general optimization phenomenon associated with long reasoning traces. The effectiveness of QK-Restore across both math and code domains further supports our central hypothesis that long-range retrieval is primarily governed by routing parameters, while task-specific reasoning capabilities are largely encoded elsewhere in the network.
7.3 Analysis on Attention Routing Recover
In this section, we visualize the attention map to characterize how CoT supervised fine-tuning alters the long-range routing behavior of the softmax-attention layers in HypeNet-5B. We conduct the analysis on the models trained on Miromind dataset.
In detail, we compare the attention matrices of a pre-SFT and a post-CoT-SFT checkpoint of HypeNet-5B across all softmax-attention layers. We then identify the most affected layer via a per-head mean attention distance metric:
| (17) |
where indexes the attention head and is the post-softmax weight from query to key . Ranking layers by the number of heads whose decreases after SFT, Layer 33 emerges as the most affected, with 25 of 32 heads shifting toward shorter attention distances. We therefore focus our analysis on this layer.
We visualize the top recall heads on a synthetic NIAH prompt of 8,192 tokens, with the target fact embedded at context depth. For each head, we display a four-panel figure: the pre-SFT attention matrix (blue), the post-SFT matrix (red), the QK-Restore matrix (green), and the element-wise log-ratio (Panel 4), where positive values indicate positions where QK-Restore recovers more long-range attention than post-SFT.
Analysis. Figure 4 shows the pattern of Layer-33 Head-8 of HypeNet-5B. From panel 4, we observe that the off-diagonal lower triangle is predominantly red, confirming that QK-Restore partially recovers the long-range routing suppressed by CoT fine-tuning: for most query positions , QK-Restore attends more strongly to distant keys than post-SFT does. Second, the lower-right block, where both the query and key positions fall after the needle , shows a blue band along the main diagonal, meaning that post-SFT assigns higher weight to nearby tokens in this region, which matches our hypothesis. These observations demonstrate the effectiveness of QK-Restore in mitigating routing collapse and restoring the long-range attention behavior characteristic.
7.4 Ablation on QK-Restore
QK-Restore transplanting the full matrices from the pre-SFT checkpoint into the post-SFT model, recovers long-range recall substantially but incurs a math performance drop. The post-SFT can be decomposed as:
where encodes harmful locality drift and encodes beneficial math ability adaptation. Since both components are entangled within , QK-Restore discards them together, recovering routing at the cost of erasing the math benefit. As expected, Figure 5 and Figure 6 show a non-negligible drift in , where .
Therefore, we model it as a Procrustes problem:
| (18) |
and we call it as QK-Pro variant (more details are in Appendix N).
Furthermore, we evaluate the roles of each components in attention (i.e., ), and CoT-SFT with frozen as a preventive counterpart (denoted as QK-Frozen), to answer a question: is it preferable to prevent routing drift during training, or to correct it post-hoc?
| Method | NIAH-Single-2 | NIAH-Single-3 | MATH500 | ||
| 64K | 128K | 64K | 128K | Avg@16 | |
| QK-Restore | 91.2 | 83.2 | 52.8 | 30.8 | 33.7 |
| QK-Pro | 94.0 | 79.0 | 51.8 | 26.0 | 34.2 |
| Q-Restore | 78.4 | 79.0 | 50.8 | 24.4 | 34.2 |
| K-Restore | 90.4 | 71.6 | 46.2 | 13.2 | 34.2 |
| V-Restore | 57.2 | 36.6 | 24.0 | 15.2 | 32.1 |
| QK-Frozen | 90.6 | 76.4 | 49.4 | 23.6 | 33.1 |
Analysis. Table 4 validates three key claims. V-Restore fails to recover NIAH while degrading MATH500, confirming that does not contribute to routing recovery. Q-Restore and K-Restore each achieve partial recall recovery but fall significantly short at 128K (e.g., K-Restore collapses on NIAH-Single-3 (13.2 vs. 30.8), indicating that isolated restoration introduces a - mismatch and that joint restoration is necessary for coherent routing geometry. QK-Pro achieves the highest performance on MATH500, but trading long-range routing recovery for math preservation. QK-Frozen consistently underperforms QK-Restore in both NIAH and mathematical tasks, indicating that unconstrained SFT is for better adaptation.
8 Conclusion
In this work, we identify CoT-SFT-induced long-context recall degradation as a critical post-training failure mode in distilled hybrid models. We show that reasoning-oriented SFT biases optimization toward short-range attention patterns, perturbing the query-key routing geometry of retained softmax-attention layers and weakening long-range retrieval. Motivated by this observation, we propose QK-Restore, a training-free method that restores only the query and key projections from the pre-SFT checkpoint while preserving all other parameters. Experiments across multiple architectures show that QK-Restore substantially recovers long-context capability while largely preserving SFT reasoning gains. Overall, our findings highlight routing stability as a key factor for efficient long-context models that jointly support retrieval and reasoning.
References
Appendix A Discussion of the CoT-Markov Structure Assumption
Discussion 1. The 1-step Markov condition is an idealisation for analysis. Any -step Markov chain is equivalent to a 1-step chain on a -state product space, yielding the same exponential decay with a larger reasoning horizon .
Discussion 2. Reversibility is invoked in the proof of Lemma 4.5 at one step: it implies that the Markov operator is self-adjoint on . We acknowledge that a causally generated sequence such as a CoT reasoning chain does not, in general, satisfy detailed balance: the latent transition probabilities are asymmetric by construction. Reversibility is, however, a sufficient condition rather than a necessary one. For a general ergodic chain (possibly non-reversible), the identical bound holds with replaced by the second-largest singular value of , obtained via its singular value decomposition in place of the eigendecomposition. Since singular values are always real and non-negative, the proof proceeds without modification, and the form of the final bound is unchanged. We therefore state Lemma 4.5 under reversibility for clarity of exposition, with the understanding that all subsequent results hold for the general ergodic case by replacing eigenvalues with singular values throughout.
Appendix B Justification for Norm Bounds Assumption
Bound on . The value vector is , where denotes layer normalisation. Layer normalisation standardises each coordinate so that , where and are the empirical mean and standard deviation of the entries of . After normalisation, the pre-affine vector satisfies , giving
| (19) | ||||
| (20) | ||||
| (21) |
for a constant depending only on the learned scale and bias . The projection matrix is a fixed trained weight; its operator norm is a finite model-dependent constant. Combining we have:
| (22) | ||||
| (23) |
Bound on . The gradient back-propagates through the final layer normalisation (with Jacobian ) and the language-model head , giving
| (24) |
where is the predicted distribution and is the one-hot target. By submultiplicativity of the operator norm:
| (25) |
Since is a difference of probability vectors, . For a fixed trained model, both and are finite constants. Therefore:
| (26) |
Appendix C Justification for Score-function Identity Assumption
At the optimum , the cross-entropy loss with respect to the attention output satisfies, for any fixed context :
| (27) | |||
| (28) | |||
| (29) |
since the cross-entropy gradient vanishes at its own minimum. Applying the law of total expectation over then yields . More precisely, because the gradient is computed at a given context and then averaged over the target drawn from the conditional, the identity holds conditionally:
| (30) |
Appendix D Preliminary Definition and Lemma
We list the following necessary definitions and lemmas for the proof.
Definition D.1( Space).
For a Markov chain with finite state space and stationary distribution , define the weighted function space:
| (31) |
with inner product
| (32) | ||||
| (33) |
and norm .
In addition, mean-zero means , i.e., the function has zero expectation under .
Definition D.2(Operator acting on functions).
We define the Markov operator acting on functions:
| (34) | ||||
| (35) |
Lemma D.3.
For a reversible chain, detailed balance holds: , which makes symmetric in .
Proof.
According to the definition of Markov operator in D.2, we have:
| (36) | ||||
| (37) | ||||
| (38) | ||||
| (39) | ||||
| (40) |
Since , the finite-dimensional spectral theorem for self-adjoint operators yields a real orthonormal eigenbasis of with eigenvalues , and spectral gap where :
| (41) |
Further, for mean-zero :
| (42) |
which means has no component along . Its expansion uses only :
| (43) |
Appendix E Proof for Spectral Correlation Decay Lemma
Proof.
At stationarity (), the Tower Property tells:
| (44) |
Note that the inner expectation is the -step prediction: for a fixed starting state , the expected value of after steps is exactly :
| (45) |
Therefore:
| (46) | ||||
| (47) | ||||
| (48) |
Then we expand by the orthonormal eigenbasis:
| (49) |
Using the orthonormality , we further simplify :
| (50) | ||||
| (51) |
Then we bound it:
| (52) | ||||
| (53) | ||||
| (54) |
Applying the Cauchy-Schwarz Inequality:
| (55) | ||||
| (56) |
Now combining all together:
| (57) | ||||
| (58) |
Appendix F Proof of the Gradient Locality Theorem
Proof.
For clarity, we drop the layer/head superscripts in the following derivation. First, we bound the gradient:
| (59) |
Then we consider how to bound . Applying the Tower Property:
| (60) | ||||
| (61) |
where we define:
| (62) |
According to Assumption 4.4:
| (63) | ||||
| (64) | ||||
| (65) | ||||
| (66) |
Similarly, centering does not change the covariance (since ). Both and are mean-zero functions on .
Then applying Lemma 4.5:
| (67) | ||||
| (68) | ||||
| (69) |
Applying the same bound to the term, which satisfies as a convex combination:
| (70) |
where and . ∎
Appendix G Proof for Token Autocorrelation Decay
Define the centred indicator feature and the -weighted autocorrelation at lag :
| (71) |
For the numerator , we expand it as:
| (72) |
Considering the expectation of :
| (73) | ||||
| (74) |
Therefore we can further simplify:
| (75) |
In the HMM, tokens are conditionally independent given their latent states: . By the tower property:
| (76) | |||
| (77) | |||
| (78) |
where is a mean-zero function of :
| (79) | ||||
| (80) | ||||
| (81) |
Since both and are mean-zero, we can apply Lemma 4.5 directly:
| (82) |
where is the second-largest eigenvalue of the latent transition matrix . Now we can bound the autocorrelation:
| (83) | ||||
| (84) |
Conclusion. Observing with in the corpus confirms that the latent chain has spectral gap , which is exactly Lemma 4.5’s hypothesis.
Appendix H Discussion of Attention Gradient Decay in Pre-training Stage
Theorem 4.6 rests on Lemma 4.5, which holds for any ergodic Markov chain with a spectral gap: a property of the data-generating process, not specific to CoT structure, which is also observed in general pre-training data corpus.
Consequently, the gradient locality bound applies to any training corpus, including general pre-training data. The severity of routing collapse is governed by the mismatch : how far the data demands the model to attend vs. how far the gradient actually reinforces it. For general pre-training corpora, this mismatch is smaller than for CoT, so routing degrades more slowly but by the same mechanism.
For hybrid architectures, however, even the slower pre-training erosion meaningfully reduces available recall capacity before fine-tuning begins, leaving the model closer to the threshold. The hybrid architecture’s vulnerability therefore originates during pre-training and is further amplified by CoT fine-tuning.
Appendix I Proof for the Routing-Extraction Gradient Decoupling Theorem
Proof.
From Equation 7, we have:
| (85) |
where is the corresponding KV head. Then considering its Frobenius norm:
| (86) | ||||
| (87) | ||||
| (88) |
via Theorem 4.6 and bounded norms , .
For the Exctraction Parameters, the value vector enters the loss through . Accumulating all downstream dependencies:
| (89) |
Now we assume that for all and for all :
| (90) |
Appendix J Training Data and Configuration Details
For HypeNet, we conduct both the pre-training and SFT.
Pre-training Data. Since we mainly focus on the long-context recall and math reasoning ability, during the pre-training stage, we mix the FineWeb-Edu [lozhkov2024fineweb-edu] and UltraData-Math [ultradata-math]. Following observations from recent pretraining studies that maintain web corpora as the dominant source while increasing structured reasoning data [allal2025smollm2smolgoesbig], we construct a mixture with general web data and math-focused data. Prior evidence indicates that web-scale corpora preserve broad linguistic and factual competence, while specialized mathematical corpora improve reasoning and STEM performance.
Pre-training Config. We adopt the 3-stage training for HypeNet, keeping the same setting in chen2026hybridlinearattentionright, shown in Table 5, Table 6, and Table 7.
| Stage | Tokens | LR | LR Scheduler | Context Len. | Batch | Training Steps |
| 1 | 320M | 1e-3 1e-5 | Cosine | 512 | 32 | 20,000 |
| 2 | 1B | 1e-4 1e-5 | Cosine | 512 | 96 | 20,000 |
| 3 | 1B | 1e-5 | Constant | 16,384 | 128 | 500 |
| Stage | Tokens | LR | LR Scheduler | Context Len. | Batch | Training Steps |
| 1 | 320M | 1e-3 1e-5 | Cosine | 512 | 32 | 20,000 |
| 2 | 1B | 5e-5 1e-5 | Cosine | 512 | 96 | 20,000 |
| 3 | 1B | 1e-5 | Constant | 16,384 | 128 | 500 |
| Stage | Tokens | LR | LR Scheduler | Context Len. | Batch | Training Steps |
| 1 | 320M | 1e-3 1e-5 | Cosine | 512 | 32 | 20,000 |
| 2 | 1B | 3e-5 1e-5 | Cosine | 512 | 96 | 20,000 |
| 3 | 1B | 1e-5 | Constant | 16,384 | 128 | 500 |
| LR | LR Scheduler | Context Len. | Batch | Training Steps |
| 1e-5 | Constant | 16,384 | 128 | 100 |
| LR | LR Scheduler | Context Len. | Batch | Training Steps |
| 1e-5 1e-6 | Cosine | 16,384 | 128 | 100 |
SFT Data. To further improve the model’s math reasoning performance, we conduct CoT-SFT with MiroMind-M1 dataset [li2025miromindm1opensourceadvancementmathematical], which is collected from OpenR1 [openr1], Open-thoughts [guha2025openthoughts], Light-R1 [wen2025light], and Synthetic-1 [2025synthetic1]. SFT Config. The detailed configurations are shown in Table 8 and Table 9.
Appendix K Evaluation Details
For evaluating NIAH, we apply LM Evaluation Harness222https://github.com/EleutherAI/lm-evaluation-harness [eval-harness] for official test. For evaluating MATH500 and GSM8K, we report average pass@1 over 16 independent generations and Maj@16 (i.e., majority vote as the final prediction) per problem. In detail, for MATH500, the generation length is 8,192; for GSM8K, we set it as 2,048, since it is easier. For LiveCodeBench, we report the average pass@1 over 8 independent generations with 16,384 response length. To measure the model’s performance after fine-tuning on Tulu3, we report the Prompt-level strict accuracy on IFEval task [zhou2023instructionfollowingevaluationlargelanguage].
Appendix L Analysis on Pure Softmax-Attention Model
Our main experiments concentrate on the hybrid models, and in this section, we explore whether SFT can degrade the pure softmax-attention models as well. We include Qwen2.5-3B, Qwen2.5-7B [qwen2025qwen25technicalreport] and Misrtral-7B-Instruct-v0.3 [jiang2023mistral7b] for investigation. We include the details on how we identify the top layers for retrieval in Appendix M.
| Model | Method | NIAH-Single-2 | NIAH-Single-3 | ||
| 32K | 64K | 32K | 64K | ||
| Qwen2.5-7B | Pre-train | 100.0 | 95.8 | 99.8 | 98.6 |
| +SFT | 100.0 | 92.6 | 99.8 | 99.4 | |
| Qwen2.5-3B | Pre-train | 100.0 | 94.6 | 100.0 | 99.6 |
| +SFT | 100.0 | 93.8 | 100.0 | 90.8 | |
| +QK-Restore | 100.0 | 93.8 | 100.0 | 93.4 | |
| Mistral-7B | Pre-train | 99.8 | 90.2 | 99.8 | 81.0 |
| +SFT | 99.6 | 59.6 | 65.8 | 4.80 | |
| +QK-Restore | 99.8 | 62.0 | 70.0 | 5.40 | |
Analysis. Table 10 shows that for Mistral-7B, it exhibits strong sensitivity to SFT at longer contexts. SFT causes severe degradation on NIAH-Single-3 at 64K, dropping from to , suggesting a breakdown of long-range retrieval. QK-Restore recovers this degradation limitedly, indicating weaker long-context representations in the pre-trained model.
For Qwen2.5-7B and Qwen2.5-3B, long-context performance is highly stable across NIAH-Single-2 and Single-3. SFT induces only minor fluctuations (e.g., Qwen2.5-7B drops from to at 64K), indicating that these models already learn robust long-context retrieval during pre-training. Accordingly, QK-Restore yields only marginal gains, consistent with a near-saturation regime where little capability is lost. To investigate the underlying mechanism, we measure the attention gradient decay on Qwen2.5-3B, shown in Figure 7. Unlike HypeNet-2B, whose full-attention layers exhibit a well-fitted exponential decay, it has a clear near-flat tail. We hypothesize that long-context recall in pure softmax-attention models is maintained structurally, via the abundance of full-attention layers.
Appendix M Method of Identifying Retrieval Layers in Pure Softmax-Attention Models
To identify which softmax-attention layers are primarily responsible for long-range retrieval, we conduct a leave-one-out ablation study over the set . For each layer , we construct an ablated model by setting . Each ablated model is evaluated on several retrieval benchmarks (Ruler-QA-SQuAD [hsieh2024ruler], FDA [arora2023language], SWDE [lockard2019openceres]). The importance score of layer is defined as:
where denotes the average retrieval accuracy across benchmarks. A large indicates that layer contributes disproportionately to retrieval and that its routing weights are critical; a near-zero score indicates redundancy. Layers are ranked by in descending order, and the top- layers are selected as the target set for weight restoration. In our experiments, we keep as the same as in HypeNet.
Appendix N Details of deriving QK-Pro
The bilinear constraint has no closed-form joint solution, so we linearise by fixing one factor to its pre-SFT value, making the remaining problem a standard constrained least-squares system solvable via Lagrange multipliers.
We fix , find minimum-deviation :
| (91) | |||
| (92) |
Then applying the Lagrange multipliers:
| (93) |
Taking the derivative , and we get:
| (94) |