WildCat: Near-Linear Attention in Theory and Practice

We introduce WildCat, a high-accuracy, low-cost approach to compressing the attention mechanism in neural networks. The attention mechanism is 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. WildCat runs in near-linear $O(n^{1+o(1)})$ time while guaranteeing super-polynomial $O(n^{-\sqrt{\log(\log(n))}})$ error decay compared to exact attention. This property enables adapting the runtime parameter at the time of deployment depending on resource constraints and quality requirements. 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.