AVQ-Attention: Adaptive Vector-Quantized Attention

Standard Vector-Quantized (VQ) attention reduces the quadratic complexity of transformers to linear O(MN) by clustering keys into a fixed set of M codewords. However, this approach applies uniform precision across the entire key space, often wasting capacity on low-attention regions while providing insufficient resolution for high-attention regions. The authors ask: can we dynamically allocate codebook capacity based on the attention importance of specific regions to improve the accuracy-efficiency trade-off?

To address this, the authors introduce Adaptive Vector-Quantized (AVQ) Attention. The method utilizes a hierarchical codebook structure where each parent codeword is associated with a set of pre-learned child codewords. During the forward pass, the model computes attention over parent codewords and identifies the most important ones based on the attention mass they receive. These top-P important parents are then refined by their corresponding children.

To ensure this process remains efficient, the authors developed custom Triton kernels that integrate into the Flash Attention tiled computation paradigm. They enforce a geometric constraint where each parent is the mean of its children, which allows the model to perform in-register corrections during the attention computation without needing to revisit global memory. This enables the adaptive refinement to occur with minimal overhead.

AVQ-Attention provides a way to maintain the linear complexity benefits of VQ-attention while significantly improving the quality of the approximation. By concentrating computational resources where they are most needed—specifically in regions of the key space that receive high attention—the model achieves better performance on tasks like image classification and generation compared to fixed-codebook methods. This approach is particularly useful for long-sequence processing where the O(N^2) bottleneck of standard attention is prohibitive.