📊 statistics

WildCat: Near-Linear Attention in Theory and Practice

WildCat 是一种高精度、低成本的注意力压缩方法，它通过利用随机枢轴乔莱斯基（Cholesky）子采样来选择并加权一个小的核心集（coreset），从而实现了近乎线性的时间复杂度以及超多项式级的误差衰减，在图像和语言任务的理论保证与实际性能方面均优于先前的近似方法。

原作者： Tobias Schröder, Lester Mackey

发布于 2026-06-02

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Tobias Schröder, Lester Mackey

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图回忆一段长达数小时的朋友间的对话。在标准的 AI 模型（例如驱动当今聊天机器人的模型）中，试图一次性记起整场对话中的每一个字，就像是试图在脑海中同时装下一座图书馆的藏书，并同时撰写一本新书。你添加的词汇越多，难度就越大，所需的记忆力呈**二次方（quadratic）**增长。这意味着，如果你将对话长度增加一倍，记忆所需的精力并不仅仅是增加一倍，而是会变为原来的四倍。最终，计算机会耗尽内存或在思考上花费过长时间。

这篇论文介绍了一种名为 WILDCAT（加权迭代低秩分解核心集注意力，Weighted Iterative Low-rank Decomposition for Coreset ATtention）的新方法来解决这个问题。以下是它的工作原理，使用简单的类比进行说明：

问题所在：“图书馆”瓶颈

把 AI 的记忆想象成一座巨大的图书馆。当 AI 需要回答一个问题时，它通常需要扫描每一本书（对话中的每一个词）来寻找相关信息。

旧方法： 如果图书馆有 1,000 本书，AI 就要检查 1,000 本；如果图书馆增长到 10,000 本，AI 就得检查 10,000 本，但其中的“精力消耗”会爆炸式增长。这就像是通过对比每一根稻草与其它所有稻草，试图在一堆干草中寻找一根特定的针。

解决方案：“专家小组”（WILDCAT）

WILDCAT 改变了策略。它不再试图完美地记住所有内容，而是意识到对话中的并非每个词都同等重要。有些词至关重要，而有些词只是填充物。

WILDCAT 主要做了两件事：

选取“核心集”（专家小组）：
想象你面对着 10,000 人的庞大群众，而你需要了解整个房间的情绪。与其采访每一个人，WILDCAT 使用了一种聪明且快速的算法（称为“随机枢轴乔莱斯基分解”，Randomly Pivoted Cholesky）来挑选出一个极小的、具有代表性的群体，比如只有 50 人。
- 神奇之处： 它并不是随机挑选，而是挑选出那些最能代表整体的最重要的成员。这就像是从一个城镇中选出最有发言权且最具多样性的成员来组成一个议会，让他们代表所有人发声。
赋予声音权重：
一旦拥有了这个 50 人的小团体，WILDCAT 并不会平等对待他们。它会为每个人分配“权重”。
- 类比： 如果这个小组成员中有一位非常大声、有主见的领导者，那么他的声音分量就更重；如果另一位成员很安静，他的声音分量就较轻。WILDCAT 会计算出这 50 个人完美的“音量”，使得当你只听取这 50 人的声音时，听起来的效果与聆听整个 10,000 人的群众完全一致。

为什么这意义重大

论文声称 WILDCAT 实现了三大突破：

速度（近线性）： 因为它只倾听 50 人的小团体而不是 10,000 人的大群众，所以思考所需的时间增长得非常缓慢。如果你将对话长度增加一倍，所花费的时间只会略微增加，而不是大幅度增加。这就像是从阅读一本书的每一页，转变为只阅读一份高度准确的摘要。
准确度（超多项式）： 通常情况下，当你对事物进行总结时，会丢失细节。但 WILDCAT 很特别，它声称几乎不会丢失任何重要细节。论文从数学上证明，随着对话变得越来越长，误差（缺失的细节）缩减的速度极其之快——比目前使用的几乎任何其他方法都要快。
实用性： 作者不仅做了数学推导，还在强大的图形处理器（GPU）上构建了一个可运行的版本。他们在以下领域进行了测试：
- 生成图像： 它生成的图像质量高，且比其他方法更快，且不会让图像看起来模糊或扭曲。
- 图像分类： 它识别照片中物体的能力与完整的、缓慢的方法一样出色，但速度更快。
- 长对话： 它允许语言模型在不耗尽内存的情况下，记住更长的聊天历史，其表现优于其他“记忆压缩”技巧。

