Attention Amnesia：CoT 后训练如何改坏 Hybrid LLM 的长程检索路由

核心判断：这不是泛泛的“记忆预算不够”，而是 Q/K routing 被后训练改偏

如果只读推文，容易得到一个模糊结论：模型学会 CoT 后，长上下文能力就会被挤掉。论文真正有信息量的地方在于把这个现象拆到了参数功能层面：attention 里的 Q/K 更像“索引和寻址系统”，V/O 更像“被索引后拿回什么内容”。CoT-SFT 主要破坏前者，而后者仍可承载后训练带来的 reasoning adaptation。

现象

CoT-SFT 后，HypeNet、Jet-Nemotron 等 hybrid 模型在困难 NIAH 长上下文召回上明显下降，且 context 越长、检索越难越明显。

机制

CoT trace 的训练梯度集中在短距离 token pair，导致少数 full-attention 层的 W_Q/W_K 发生 locality-biased drift。

修法

QK-Restore 从 pre-SFT checkpoint 恢复 Q/K，保留 post-SFT 的其他参数，从而在零训练成本下恢复一大部分长程召回。

一句话重述

这篇论文不是说 CoT prompt 会立刻伤害检索，也不是说所有模型一做 reasoning SFT 都会崩；它说的是：在依赖少数 full-attention 层做全局召回的 hybrid 架构中，CoT-SFT 会把负责 long-range selection 的 Q/K 几何调成更局部的 reasoning transition 几何。

问题背景：便宜长上下文与强推理后训练为什么会冲突

长上下文服务的成本瓶颈来自 attention 和 KV cache。为了把 128K、256K 甚至更长上下文做得便宜，许多模型路线会从纯 softmax attention 走向 linear attention、state-space model 或 hybrid attention。Hybrid 模型的典型折中是：大部分层换成线性或递归机制，少数层保留 full softmax attention，用它们承担全局检索和远距离 token routing。

这个设计把成本降下来，但也把长程召回能力集中到了少数“全局路由层”上。纯 Transformer 中的远距离索引有很多冗余层和头；hybrid 模型里，如果少数保留的 full-attention 层被后训练改坏，模型不一定还有足够冗余来补救。

架构压缩大多数层不再做完整 softmax attention

稀缺路由少数 full-attention 层负责全局召回

CoT-SFT训练信号来自局部推理步骤

Q/K 漂移路由几何偏向近邻 token

召回失灵远处 needle 难以被稳定选中

这就是推文中“teach it to reason / forgets how to retrieve”的工程背景：cheap long-context serving 和 CoT reasoning 恰好都很重要，但它们会在 hybrid 模型的稀缺 routing capacity 上相遇。

机制拆解：Q/K 决定“去哪找”，CoT 梯度却主要强化“附近找”

1. Attention routing 与 extraction 的功能分工

在一个 softmax-attention head 中，查询位置 \(t\) 对历史位置 \(s\) 的注意力强度由 pre-softmax logit 决定：

e_{ts}^{(\ell,h)}=\frac{\left(q_t^{(\ell,h)}\right)^\top k_s^{(\ell,h)}}{\sqrt{d_h}},\qquad q_t=W_Q h_t,\quad k_s=W_K h_s,\quad v_s=W_V h_s.

这里 W_Q 和 W_K 决定 query 与 key 的匹配，也就是模型要把注意力路由到哪里；W_V 与输出投影更多决定从被选中的位置抽取什么内容。论文把这称为 routing-extraction decomposition。它不是一个严格的模块隔离定理，但在解释长程检索时很有用：如果 Q/K 几何坏了，模型可能根本看不到远处信息；如果 V/O 继续适应，模型仍可能在短程 reasoning benchmark 上表现变好。

2. CoT trace 为什么天然偏局部

论文把 CoT 推理轨迹抽象成 latent Markov process：当前推理状态主要由上一个或近几个状态转移而来。这个假设是理想化的，但抓住了数学和代码推理 trace 的一个共同结构：模型在写“因此”“下一步”“代入上式”时，最直接的监督信号通常来自邻近步骤，而不是十万 token 之前的事实。

在这个视角下，distance-conditioned attention-logit gradient 会随距离衰减：

g(\tau)=\mathbb{E}_{t,\ell,h}\left[\left|\frac{\partial \mathcal{L}}{\partial e_{t,t-\tau}^{(\ell,h)}}\right|\right] \le C'\exp(-\tau/W).

