WildCat让注意力机制近似线性加速，精度却远超以往。
WildCat: Near-Linear Attention in Theory and Practice
- 用快速谱准确采样选小核心集，再优化加权降低误差。
- 在输入长度n下，误差以超多项式速度衰减，耗时接近线性。
- 适合图像生成、分类和语言模型缓存压缩等场景。
我们提出WildCat，一种高精度、低成本的神经网络注意力机制压缩方法。尽管注意力是现代网络架构的核心组件，但其资源消耗随输入序列长度n呈二次增长，部署成本极高。WildCat通过仅对一个加权小核心集进行注意力计算，避免了二次开销。关键在于，我们采用快速且谱准确的采样算法——随机轮换的乔列斯基分解，来选择核心集，并对元素进行最优加权以最小化重构误差。令人惊讶的是，在输入有界条件下，WildCat以超多项式速度 $O(n^{- oot{ ext{log(log(n))}}{}})$ 逼近精确注意力，同时运行时间仅为近线性 $O(n^{1+o(1)})$。相比之下，以往实用方法要么缺乏误差保证，要么需二次时间才能达到如此高保真度。我们还提供了针对GPU优化的PyTorch实现，并通过一系列基准实验验证了WildCat在图像生成、图像分类及语言模型键值缓存压缩中的优势。
原文摘要 · Abstract (English)
We introduce WildCat, a high-accuracy, low-cost approach to compressing the attention mechanism in neural networks. While attention is a staple of modern network architectures, it is also notoriously expensive to deploy due to resource requirements that scale quadratically with the input sequence length $n$. WildCat avoids these quadratic costs by only attending over a small weighted coreset. Crucially, we select the coreset using a fast but spectrally-accurate subsampling algorithm -- randomly pivoted Cholesky -- and weight the elements optimally to minimise reconstruction error. Remarkably, given bounded inputs, WildCat approximates exact attention with super-polynomial $O(n^{-\sqrt{\log(\log(n))}})$ error decay while running in near-linear $O(n^{1+o(1)})$ time. In contrast, prior practical approximations either lack error guarantees or require quadratic runtime to guarantee such high fidelity. We couple this advance with a GPU-optimized PyTorch implementation and a suite of benchmark experiments demonstrating the benefits of WildCat for image generation, image classification, and language model KV cache compression.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。