[ArXiv 2025] IndexCache：打破稀疏注意力的二次方“隐形”瓶颈

TL;DR

清华大学与 Z.ai 团队联合发布的 IndexCache 揭示了稀疏注意力模型（如 DeepSeek-V3 架构）中的一个巨大优化空间：负责筛选 Token 的“索引器”在每一层都在做重复劳动。通过简单的跨层索引复用，IndexCache 杀掉了 75% 的索引计算，在 200K 上下文下实现了近 2 倍的加速，且模型能力几乎无损。

1. 痛点：被忽视的 $O(L^2)$ “索引税”

在处理长文本时，稀疏注意力（Sparse Attention）是当前的标配。以 DeepSeek 为例，它引入了一个轻量级的 Indexer（索引器），先为每个 Query 选出最相关的 Top-k 个 Token，再进行核心计算。

虽然核心计算变快了，但 Indexer 本身依然是 $O(L^2)$ 复杂度。当文本长度 $L$ 增加到 200K 甚至 1M 时，每一层都要跑一遍 Indexer 的开销变得不可忽视。 如图 1 所示，随着长度增加，Indexer 的耗时占比急剧上升，成为新的工程瓶颈。

2. 核心直觉：跨层冗余与稳定选择

作者发现：相邻两层选出的 Top-k Token 高度重合（重合度 70%-100%）。 这产生了一个直观的构想：我们为什么不让第 $N+1$ 层直接“抄”第 $N$ 层的索引结果呢？这就是 IndexCache 的物理直觉。

• Full 层：正常计算 Indexer 并把结果存入 Cache。
• Shared 层：直接读取 Cache 里的索引，跳过本层的 Indexer 计算。

左图为标准 DSA，右图为 IndexCache。唯一的区别是一个简单的条件分支。

3. 方法论：如何优雅地“偷懒”？

直接乱跳层会导致精度崩塌。论文提出了两套方案：

3.1 训练无关（Training-free）：贪心搜索

如果模型已经训练好了，作者发现不能简单地“每隔几层跳一次”（Uniform Interleaving）。

• Greedy Selection：在校准集上，一层一层地试：如果把这一层的 Indexer 删掉，对最终 Loss 影响大吗？通过这种方式选出最不敏感的层进行合并。
• Insight：早期层通常比深层更敏感，因为误差会沿网络积累。

3.2 训练相关（Training-aware）：多层联合蒸馏

如果在训练期就知道要复用索引，效果会更好。 作者提出了 Multi-layer Distillation Loss。让 Full 层的那个 Indexer 同时去拟合它后面所有 Shared 层的注意力分布（寻找“公约数”）。数学推导证明，这等同于将 Indexer 蒸馏到这几层注意力分布的中心点。

4. 实验战绩：速度翻倍，智力在线

在 30B 规模的模型上，IndexCache 展示了统治级的性能：

• Prefill 加速：200K 长度下加速 1.82x。
• Decode 吞吐：由于减少了每步的 Indexer 扫描开销，吞吐量提升了 51%。
• 精度保持：在 LongBench v2、AIME（数学）和 GPQA（逻辑推理）等任务上，保留 1/4 索引层后的表现与全索引模型几乎一致。

注意：在更长的 Context 场景下，IndexCache 的边际收益更高。

5. 深度洞察与总结

IndexCache 的成功告诉我们，深度学习模型在筛选信息时存在大量“认知冗余”。

• 局限性：虽然 Indexer 计算省了，但核心稀疏注意力（Sparse Attention）的算力开销和 KV Cache 的 IO 开销依然存在。它解决的是“如何选 Token”的成本，而非“处理 Token”的全部成本。
• 行业价值：随着 DeepSeek-V3 等模型将 DSA 这种动态稀疏架构推向主流，IndexCache 提供了一种极其廉价且高效的端侧/云端部署优化手段。它不再依赖复杂的 Full Attention Oracle，是真正面向生成环境的工程进化。

未来，我们可能会看到这种“索引缓存”机制不仅在层间复用，甚至在时间步（Auto-regressive steps）之间进行复用，进一步榨干推理性能。