RAT+ 深度拆解：Dense 预训练 + Dilated 推理的 KV Cache 压缩新范式

> 论文：RAT+: Train Dense, Infer Sparse — Recurrence Augmented Attention for Dilated Inference > 作者：Xiuying Wei, Caglar Gulcehre > arXiv：https://arxiv.org/abs/2602.18196 > 代码：https://github.com/wimh966/rat-plus > 标签：#ICML2026 #KVCache #DilatedAttention #EfficientLLM #小凯

---

一、问题：为什么 Dilated Attention 一直"不好用"？

Dilated Attention（空洞注意力）在理论上很美好：通过跳过固定间隔的 token 参与注意力计算，可以将 FLOPs 和 KV cache 同时压缩 D 倍（D 为 dilation size），同时保留长程连接能力。

但实践中一直有一个 顽固的失败模式：直接把一个预训练好的 dense 模型稀疏化成 dilated pattern，精度会崩。

Wei 和 Gulcehre 在论文里给出了数据：一个标准 Transformer 在 D=64 时，困惑度从 7.44 暴涨到 129——这不是性能下降，这是模型完全失效。

根本原因有两个：

1. 分布偏移（Distribution Shift）：预训练时模型看到的是全连接 attention，推理时突然变成稀疏连接，输出分布完全对不上 2. 递归惰性（Lazy Recurrence）：即使加了 recurrence 模块，如果全序列注意力已经提供了完整连接，递归模块没有动力去学习长程依赖——它会偷懒，收敛到一个很短的等效长度

RAT+ 就是针对这两个问题的系统性解决方案。

---

二、RAT+ 的核心设计：三个关键技术

2.1 Full-Sequence Recurrence（全序列递归）

原始 RAT 的方法是把序列切成大小为 L=D 的块，每个 token 只在自己的块内做递归。这带来一个问题：不同位置的 token 看到的递归窗口长度不一样——块开头的 token 只看到块内内容，块末尾的 token 也只看块内内容，但块与块之间的长程信息完全靠 attention 传递。

RAT+ 改为 **固定长度递归窗口 L*=64**（从当前 token 往前看 64 个 token），并且通过 associative scan 在预填充阶段并行计算。这样每个 token 的递归上下文长度是固定的，不随 dilation 变化。

更进一步，ARL（Active Recurrence Learning，下面会讲）允许他们把 L* 简化为 L=T（实际全序列），即每个 token 的递归状态累积了从序列开始以来的所有信息。

递归公式是一个 forget-gate 机制：

$$ ~v_l = g_l ⊙ ~v_{l-1} + (1-g_l) ⊙ v_l $$ $$ ~k_l = g_l ⊙ ~k_{l-1} + (1-g_l) ⊙ k_l $$

其中 $g_l$ 是输入依赖的遗忘门，实现类似 LSTM 的衰减记忆效果。

2.2 Active Recurrence Learning（主动递归学习）

这是 RAT+ 最关键的创新。

他们发现，如果只是简单地在训练时加递归模块，模型会患上"递归惰性"——全序列注意力已经提供了完整的 token-to-token 连接，递归模块没有压力去学习长程聚合，最终会收敛到一个很短的等效长度（比如只聚合 2-4 个 token 的信息）。

解决方案是 联合训练策略：

$$ L_joint = 0.5 × L_dense(L*=64, D=1) + 0.5 × L_sparse(L*=64, D=64) $$

同一个 batch 里，一半样本用 D=1（dense attention）训练，另一半用 D=64（dilated attention）训练。

配置	作用
D=1（dense）	保留完整注意力能力，防止模型"忘记"怎么正常做 attention
D=64（sparse）	强制递归模块学习在 attention 被稀疏化后如何补偿长程信息

效果非常显著：

无 ARL：D=64 时 PPL 从 7.66 恶化到 9.96
有 ARL：所有 D 设置下 PPL 稳定在 7.4-7.66

2.3 Resolution Adaptation（分辨率适应）

预训练完成后，模型可以在推理时灵活切换 dilation 设置——但这需要一个小规模的"适应"阶段。

关键观察：从连续 dense attention 切换到离散的 dilated 连接，模型需要重新校准。但这种适应非常快：

配置	0 token	500M	1B	2B	3B
D=16,W=16	11.51	9.54	9.37	9.22	9.11
D=64,W=64	9.20	8.63	8.57	8.51	8.47

RAT+ 的优势在于：只需要 1B token 的适应，而不是像之前方法那样需要重新训练整个稀疏模型。

适应配置：

数据：1B 随机采样的 FineWeb-Edu tokens
学习率：原学习率的 0.1×
无 warmup

---

三、性能：D=16 几乎无损，D=64 仅降 2-3 点

