注意力遗忘症：思维链微调如何悄悄抹掉混合模型的长程记忆，以及怎么一行不练把它修回来
arXiv ID：2606.11052 ｜ 提交日期：2026-06-09
如果你正在用混合线性注意力架构（hybrid linear-attention）的大模型，又恰好给它做过思维链监督微调（CoT-SFT），那这篇论文值得你停下来看一眼。它讲的是一件相当反直觉的事：你辛辛苦苦把模型的数学推理能力调上去了，结果它在长文档里"翻书找答案"的能力却悄悄塌了一大块，而且塌得很有规律。更要命的是，作者证明这件事不是偶然，是训练机制层面的必然后果——然后给出了一个连一步反向传播都不需要的修复办法。
一句话先说清楚这篇论文在干嘛
混合架构模型里那几层稀有的 softmax 注意力，本来扛着"长程精确召回"这个重任。CoT-SFT 这种偏向局部、短程推理的训练数据，会在反向传播里系统性地饿死这些长程连接，导致路由（routing）几何结构漂移、长程召回崩溃。作者把这个现象命名为注意力遗忘症（Attention Amnesia），从理论上推导出它发生的精确条件，再提出 QK-Restore：只把 softmax 层的 \(W_Q, W_K\) 两个权重从微调前的 checkpoint 搬回来，其他参数原封不动，免训练，长程召回当场恢复，数学能力还基本不掉。
一、先理解背景：混合架构为什么会有这个软肋
纯 softmax 注意力的 Transformer，每一层都能看全程，长程召回不是问题，问题是算力和显存随上下文长度平方级爆炸。于是这两年很流行一种折中：混合线性注意力架构。它的做法是把绝大多数层换成高效的线性注意力（比如 Lightning Attention），只在少数几层保留完整的 softmax 注意力。线性层负责便宜地处理大部分信息流，那几层 softmax 则像稀缺的"精密仪器"，专门负责需要全局精确定位的活儿——典型代表就是大海捞针式的长程召回。
论文里的实验主力是 HypeNet 系列（2B / 5B / 9B 三个规模）和 Jet-Nemotron-2B，这些都是混合架构。对照组则是纯 softmax 的 Mistral-7B、Qwen2.5-3B/7B。
这套设计平时跑得好好的。麻烦出在你想让它会推理、于是上了 CoT-SFT 之后。
二、现象：微调后长程召回的结构化崩塌
作者用 NIAH（Needle-In-A-Haystack，大海捞针）这个经典长程召回评测来量化。NIAH 的玩法是：在一大段无关文本里藏一句"针"（关键信息），然后在很远的地方问模型这句针的内容，看它能不能精确捞回来。难度分 Single-1/2/3 三档，数字越大要捞的针越多、越绕；上下文长度从 32K 一路拉到 256K。
下面这张图是整篇论文的"案发现场"：
蓝色实线是预训练版（Pre-train）的 HypeNet-2B，红色虚线是做过 CoT-SFT 的同一个模型。左图 NIAH-Single-2、右图 NIAH-Single-3。你能很清楚地看到：随着上下文拉长，微调版的召回分数掉得触目惊心——Single-2 上掉了 40%，Single-3 上掉了 23%。注意，这两个模型的数学推理能力，微调版是明显更强的。也就是说，你花训练换来的推理增益，背地里是拿长程记忆抵押的。
更关键的是，作者强调这种退化不是随机噪声，而是结构化的：它专门发生在需要长距离依赖的样本上，短程的基本不受影响。这个"结构化"特征，正是后面理论分析能站住脚的前提。
三、理论：为什么 CoT-SFT 一定会饿死长程注意力
这是本文最硬核、也最有说服力的部分。作者没有停在"我们观察到一个现象"，而是把它推到了"在这些假设下，这个现象必然发生"。
3.1 两把尺子：相关性衰减 vs 梯度衰减
作者引入了两个刻画"距离"的特征宽度。
第一把尺子是 token 自相关衰减宽度 \(W_{corr}\)。它衡量的是：数据本身要求模型关注多远。如果一段语料里，相隔 \(\tau\) 个 token 的内容还高度相关，那 \(W_{corr}\) 就大。论文拿两类语料做对比测出来：CoT 推理数据（MiroMind）的 \(W_{corr} \approx 172\)（拟合 \(\hat\rho=0.9942\)，\(R^2=0.98\)），而一般文本（FineWeb）的 \(W_{corr} \approx 115\)（\(\hat\rho=0.9914\)，\(R^2=0.90\)）。也就是说，CoT 数据本身的内容相关性反而拖得更长，需要更长的注意力跨度才能覆盖。
第二把尺子是注意力梯度衰减宽度 \(W_{grad}\)。它衡量的是：训练时，反向传播的梯度信号能把"关注某个远处 token"这件事强化到多远。作者证明了一个关键结论（Theorem 4.6）：在 CoT-Markov 结构假设下，注意力的梯度局部性是指数衰减的——
也就是说，距离 \(\tau\) 越远，梯度对那条长程连接的强化越是指数级地微弱。实测下来，CoT-SFT 后的梯度衰减宽度 \(W_{grad} \approx 41\)（\(R^2=0.82\)），比预训练阶段一般数据的 \(W_{grad} \approx 48\)（\(R^2=0.84\)）还要更小——微调不仅没把梯度的有效作用范围撑大，反而进一步收窄了。
3.2 致命的"距离带"
把两把尺子摆在一起，崩塌的根因就浮出水面了。当 \(W_{grad} \lt W_{corr}\) 时，存在一个距离区间：
在这个带里，数据明明需要长程注意力（因为 \(d \leq W_{corr}\)），但训练梯度已经弱到不再强化它了（因为 \(d \gt W_{grad}\)）。代入实测的两个宽度，这条致命带就是 \(d \in (41, 172]\)——足足一百多个 token 的跨度，落在里面的长程连接每过一步 SFT 都得不到梯度维护，权重一点点漂走，最终路由崩溃。这就是注意力遗忘症的微观成因。
下面这张图验证了这两条衰减曲线在真实语料上确实如理论所料：
3.3 路由与抽取的梯度解耦
光知道"长程连接会饿死"还不够，得知道该救哪个权重。这就引出第二个核心定理（Theorem 5.1）：路由-抽取梯度解耦。
softmax 注意力可以拆成两个功能：
- 路由（Routing）：由 \(W_Q, W_K\) 决定，负责"决定关注谁"，也就是注意力的几何结构 \(R = W_Q W_K^\top\)；
- 抽取（Extraction）：由 \(W_V\) 决定，负责"关注到了之后把什么信息搬出来"。
作者证明，在 CoT-SFT 的梯度里，路由部分和抽取部分是解耦的，而前面那套指数衰减恰恰主要打击的是路由。换句话说，微调真正破坏的是"该看哪里"的能力（\(W_Q, W_K\) 漂移），而"看到后怎么用"（\(W_V\)）基本没坏。
这个结论的操作意义极其直接：要修，就只修 \(W_Q\) 和 \(W_K\) 就够了。
下面两张图直接量化了这种"路由几何漂移"——逐层统计微调前后路由矩阵 \(R = W_Q W_K^\top\) 的相对变化 \(\|\Delta R\|_F / \|R_{pre}\|_F\)，而且在 HypeNet-2B 和 Jet-Nemotron-2B 两个不同架构上都成立：
还可以从注意力图（attention map）这个更直观的视角看同一件事。下面这张图取 layer 33、head 8，把针放在位置 4096，横轴是被关注的 key 位置、纵轴是发起查询的 token 位置。四个面板从左到右分别是 Pre-SFT（预训练）、Post-CoT-SFT（微调后）、QK-Restore（修复后）三态的注意力热力图，外加一张 QK-Restore 与 Post 的对数比值图。对照前两张能很清楚地看到：微调后长程的注意力质量明显被抽干、塌缩到局部，而 QK-Restore 把指向针位置的那条长程注意力又点亮了回来：
作为对照，纯 softmax 的 Qwen2.5 也被拿来验证梯度衰减规律，CoT 微调后（红）的梯度衰减明显快于预训练阶段（蓝），说明这套理论不只对混合架构成立：
四、方法：QK-Restore，把"该看哪里"的记忆搬回来
理论铺到这儿，方法就简单到近乎"废话"了——但这恰恰是它优雅的地方。
4.1 核心做法
既然坏的是 softmax 层的路由几何，而路由几何由 \(W_Q, W_K\) 决定，那就只把这两个权重从微调前的 checkpoint 移植回微调后的模型，其余所有参数（FFN、线性注意力层、softmax 层的 \(W_V\) 输出投影等）全部保留 post-SFT 的版本。
一句话：\(W_Q^{post}, W_K^{post} \leftarrow W_Q^{pre}, W_K^{pre}\)，别的不动。
这就是 QK-Restore。它不需要任何额外训练、不需要数据、不需要梯度，纯粹是权重的"零成本器官移植"。
下面这张方法总览图把整件事讲得很清楚：
图分三栏。A 栏画的是 post-SFT 的混合模型结构（FFN、softmax 注意力、线性注意力交错堆叠）；B 栏是 QK-Restore 的操作——把发生了 routing drift 的 \(W_Q^{post}, W_K^{post}\) 换成绿色的微调前版本 \(W_Q^{pre}, W_K^{pre}\)，而 \(W_V^{post}\) 保留不动；C 栏是效果柱状图，分别对比 Pre-SFT、Post-SFT、QK-Restore 三者在数学推理（Math Reasoning）和长程召回（Long-Context Recall）上的表现。结论一目了然：CoT-SFT 让推理↑但召回↓，QK-Restore 则让推理基本持平、召回拉回来。
4.2 进阶变体：QK-Pro
直接搬回 \(W_Q^{pre}, W_K^{pre}\) 有一个潜在顾虑：微调后的其他参数已经适应了新的推理模式，硬塞回旧的路由权重可能不是最优配合。于是作者又提出一个 Procrustes 变体 QK-Pro。
它的思路是：在严格保持微调前那套路由几何 \(R_{pre} = W_Q W_K^\top\) 不变的约束下，去寻找一组在数值上尽可能接近 post-SFT 权重的 \(W_Q, W_K\)。这等价于一个带正交约束的 Procrustes 问题——既保住了"该看哪里"的几何，又让权重尽量贴合微调后的整体状态，在路由保留和推理适应之间取一个更细的平衡。
五、实验：到底修回来多少
下面把论文 Table 1–4 的核心数值整理出来。结论很统一：QK-Restore / QK-Pro 在几乎不牺牲数学推理的前提下，把长程召回大幅拉回，且在 2B/5B/9B 多个规模和两种混合架构上都稳定复现。
5.1 长程召回的恢复（NIAH）
CoT-SFT 后 Single-2 掉 40%、Single-3 掉 23% 的窟窿，QK-Restore 基本补平，把召回曲线拉回到接近预训练版的水平，且上下文越长优势越明显。这印证了理论预测——被饿死的正是落在 \((W_{grad}, W_{corr}]\) 距离带里的长程连接，而把路由权重搬回来恰好直接复原了这条带上的注意力。
5.2 数学推理几乎不掉
关键的代价问题：把 \(W_Q, W_K\) 换回旧版，会不会把辛苦练来的推理能力也一起退回去？答案是基本不会。因为 Theorem 5.1 的解耦结论保证了——推理能力的增益主要沉淀在 \(W_V\)（抽取）和 FFN 里，而这些 QK-Restore 一概不动。所以推理分数维持在 post-SFT 的高位，召回却回来了，鱼和熊掌罕见地同时拿到。
5.3 跨规模、跨架构的稳健性
- 规模维度：HypeNet 在 2B、5B、9B 三档上，退化与修复的规律一致，说明这不是某个小模型的偶发问题，而是混合架构的结构性特征。
- 架构维度：换到 Jet-Nemotron-2B（不同的混合设计），路由漂移和 QK-Restore 的修复效果同样成立（见 Figure 6）。
- QK-Pro 对比：在需要更精细平衡推理与召回的设定下，QK-Pro 比直接 Restore 更稳，给出了一个可调的中间档。
六、我的几点解读
第一，这篇论文真正的价值在"机制可解释"。 现在很多关于微调副作用的工作停在"我们发现微调会损害某能力"，靠的是经验观察。这篇不一样，它把现象归因到了一个可计算的不等式 \(W_{grad} \lt W_{corr}\) 上，还分别测出了两个宽度的数值（172 vs 115 那组对比尤其有说服力）。当你能写出导致问题的精确条件，修复方案往往就自动浮现了——QK-Restore 之所以"只修 QK"，不是试出来的，是 Theorem 5.1 解耦结论直接指出来的。这种从理论到方法的闭环，是我觉得最漂亮的地方。
第二，免训练修复的工程吸引力被低估了。 在真实落地里，你常常拿到的是一个已经做过 CoT-SFT 的 checkpoint，重训成本高、数据也未必还在手边。QK-Restore 只要你还留着微调前的权重，就能零成本把长程召回补回来，没有任何额外算力开销。对于混合架构的部署方来说，这几乎是"白捡"的能力恢复。
第三，它其实在提醒一个更普遍的问题：训练数据的"距离结构"会塑造模型的"注意力结构"。 CoT 数据偏短程，于是把长程连接饿死了。反过来想，如果你的下游真的需要长程能力，那 SFT 数据里就得有意识地掺入长距离依赖的样本，去把 \(W_{grad}\) 撑大、覆盖住 \(W_{corr}\)。QK-Restore 是事后补救，而这个洞察指向的是事前预防。
第四，边界也要说清楚。 QK-Restore 的前提是你手里得有微调前的 \(W_Q, W_K\)——如果只拿到一个 post-SFT 的模型、没有原始 checkpoint，这招就使不出来。另外它修的是"路由几何漂移"这一类退化，对于那些确实需要靠新训练才能获得的长程能力（而非被破坏后待恢复的），它无能为力。它是"复原"，不是"凭空创造"。
七、一句话总结
混合线性注意力模型里那几层珍贵的 softmax 注意力，是长程记忆的命根子；CoT-SFT 会因为梯度的指数级局部衰减，在 \((W_{grad}, W_{corr}]\) 这条距离带上悄悄饿死它们的路由连接，造成结构化的注意力遗忘症。而既然坏的只是"该看哪里"的 \(W_Q, W_K\)，那把它们从微调前搬回来就够了——QK-Restore，一步不练，召回归位，推理无损。论文 arXiv ID 是 2606.11052，做混合架构的同行建议读一遍理论部分，那才是它的精华。