直观解释是：当 \(\tau\) 很小，当前 token 与前面推理步骤的关联强，训练会反复强化这种邻近模式；当 \(\tau\) 很大，梯度很弱，长距离 routing 得不到足够维护。于是即使模型的上下文窗口标称支持 256K，post-training 也可能把“应该能跳远”的 Q/K 投影调成“更爱看附近”的投影。

3. 为什么问题集中在 hybrid 模型

纯 softmax Transformer 有多层、多头的全局 attention 冗余；局部化一部分头不一定立刻摧毁长程检索。Hybrid 模型不同：它把多数层替换为 linear / recurrent state，长程召回更依赖少数保留 full-attention 层。论文的注意力可视化也支持这一点：CoT-SFT 后，受影响层的平均注意力距离缩短，QK-Restore 后远距离 off-diagonal attention 被部分恢复。

关键 insight

论文的机制不是“模型 memory budget 被 reasoning 占满”这么粗，而是“负责远距离寻址的低冗余 Q/K 路由系统，被 CoT 局部梯度重写”。这使得修复可以非常局部：不必重新训练整个模型，只要恢复一小组参数的 routing geometry。

QK-Restore：用 pre-SFT 的路由，保留 post-SFT 的推理

QK-Restore 的方法简单到像一个 diagnostic intervention：从 post-SFT 模型出发，只把 full-attention 层的 W_Q 和 W_K 换回 pre-SFT checkpoint，其余参数保持 post-SFT。形式化地说：

\theta_{\mathrm{rep}}\leftarrow \theta_{\mathrm{post}},\qquad W_{Q,\mathrm{rep}}^{\ell}\leftarrow W_{Q,\mathrm{pre}}^{\ell},\quad W_{K,\mathrm{rep}}^{\ell}\leftarrow W_{K,\mathrm{pre}}^{\ell},\quad \ell\in L_{\mathrm{attn}}.

这个方法的意义有两层。第一，它是实用修复：只要保存了 pre-SFT 和 post-SFT 权重，就可以零训练成本生成 repaired checkpoint。第二，它是因果诊断：如果只替换 Q/K 就能恢复召回，同时保留 reasoning，那么说明“召回损坏”主要不是全模型灾难性遗忘，而是可定位的 routing drift。

输入pre + post需要 CoT-SFT 前后的 checkpoint

替换对象Q/K只作用于保留的 full-attention 层

训练成本0不需要重新 SFT 或 RL

目标routing恢复“去哪找”，不回滚全部 reasoning adaptation

Procrustes 变体：在恢复路由和保留 reasoning adaptation 之间折中

直接恢复 Q/K 会把 post-SFT 中可能有用的 Q/K 变化也一起抹掉。论文进一步提出 QK-Pro / Procrustes 变体，把目标写成“尽量接近 post-SFT 权重，同时让 QK routing product 回到 pre-SFT 几何”。核心约束可以理解为：

\min_{W_Q^{\mathrm{new}}, W_K^{\mathrm{new}}} \left\|W_Q^{\mathrm{new}}-W_Q^{\mathrm{post}}\right\|_F^2+ \left\|W_K^{\mathrm{new}}-W_K^{\mathrm{post}}\right\|_F^2 \quad \mathrm{s.t.}\quad W_Q^{\mathrm{new}}\left(W_K^{\mathrm{new}}\right)^\top = R_{\mathrm{pre}}.

实验上，QK-Pro 往往在 math 保留上略好，但长程召回恢复不如完整 QK-Restore。这说明 beneficial math adaptation 和 harmful locality drift 在 Q/K 中确实有纠缠；越强行保持 post-SFT Q/K，越可能保留一部分局部化路由。

关键证据：退化、恢复、对照与消融

1. CoT-SFT 后，困难长上下文 NIAH 明显退化

论文最有说服力的不是单个数字，而是模式：简单 needle 场景常常接近饱和，困难场景和更长上下文才暴露问题。下面列的是正文中最能说明趋势的样本。

模型与任务	Pre-SFT	CoT-SFT 后	QK-Restore 后	解读
HypeNet-9B · NIAH-S2 @ 256K	67.2	9.4	19.6	最醒目的“attention amnesia”：CoT-SFT 几乎摧毁困难长程检索，QK-Restore 有恢复但不完全。
HypeNet-9B · NIAH-S3 @ 256K	52.0	22.8	42.6	恢复幅度接近 20 点，说明 Q/K 确实承载了大量长程 routing。
HypeNet-2B · NIAH-S3 @ 128K	30.4	7.6	30.8	几乎回到 pre-SFT 水平，是 QK-Restore 作为 causal probe 的强证据。
HypeNet-5B · NIAH-S3 @ 256K	75.2	65.4	76.4	恢复后超过 pre-SFT，说明不是简单回滚模型，而是组合了 pre-SFT routing 与 post-SFT 表征。