核心结论

WILDCAT 就像一位超级高效的秘书，能够听完一场 10 小时的会议，瞬间挑出最重要的 50 个时刻，并为它们分配正确的权重，然后生成一份如此准确的会议摘要，以至于老板感觉自己听到了全程，但秘书其实只写了几页纸。

这使得 AI 模型能够在不需要超级计算机的情况下，处理更长的对话和更大的任务，让它们运行得更快、成本更低，并具备记忆更长上下文的能力。

技术摘要：WILDCAT：理论与实践中的近线性注意力机制

问题陈述

注意力机制是现代 Transformer 架构的基石，能够对长距离依赖进行建模。然而，其部署成本极高，因为资源需求随输入序列长度 $n$ 呈二次方级增长 ( $\Theta(n^2d)$ )。尽管存在许多近似注意力方法，但在理论保证与实际性能之间仍存在显著差距：

实际近似方法（如 Reformer、Performer、ScatterBrain）通常缺乏严格的误差保证，或者需要二次方运行时间才能确保高保真度。
理论进展（如 Alman & Song, 2023）表明，对于有界输入，实现近线性时间 ( $O(n^{1+o(1)})$ ) 且具有超多项式误差衰减是可能的，但现有的实际实现要么无法在近线性时间内达到这种衰减，要么需要二次方时间才能保证这一点。

本文指出，需要一种既具备计算效率（近线性时间）、又具备谱准确性（超多项式误差衰减），且能在现代硬件上高效实现的方案。

方法论

作者引入了 WILDCAT（基于加权迭代低秩分解的核集注意力，Weighted Iterative Low-rank Decomposition for Coreset ATtention），该方法通过仅对一小部分加权的输入键（keys）进行注意力计算来近似 Softmax 注意力机制。该方法由三个核心组件组成：

1. 通过随机枢轴 Cholesky (RPC) 进行加权核集选择

WILDCAT 不会对所有 $n$ 个键进行注意力计算，而是选择一个大小为 $r$ 的键子集 $K_S$ 来构成核集（coreset）。

算法： 该选择过程使用并行化的随机枢柄 Cholesky (RPC) 算法（改编自 Chen et al., 2022a）。该算法构建了核矩阵 $H = \exp(\beta K K^\top)$ 的部分 Cholesky 分解。
枢柄规则： 在每次迭代中，从当前残差核矩阵的对角线上采样一个枢柄索引。这确保了所选的点在频谱上能代表整个数据集。
效率： 选择过程的运行时间为 $O(nr^2)$ 。为了最大化 GPU 吞吐量，作者采用了分而治之的策略，将输入划分为 $B$ 个分箱并并行处理，从而将运行时间降低至 $O(nr^2/B^3 + nrd/B^2)$ 。

2. 通过 Nyström 近似进行最优重加权

一旦选定了核集 $K_S$ ，该方法会计算最优权重，以最小化注意力矩阵的重建误差。

Nyström 权重： 权重 $W$ 使用 Nyström 方法导出： $W = h(K_S, K_S)^+ h(K_S, K)$ 。这提供了最优的线性组合，利用核集特征来近似完整的核特征。
值压缩： 值矩阵 $V$ 被压缩为 $V_S = WV$ ，使得最终的注意力输出可以仅通过这 $r$ 个核集元素进行计算。

3. 不变性与重缩放