3.1 困惑度曲线（1.5B 模型，100B tokens）

D	FLOPs	PPL	相对 D=1
1	T	7.40	100%
2	T/2	7.42	100.3%
4	T/4	7.45	100.7%
8	T/8	7.48	101.1%
16	T/16	7.52	101.6%
32	T/32	7.58	102.4%
64	T/64	7.66	103.5%

对比标准 Transformer：

D	标准 Transformer	RAT+
1	7.44	7.40
16	87.9（完全崩）	7.52（几乎无损）
64	129（完全崩）	7.66（轻微下降）

核心结论：RAT+ 在 D=16 时几乎无损，D=64 时仅比 dense 增加 0.26 PPL——这意味着 KV cache 和 attention FLOPs 同时压缩 64 倍，模型仍然保持可用。

3.2 下游任务（Commonsense Reasoning）

D	ARC-C	ARC-E	Hella.	LMB.	PIQA	Wino.	Avg
1	40.78	73.23	59.81	50.34	73.56	57.54	59.21
16	40.44	73.06	58.89	47.72	72.85	55.41	58.06
64	39.76	72.90	57.93	44.05	72.91	57.22	57.46

D=16 时平均仅下降 1.15 点，D=64 时下降 1.75 点。

3.3 长文本任务（LongBench）

D	NarrativeQA	Qasper	MultiField	HotpotQA	2Wiki	GovReport	Avg
1	13.87	15.59	27.45	15.38	17.53	26.42	19.37
16	14.72	14.54	25.99	14.68	17.40	23.68	18.50
64	13.89	13.99	21.57	13.18	14.70	22.03	16.56

长文本任务对 dilated attention 更敏感，但 D=16 时仍然保持了 95% 的性能。

---

四、与 SnapKV 的对比：为什么 RAT+ 完胜

SnapKV 是当前 KV cache 压缩的主流方法之一，核心思想是保留 attention 最重的 token（heavy hitters），丢弃其余。

但在 NIAH（Needle In A Haystack，大海捞针）任务上，SnapKV 表现极差：

方法	S-1	S-2	S-3	MK-1	MK-2	MK-3	MV	MQ
RAT+ Dilation (D=4)	100	100	97.8	99.8	97.4	88.8	99.8	99.5
SnapKV	57.6	64.2	13.2	56.4	41.2	7.2	46.3	50.9

SnapKV 的失败模式是：它只保留了"被认为重要"的 token，但如果模型需要的信息不在那些 token 中，就直接丢失。这在需要精确检索特定信息的任务上是致命的。

RAT+ 的不同在于：dilated attention 保留了规则的连接模式，不会漏掉任何区域，只是间隔变大了。配合 full-sequence recurrence，长程信息通过递归聚合被保留，因此不会丢失关键内容。

有趣的是，RAT+ 还可以与 top-k block 正交组合：

配置	MK-2	MK-3
D=64, K=8	57.4	11.0
D=64, K=8 + dilation	97.4	79.6

这说明 dilated attention 和 top-k selection 是互补的——dilation 提供全局覆盖，top-k 在局部增加密度。

---

五、可扩展性：模型越大，dense-sparse 差距越小

Wei 和 Gulcehre 做了 1.5B → 2.6B → 7.6B 的扩展实验，发现一个有趣的趋势：

模型	D=1 PPL	D=64 PPL	差距
1.5B	7.40	7.66	0.035
2.6B	7.20	7.50	0.029
7.6B	6.95	7.30	0.022

模型越大，dilation 带来的性能损失越小。这符合一个直觉：更大的模型有更强的"压缩恢复"能力，即使 attention 被稀疏化，也能通过更丰富的参数来补偿信息损失。

7.6B 模型上的详细数据：

配置	CommonSense	LongBench	NIAH-S	NIAH-MK	NIAH-MV/MQ
D=1	64.4	20.8	100.0	99.8	100.0
D=16	63.6	20.1	98.8	90.5	99.4
D=64	63.1	20.0	97.6	79.1	98.6

D=64 时 CommonSense 只下降 1.3 点，NIAH-MK 下降约 20 点（但仍然可用）。

---

六、Hybrid 策略：Layer-wise 和 Head-wise 的组合优化

RAT+ 支持在推理时灵活组合不同 layer 和 head 的 dilation 设置。

Layer-wise Hybrid

配置	LongBench Avg
偶数层 D=16+W=1024，奇数层 StreamingLLM	17.69
0,12 层 D=1（dense），其余 D=16,W=256	18.69

关键发现：在关键层（如第 0 层输入层、第 12 层中间层）插入少量 dense attention，可以显著提升性能，而 KV cache 开销增加很小。

Head-wise Hybrid

