AVQ-Attention: Adaptive Vector-Quantized Attention
arXiv:2607.12789
The paper introduces Adaptive Vector-Quantized (AVQ) Attention, which dynamically allocates codebook capacity to the most important regions of the key space, preserving O(MN) complexity while improving the accuracy-efficiency trade‑off of transformer attention.
Abstract
The complexity of attention over tokens remains a computational bottleneck in transformer models. Vector-Quantized (VQ) attention reduces this to by representing keys with codewords, but applies uniform codebook capacity regardless of where attention mass concentrates: high-attention regions of key space may be coarsely approximated while low-attention regions waste representational capacity. We propose Adaptive Vector-Quantized (AVQ) Attention, which adaptively allocates codebook capacity based on attention importance. Starting from a small set of codewords, our method identifies the most important codes during the forward pass and refines them with pre-learned child codewords, achieving fine-grained quantization where it matters most while maintaining coarse quantization elsewhere. We develop an implementation using custom Triton kernels that enables the full adaptive refinement process, including importance scoring, child codeword insertion, and parent contribution replacement, to be carried out within the tiled computation paradigm of Flash Attention with minimal overhead. Our approach maintains complexity while achieving improved accuracy-efficiency trade-offs compared to fixed-codebook VQ-attention.
Accepted at ECCV 2026