作者利用了 Softmax 注意力在键的全局重定中心（recentering）以及键与查询（queries）的重缩放下的不变性。

重定中心： 通过减去均值对键进行重定中心处理（ $K \leftarrow K - \bar{k}$ ）。
温度缩放： 推导出了一个闭式重缩放参数 $\tau$ ，用于平衡两个权衡：一是通过增加 $\tau$ 使核矩阵更易于低秩近似，二是最小化基于查询的误差膨胀因子。该参数取决于 Lambert-W 函数以及查询与键的范数。

核心贡献

本文为 WILDCAT 确立了以下理想特性：

近线性运行时间： WILDCAT 避免了精确注意力的二次方成本。通过对大小为 $r \in n^{o(1)}$ 的核集进行注意力计算，算法实现了近线性时间复杂度 $O(n^{1+o(1)})$ 。
超多项式误差衰减： 在有界输入假设下，WILDCAT 以 $O(n^{-\sqrt{\log(\log(n))}})$ 的超多项式误差衰减来近似精确注意力。这是在保持近线性运行时间的同时实现的，而此前的实际近似方法无法在不使用二次方时间的情况下保证这一点。
对增长的鲁棒性： 理论分析（推论 2）表明，即使输入项或维度随 $n$ 呈超对数增长，WILDCAT 仍能保持超多项式误差衰减，从而超越了以往关于注意力机制计算困难性的限制。
KV Cache 压缩： 该方法包含一个 COMPRESSKV 模块，可将 KV 缓存的内存占用从线性 $\Theta(nd)$ 降低到近常数级 $O(rd)$，从而促进长上下文推理。
实际落地实现： 作者提供了一个经过 GPU 优化的 PyTorch 实现，弥合了理论谱保证与现实性能之间的鸿沟。

实验结果

作者在三个领域对 WILDCAT 进行了评估：

图像生成 (BigGAN)： WILDCAT 相比精确注意力实现了 4.33 倍的加速。它表现出最小的 Inception Score 退化（1.21%）以及与 ImageNet 验证集相比 零退化 的 Fréchet Inception Distance (FID)，性能优于 Reformer、Performer、KDEformer 和 Thinformer。
图像分类 (T2T-ViT)： 在 ImageNet 分类基准测试中，WILDCAT 在所有近似方法中取得了最高的平均 Top-1 准确率（82.19%，精确注意力为 82.55%）。对于计算量占主导地位的第一层，它实现了 11.59 倍的加速。
长上下文语言模型 (KV Cache 压缩)： 使用 Qwen2.5-7B-Instruct 模型在 13 个 LongBench-E 任务上进行测试，WILDCAT（通过 COMPRESSKV）在不同压缩水平（75%、87.5%、93.75%）下，在平均压缩质量方面始终优于五种领先的 KV 缓存压缩方法（包括 StreamingLLM、PyramidKV 和 SnapKV）。
与 FlashAttention 2 的对比： 基准测试显示，随着序列长度增加（从 $2^{13}$ 到 $2^{18}$ ），WILDCAT 的加速比从 1.1 倍增长到 68 倍，同时近似误差稳步下降。

重要性与声明

本文声称成功弥合了注意力近似中的理论与实践差距。

理论突破： 它证明了对于实际的有界输入，可以在近线性时间内实现超多项式误差衰减，这一保证此前仅存在于理论层面或需要二次方运行时间。
实际效用： 该方法不仅是一个理论构造；它在 GPU 上得到了高效实现，并在生成、分类和语言理解的先进模型上得到了验证。
可扩展性： 通过将 KV 缓存的内存需求降低到近常数规模，WILDCAT 使在资源受限的硬件上部署长上下文语言模型成为可能。

作者指出了局限性，具体而言，目前的工作尚未解决带有因果掩码（causal masking）的流式生成问题（尽管他们建议未来可以通过前缀和或分而治之的方法进行扩展），并且 RPC 中枢柄选择的顺序性可以通过未来的过采样机制进行改进。