2. 普通 instruction SFT 没有同样破坏长上下文

论文用 Tulu-3 做 non-CoT instruction-following SFT 对照。结果并不支持“任何 post-training 都会伤长上下文”：HypeNet-5B 和 HypeNet-9B 在一些长程 NIAH 指标上反而提升。这一点很重要，因为它把问题从“后训练泛化风险”缩窄到“CoT-style traces 的局部监督结构”。

设置	代表结果	含义
HypeNet-5B · Tulu-3 SFT	NIAH-S3 @ 256K：75.2 → 83.8	普通 instruction SFT 可以提升指令能力，同时不伤长程召回。
HypeNet-9B · Tulu-3 SFT	NIAH-S3 @ 256K：52.0 → 86.8	退化不是 hybrid 模型 post-training 的必然后果。
OpenThoughts-3 / coding CoT	HypeNet-9B NIAH-S3 @ 256K：40.6 → 61.2，经 QK-Restore 恢复，同时 LCB 基本保留	CoT 引发的 routing 问题不只出现在数学，也可出现在代码 reasoning 数据。

3. 消融支持“Q 和 K 必须一起恢复”

Q-Restore 和 K-Restore 都有帮助，但不如 QK-Restore；V-Restore 不仅不能恢复召回，还可能伤 math。这个结果符合功能分工：Q/K 是 pairwise routing geometry，单独恢复一个会造成 Q-K mismatch；V/O 更多是 extraction，不是远距离寻址核心。

变体	NIAH-S2 @ 128K	NIAH-S3 @ 128K	MATH500 Avg@16	机制解读
QK-Restore	83.2	30.8	33.7	联合恢复 Q/K，路由几何最一致。
QK-Pro	79.0	26.0	34.2	更保留 post-SFT adaptation，但召回恢复略弱。
Q-Restore	79.0	24.4	34.2	只恢复 query，不足以重建匹配几何。
K-Restore	71.6	13.2	34.2	只恢复 key 更容易出现 Q/K mismatch。
V-Restore	36.6	15.2	32.1	value 侧不是长程 routing 的关键恢复点。

4. Pure softmax 模型不是同一种脆弱性

附录中，Qwen2.5-3B/7B 在 32K/64K NIAH 上相对稳定，Mistral-7B 在 S3 @ 64K 上从 81.0 掉到 4.8，但 QK-Restore 仅小幅恢复到 5.4。这个结果提醒我们：QK-Restore 不是万能补丁；当模型本身长上下文表征或 softmax 层冗余结构不同，Q/K 替换的效果可能很有限。

批判性阅读：这篇论文强在哪里，弱在哪里

强点 1：干预非常局部

只替换 Q/K 就能恢复召回，是比相关性可视化更强的证据。它把 failure mode 从“泛化退化”定位到“routing 参数漂移”。

强点 2：有 non-CoT 对照

Tulu-3 instruction SFT 对照排除了“后训练必然伤长上下文”的粗糙解释，使结论更聚焦于 CoT 数据结构。

强点 3：消融闭环完整

Q-only、K-only、V-only、QK-Pro、QK-Frozen 等消融形成了较清晰的功能分工证据。

弱点：理论更像解释框架

latent Markov、stationarity、score-function identity 等假设较强，梯度衰减推导更适合作为机制解释，而不是完全决定性的数学证明。

最容易误读的地方

不要把本文结论扩展成“CoT 一定破坏所有长上下文模型”。论文的主战场是 distilled hybrid linear-attention models；对于原生 hybrid、大规模商业模型、纯 softmax 模型、真实 RAG 工作负载，还需要更广泛验证。

NIAH 是诊断，不是真实长文理解的全部

Needle-In-A-Haystack 很适合测 exact recall 和远距离寻址，但真实任务可能包含多跳聚合、冲突证据、上下文压缩、工具检索、生成约束等复杂因素。NIAH 退化是重要红灯，不等于所有真实长上下文任务同幅度下降；反过来，NIAH 不退化也不能证明真实长文推理可靠。

QK-Restore 的适用条件很苛刻

