WildCat：理论与实践中的近线性注意力机制

我们提出 WildCat，这是一种高精度、低成本的压缩神经网络注意力机制的方法。虽然注意力是现代网络架构的基本组件，但由于注意力机制随输入序列长度为 $n$ 而呈二次增长的资源需求，部署时却非常昂贵。WildCat 通过仅注意力于小型加权核心组（coreset）来避免这些二次成本。关键在于，我们使用一种快速且在谱上精确的抽样算法——随机枢轴柯尔西 Cholesky 方法——选择核心组，并对元素进行最优加权以最小化重构误差。惊人的是，在给定有界输入的情况下，WildCat 对精确注意力的近似误差以超多项式 $O(n^{-\sqrt{\log(\log(n))}})$ 衰减，而运行时间为近线性 $O(n^{1+o(1)})$。相比之下，先前的实际近似方法要么缺乏误差保证，要么需要二次运行时间才能保证如此高的保真度。我们结合这一突破，配合 GPU 优化的 PyTorch 实现，并进行一系列基准实验，展示 WildCat 在图像生成、图像分类和语言模型 KV 缓存压缩方面的优势。