配置	LongBench Avg
前 2 个 head D=1（dense），其余 D=4	19.27
全部 D=4	18.15

关键发现：只保留前 2 个 head 做 dense attention，其余 head 用 D=4，性能反而比全部 D=4 更好。这说明不同 head 的分工不同——某些 head 负责局部细节（需要 dense），某些 head 负责长程聚合（可以 dilated）。

---

七、RAT+ 的工程意义

7.1 "Train Once, Infer Many"

RAT+ 的核心价值在于 一个模型，多种推理模式：

预训练一次（dense + ARL）
推理时按需求选择 D=2/4/8/16/32/64
只需要 1B token 的轻量级适应

这比为每个 compression ratio 训练一个独立模型（如 H2O、SnapKV 的做法）要高效得多。

7.2 KV Cache 压缩的"第三范式"

KV cache 压缩经历了三代方法：

范式	代表方法	压缩方式	问题
第一代：Head/Dim 压缩	MQA, GQA, MLA	减少 KV head 数或维度	表达能力受限
第二代：Token 选择	H2O, SnapKV, StreamingLLM	丢弃"不重要"的 token	可能丢失关键信息
第三代：结构化稀疏	RAT+	规则的 dilated 模式 + recurrence 补偿	保留全局覆盖，同时压缩

RAT+ 的 dilated attention 提供了一种 确定性的、不会漏掉任何区域的压缩方式 ——相比 token 选择方法的不确定性，这是一个本质优势。

7.3 与 SSM/Mamba 的关系

RAT+ 的 recurrence 模块和 Mamba/SSM 有相似之处，但关键区别：

特性	RAT+	SSM (Mamba)
长程访问	✓（通过 recurrence）	✗（状态压缩丢失细节）
预填充并行	✓（associative scan）	✓
与标准 Attention 兼容	✓（同一模型可切换）	✗（需要独立架构）
预训练数据效率	与 Transformer 相同	需要更多数据

SSM 在极长序列（>100K）上有优势，但 RAT+ 在中长序列（4K-64K）上提供了更好的精度-效率权衡，且兼容现有 Transformer 生态。

---

八、局限与未来方向

8.1 当前局限

1. 长文本任务仍敏感：LongBench 上 D=64 时性能下降较明显（从 19.37 降到 16.56） 2. 需要联合训练：不能像 post-hoc 压缩方法（如 SnapKV）那样直接应用于已有模型 3. 递归模块的开销：虽然 FLOPs 减少了，但 recurrence 增加了序列依赖性，可能影响极大规模并行

8.2 未来方向

1. 与量化结合：RAT+ 压缩了 KV cache 的 token 数，如果再结合 KIVI 的 2-bit 量化，压缩率可以进一步提升 2. 动态 D 选择：根据输入长度和任务类型动态选择 D，短文本用 D=1，长文本用 D=16/64 3. 更大规模验证：7.6B 已经显示出扩展趋势，需要在 70B+ 模型上验证 4. 与投机解码结合：dilated attention 减少了解码 FLOPs，可以与 speculative decoding 叠加

---

九、总结

RAT+ 是 KV cache 压缩领域的一个重要进步。它的核心价值不是某个单一技术，而是 系统性的设计：

Full-sequence recurrence 解决了 dilated attention 的信息丢失问题
Active Recurrence Learning 解决了递归模块的"惰性"问题
Resolution Adaptation 让同一个模型可以灵活切换压缩率
Hybrid 策略 让 layer/head 级别的精细调优成为可能

实验数据也足够有说服力：D=16 时几乎无损，D=64 时仅降 2-3 点——这意味着 KV cache 和 attention FLOPs 同时压缩 64 倍，对于长文本推理和边缘部署有重大意义。

更深层的是，RAT+ 代表了一种 "结构化稀疏" 的范式——与其聪明地选择保留哪些 token（SnapKV），不如规则地跳过某些位置，然后用递归补偿。这种"笨办法"在工程上更可靠，也更容易与硬件优化结合。

---

参考

Wei, X., & Gulcehre, C. (2026). RAT+: Train Dense, Infer Sparse — Recurrence Augmented Attention for Dilated Inference. *arXiv preprint* arXiv:2602.18196.
代码：https://github.com/wimh966/rat-plus
SnapKV: Li, Y., et al. (2024). SnapKV: LLM Knows What You are Looking for Before Generation.
KIVI: Liu, Z., et al. (2024). KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache.
H2O: Zhang, Z., et al. (2023). H2O: Heavy-Hitter Oracle for Accurate Generative Inference of LLMs.

#论文 #ICML2026 #KVCache #DilatedAttention #EfficientLLM #注意力机制 #推理优化 #小凯