When Hidden States Drift: Can KV Caches Rescue Long-Range Speculative Decoding?

1 · 动机:long-range decay 的第二种解释

1.1 历史脉络:从 hidden-state reuse 的胜利到长程墙

2024 年开始,hidden-state reuse 已经是 speculative decoding 的事实标准:EAGLE / EAGLE-2 / EAGLE-3 把 target 模型的 top-layer hidden state 作为 drafter 的输入,Medusa / MTP / DeepSeek-V3 内部的 multi-token prediction head 也走同一条路。这条路赢在两点:① drafter 只要 1 层就足够预测,因为最重的 prefix 编码已经在 target 那里做完了;② 训练目标可以直接用 hidden state 的 L2 + cross-entropy,信号丰富。

但所有这些 drafter 都共享同一个失败模式 —— long-range decay:接受率 α_k 随推测步 k 单调下降。论文 Figure 1 给出的 Qwen3.5-4B 数字非常具体:

这个曲线决定 draft tree 能伸多深。如果 α 在 k=6 已经低于 0.6,你再多猜也几乎全被 reject。所以"修长程"是 speculative decoding 的核心瓶颈。

1.2 别的方案为什么不够

领域里对 long-range decay 的标准解释只有一个: train-inference mismatch。训练时 drafter 看到的是 target 的真实 hidden state,推理时它必须 condition 在自己递归生成的 hidden state 上,这两者会逐步 drift。 HASS / EAGLE-3 引入 autoregressive TTT(test-time training)把训练时也喂 drafter 自己的 drift trajectory,确实把 α 抬高一截。

但作者点了一刀:就算用了 TTT,decay 依然存在(Figure 1 那两条曲线都是 TTT-trained 的)。这意味着 drift 只解释了一部分,还有一个独立的损失源没人讨论。下表把所有候选解释摆出来:

1.3 为什么这事不平凡:context preservation 是另一种透镜

核心 insight 用一句话讲透:target 的 ℎ_𝑡^L 是一个 query-biased compression,它对 t+1 优化,所以对 t+k 而言,被它压扁掉的那些 token 信息可能正好是关键。

这是一个 information-theoretic 视角而不是 distributional 视角。前人都在解决"训练分布和推理分布不一样"这个 distributional gap;作者改去问"reuse 哪种 representation 才能保留足够 prefix 信息"。一旦把问题这么 reformulate,KV cache 就成了一个非常自然的 candidate:它是 attention 之前的 raw token-wise 表征,没有被 q_t 加权聚合过,等价于把"recovery 问题"换成"function approximation 问题"——drafter 不必"挖出"被压扁的信息,而只需要估计未来位置的 q′ 然后重新对 KV 做 attention。

这个角度的非平凡之处在于 reframe 而不是新方法本身。GLIDE+CAPE / LongSpec 早就用过 cross-attention 接 target KV,但都把它当作"减少 drafter 计算"或"长 context 内存效率"的工程手段;没人把它当作长程衰减的解释来做受控对照。本文的贡献是先建一个统一对照框架(KVShot)把 hidden-only / KV-only / 混合放在同一 TTT pipeline 下比 ,然后把"是不是 representation 问题"这个 hypothesis 实证出来。

2 · 背景速查

关键公式回顾

这是 self-attention 的标准式。论文的整个观点都建在这一个公式上:α_𝑖 由 𝑞_𝑡 决定,而 𝑞_𝑡 是为了预测 𝑡+1 优化的;只要你信这一句,后面的 KV-Reuse Hypothesis 就是它的逻辑后果。

HS 误差有两块:压缩损失(k 无关、靠 target 自己产生)+ drift(k 越大越严重);KV 误差则把第一块清零,只剩下 query 估计这一块。TTT 只能压第二块,所以 hidden-only 即便 TTT 训得再好,也压不掉 E_comp。

3 · Hidden state 是有偏压缩 (concrete walk-through)

这一节把 §1.3 的论证落到一个具体例子。设想一段 prefix:

「The patient was prescribed aspirin for chest pain.
 He visited the clinic on Monday.
 Today he reports nausea and tells the doctor: ___」

target 在最后一个位置 𝑡 计算 𝑞_𝑡 是为了预测下一个 token(可能是 "I", "the", "my")。它的 attention 大概率会重权 "doctor"、"reports"、"nausea"(都是局部句法 / 语义相关),而 "aspirin"、"Monday" 拿到 α≈0.01。所以 ℎ_𝑡 这个 4096 维向量里关于 aspirin 的信息已经几乎被擦掉了。

现在 drafter 想猜未来若干步。假设第 k=4 步真实下一个 token 是 "But aspirin upset my stomach"——这一步未来的 q_𝑡+4 会强烈 attend 到 "aspirin"(α=0.65)。drafter 只拿到压扁后的 ℎ_𝑡,它必须在 ℝ⁴⁰⁹⁶ 里"挖出"那个被弱权值压下的信息。这是ill-conditioned 的反问题。

反向论证:为什么"光做多层 fusion (EAGLE-3 那样)"还不够? 不同层的 q 不一样,某一层 α=0.01 的 token,在另一层可能 α=0.3,所以 multi-layer fuse 能把部分信息找回来。但只要存在跨所有 fused layer 都被弱 attend 的 token,它在 fused hidden state 里仍然是被压扁的。而 KV reuse 是不做聚合地把所有 token 都暴露出来,理论上更彻底。

4 · KV-Reuse Hypothesis 与三个 prediction

把 §3 的论证形式化,作者抛出 KV-Reuse Hypothesis 并拆出三条可证伪的预测,后面 §5–6 的所有实验都是在 verify / falsify 这三条。

三条加起来直接暗示了正确的设计形态:不是替换,而是混合——hidden-state path 锚定短程,KV path 做长程的 correction。这就是 §5 的 hybrid drafter 的来源。

5 · KVShot 框架: 三种 reuse 的受控对比

KVShot 不是一个模型,是一套诊断性测试床。所有变体都在同一条 EAGLE-3 TTT pipeline 下训练,target 都是 Qwen3-8B,这样比较的是 representation 选择而不是别的 confound。

5.1 KV-Only Reuse:把 hidden 完全踢掉

架构: drafter input 只剩 input embedding,然后做 cross-attention 到 target 的 KV。多层 KV 怎么注入需要选:

读法:linear projection 比 head concat 好(单 query attend 一个 unified KV 比 attend 多个 disjoint KV 容易);RoPE 修复几乎没区别(说明位置不是主瓶颈)。但都比 EAGLE-3 差很多:α₀ 永远过不了 0.50。这就是 P3 的实验证据——短程 hidden 提供的语义先验是 input embedding 给不出的。

具体 trace:一个 1 层 drafter 在 KV-only 下做了什么?

drafter 收到 input_emb ∈ ℝ⁴⁰⁹⁶ → 1 个 transformer block → 出 q′ ∈ ℝ¹²⁸(per head)。这个 q′ 是 input_emb 的 线性(再加 LayerNorm)投影,本质上没法 capture "我现在要预测的是 t+3 位置的 query"这种复杂依赖。所以"加深"是必需的。

5.2 Depth scaling:验证 P2

读法非常微妙:

1. 1→2 层 ΔMAT=+0.39 是最大跳跃 → 第二层让 q′ 变成 context-aware 而不是简单的 input embedding 线性变换。这正面回答了 P2。
2. 2→3, 3→4 几乎平了(+0.08, +0.03),递减很快。
3. 4 层 KV-only(MAT 2.34)仍稍输 1 层 EAGLE-3(2.37),代价还更高。
4. 但 retention 80.6% 已经显著超过 EAGLE-3 73.5%,α₆ 0.495 反超 0.469。这就是 P1 落地。

5.3 Hybrid Reuse:门控 delta correction

架构: 在 EAGLE-3 self-attention 旁挂一条平行的 cross-attention 分支接 target KV。两路输出用 gated delta rule 融合:

核心设计点:

• 不做直接 sum: 𝑛+𝑜 形式允许 cross-attention 覆盖 self-attention,可能毁掉 EAGLE-3 已有的短程精度。
• 用 delta = 𝑜 − 𝑛: 把 cross-attention 的输出当作"差异 / 修正",而不是绝对答案。
• shared gate: 𝑔 ∈ [0,1]^d 跨所有 draft step 共享(简单 + 验证 mechanism 是否本身有效)。
• warm-start: self-attention / MLP / norm 全部从 EAGLE-3 checkpoint 加载,只有新加的 cross-attention proj、gate、𝑊_Δ 是随机初始化。

5.4 Hybrid 结果:第一次正向

三条 takeaway:

• Hybrid (ckpt init) +0.17 MAT,α₆ 抬 +9.6%,retention 77.3% > 73.5%——这是论文里第一次明确正向、且每个 k 都涨。
• Init 重要: warm-start 比 random 多 +0.10 MAT,说明 self-attention anchor 给了 cross-attention 一个稳定的 baseline 去学修正。
• Self-attn anchor 不可去: cross-only(即便 ckpt init)直接掉到 EAGLE-3 baseline 之下。"KV reuse 替代 hidden reuse"这条路走不通,只能"补充"。

6 · 实验关键结果与读数

端到端的"小坑":+0.17 → +0.03

Step-wise 看一切都对——直到把训练数据从 70k(ShareGPT)放大到 280k(ShareGPT + UltraChat regenerated)再用 HF tree decoding (8,10,60) 测 MAT:

+0.03 加上 cross-attention 多带来的 5–10% drafting latency,基本 wash。这是这篇论文最诚实的一段:他们把负面结果放在正文,然后专门写了 §7 拆为什么。

为什么从 +0.17 缩到 +0.03? 三条合力:

1. 更大数据让 EAGLE-3 自己也强了:baseline MAT 4.43→5.01,那个 +0.17 里有相当一部分本来就是 hidden-state 也能从更好数据中拿到。
2. Tree verification 把分步差异压扁: (8,10,60) 已经探索很多候选路径,即使每步 α 略低也可以 fall back 到 sibling 候选。
3. Cross-attn 的 latency 直接吃掉收益。

7 · 三个结构性 bottleneck

§7 是一篇潜在 followup 的 roadmap,作者把 autoregressive TTT 不适合 KV reuse 的原因拆成三个独立 issue。

7.1 Query estimation 实际比预想更难

理论上 E_q 是函数逼近问题,听起来"更简单"。实际:target 的 q_t+k 是 L 层非线性变换 over full prefix 的产物。1 层 drafter 的 q′ 只是 input embedding 的线性变换,根本没有组合性。Tab.2 已经显示从 1→2 层是大跳,但 4 层也碰不到 1 层 EAGLE-3——hidden-state reuse 实质上"借走"了 target 那 L 层算 q 的工作。这是 KV reuse 短程一直输的根源。

7.2 Draft-side KV projection 训练信号稀疏

这是最微妙的一条,值得展开。Hybrid 的 cross-attention 分支里有两个新参数 𝑊^cross_𝐾, 𝑊^cross_𝑉。它们的工作是给"draft 自己生成的"那部分 token 算 K/V(prefix 部分的 K/V 直接 copy 自 target,免费)。

问题:autoregressive TTT 一个 batch 一般只有 K=7 个 draft token,而 prefix 长度可能几百上千。所以这两个 projection 矩阵每步只有 7 个 token 的梯度,而 𝑊_𝑄 每个位置都更新。

7.3 Gate-induced gradient starvation

Hybrid (ckpt init) 时 self-attention 是预训练好的,cross-attention 完全随机。早期 cross 输出近乎噪声,SGD 最快的 loss 下降方向就是"把 gate 𝑔 推到 0",让模型退化为 EAGLE-3。

有趣的细节:小 ḡ 不等于 cross 被忽略,因为 ḡ 是 element-wise 平均,某些 dim 的 gate 可能非常大。即便 ḡ≈0.1,Hybrid 仍能拿到 +0.17 MAT,说明这个机制的"上限"远超目前的训练状态。

