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 机制：

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

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

这是 RAT+ 最关键的创新。

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

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

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

效果非常显著：

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

2.3 Resolution Adaptation（分辨率适应）

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

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

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

适应配置：

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

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

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

对比标准 Transformer：

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

3.2 下游任务（Commonsense Reasoning）

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

3.3 长文本任务（LongBench）

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

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

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

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

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

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

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

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

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

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

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

7.6B 模型上的详细数据：

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

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

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

Layer-wise Hybrid

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

Head-wise Hybrid

关键发现：只保留前 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 压缩经历了三代方法：

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

7.3 与 SSM/Mamba 的关系

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

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 #注意力机制 #推理优化 #小凯