它依赖成对的 pre-SFT 与 post-SFT checkpoint，也依赖明确知道哪些层是 full-attention / recall-critical。如果只有 post-SFT 模型，或者训练过程中架构、tokenizer、normalization、attention 实现发生变化，直接恢复 Q/K 就不一定可行。

术语解释：先把关键词对齐

CoT-SFTChain-of-Thought supervised fine-tuning，用带显式推理步骤的数据做监督微调。这里讨论的是训练权重变化，不是推理时简单要求模型“step by step”。

Hybrid attention混合注意力架构：部分层保留 softmax full attention，其他层使用 linear attention、recurrent state 或类似高效机制，以降低长上下文成本。

NIAHNeedle-In-A-Haystack，长上下文诊断任务：把目标事实放在很长文本中，测试模型能否准确找回。

Routing注意力中“选位置”的功能，主要由 Q/K 投影和它们的相似度决定；对长程召回来说，routing 比 value extraction 更像索引系统。

Extraction从被 attention 选中的 token 中拿回什么内容、如何转换输出，通常与 V/O 投影和后续 FFN 更相关。

QK drift后训练导致 W_Q/W_K 参数相对 pre-SFT 发生漂移；本文关注的是这种漂移是否让 attention 更偏近邻、削弱远距离寻址。

QK-ProQK-Restore 的 Procrustes 变体：不是直接完整回滚 Q/K，而是尽量接近 post-SFT 权重，同时约束 QK product 回到 pre-SFT 路由几何。

工程 / 研究启发：训练强 reasoning 模型时，必须把 routing stability 纳入验收

这篇论文对工程流程的启发很直接：如果你在训练 cheap long-context model，不能只看 MATH、AIME、GSM8K、LiveCodeBench 这些 reasoning 指标。CoT-SFT 后必须重新测长程召回，尤其是 128K/256K、困难 NIAH/RULER、多 needle、多字段、多跳检索。

评测面板

把 reasoning benchmark 与 long-context retrieval benchmark 放在同一张表里；用 context length、needle depth、检索难度分层，而不是只报平均分。

训练策略

保存 pre-SFT checkpoint；对 full-attention 层 Q/K 使用小学习率、正则、冻结或后训练恢复；加入 retrieval replay 防止 routing 被局部 trace 独占。

诊断指标

监控 Q/K drift、attention mean distance、long-range head attention、NIAH failure distribution，而不是只看 loss 和 reasoning pass@k。

更深一层的启发是：模型能力不是一个统一标量。Reasoning 变强可能来自局部推理轨迹和 value-side / FFN adaptation；retrieval 变弱可能来自少数 routing 参数的几何漂移。把这两类能力拆到参数功能层面，才有机会做“既要又要”的后训练，而不是在 benchmark 上做跷跷板。

可操作清单

CoT-SFT 前后固定跑一组 long-context recall regression tests。
对 hybrid 模型记录 full-attention layer index 和 recall-critical layer 排序。
保存 pre-SFT 权重，便于 QK-Restore 或更细粒度的 routing repair。
把 Q/K 学习率、LoRA target、冻结策略作为独立超参数，而不是默认跟全模型一起训练。
上线前用真实长文任务补充 NIAH，避免把 exact needle recall 等同于全部 long-context understanding。

边界与风险：这篇论文还没有证明什么

不要过度外推

论文主要验证 HypeNet 与 Jet-Nemotron 等 hybrid/distilled 模型；对更大规模原生 hybrid 模型、生产级 RAG、工具增强 agent、商业闭源模型，结论仍是强警示而不是已完备定律。

理论证明依赖理想化假设。 CoT-Markov 视角解释力很强，但真实 CoT trace 有跳步、回看、变量复用和长程约束，不能简单等同一阶 Markov 链。
QK-Restore 可能牺牲部分 task adaptation。 消融表明 QK-Pro 在 math 上略有优势，说明完整回滚 Q/K 可能丢掉一些有用更新。
真实任务需要额外验证。 NIAH 能抓住远程寻址失败，但不能覆盖多文档综合、信息冲突、工具调用、引用准确性等问题。
没有 pre-SFT 权重时难以直接修。 这对只拿到 post-trained checkpoint 的使用者尤其重要。

证据边界与资料索引

本文依据公开 X 帖、arXiv v1 论文正文、Hugging Face paper page 与官方代码仓库进行阅读整理。文中数值来自论文表格与摘要；本文没有独立复现实验训练，也没有验证所有 checkpoint 可用性。X 平台内容、论文版本和代码仓库后续可能更新，结论应以公开版本变更为准。