7.4 共同根因:autoregressive 本身

三条 bottleneck 都源于"一次生成一个 token"这个序列结构:

• 串行生成 → drafter 没法堆深(开销线性翻倍)→ q′ 估不准
• 每步 K=7 → KV proj 信号稀疏
• warm-start + 序列训练 → gate 早期看到的 cross 输出全是噪声

作者直接指向 block-wise training(具体引用 DFlash, arXiv:2602.06036,见07_DFlash):一次并行生成一整块 draft token,既允许 drafter 加深(mask-token parallel 摊薄成本),又给 KV proj 提供 dense diverse 梯度,还能 warm-start 时同时让两路获得有意义信号。这是这篇 paper 留给社区的"下一步该这么做"。

8 · 与同类工作对比

关于 EAGLE-3 evolution 整条线和 DFlash 怎么用 block diffusion,可以参看姊妹笔记 02_EAGLE_evolution_DFlash 与 07_DFlash。

9 · 局限 / 个人 take / 待验证问题

局限

• 负面结果:核心 +0.03 HF MAT 实在太小,严格意义上无法支持端到端 speedup claim。论文也承认。
• Single target model:全部 Qwen3-8B,没在更大模型(70B+)上验证,而 hidden compression 的严重程度可能与模型规模 / depth 相关。
• KV layer 选择:L_s=3 默认值,Appendix 说"varying has negligible effect",但没系统扫,我会想看 L_s=L (全部层) 的极限。
• 没真正实现 block-wise 验证:Section 7.4 把 block training 当"可能解",但没做实验对照,只能 cite DFlash。这个 promise 是欠的。
• Tree verification 的影响没拆开:HF (8,10,60) 把分步差异压扁这个 claim 是定性的,没给 ablation,但理论上很合理。

个人 take

• 这篇论文最大的价值是 reframe:把 long-range decay 从"训练目标问题"重新定义成"reuse 表征问题",是一个干净的 conceptual contribution,即使具体方法没赢。
• 三个 bottleneck 写得非常诚实,任何想接着做 KV-aware drafter 的人都应该先读 §4。Gate starvation 是我之前没意识到的、warm-start 加 sigmoid gate 的常见陷阱,值得抽出来单独研究 mitigation(比如 GLU 风格 hard init、gate annealing、cross-branch warmup-then-merge)。
• "hidden state 是 query-biased compression"这个观点在很多其它场景也适用,比如 RAG retriever 训练、prefix tuning,都可以重新审视。
• 个人猜测真正能拿到端到端的关键是 block-wise + KV reuse(DFlash 的 5 层 + hybrid 形态)。这个 paper 实际上为这个 combination 提供了独立的 motivation。

待验证问题

1. 把 Hybrid 移植到 DFlash 的 block diffusion pipeline,KV gain 会不会从 +0.03 涨到显著?
2. 是否存在 KV-only + 大 drafter(8 层)+ block training 这条 path,能彻底超过 EAGLE-3? §5 depth scaling 已经接近,管线一改也许就过线。
3. Gate starvation 的诊断:gate 只是延迟开了,如果训练 3× 更长,Hybrid 端到端 gain 会不会回来?
4. KV proj 稀疏问题:能否在训练时人为注入更多 draft 位置(比如插入 noise tokens 模拟 draft prefix),让 W_K,V^cross 拿到更多多样性梯度,而不动 inference?
5. 更大 target model(Qwen3-72B)上,E_comp 是更严重还是更轻?直觉上 hidden 越宽信息密度越高,可能 compression loss 反而摊薄,但 layer 越多 q_𝑡+𝑘 离 q_𝑡 越远,趋势相反。需要实验。
6. Retention 这个 metric 是不是合适的"长程决策"指标?如果 α₀ 本身已经低,retention 高没意义(论文也承认 "No target info" retention 99.6% 但 MAT 1.31)。需要一个组合指标。

10 · Memory Points