AVQ-Attention: Adaptive Vector-Quantized Attention

The O(N 2 ) complexity of attention over N tokens remainsa computational bottleneck in transformer models. Vector-Quantized(VQ) attention reduces this to O(M N ) by representing keys with Mcodewords, but applies uniform codebook capacity regardless of whereattention mass concentrates: high-attention regions of key space maybe coarsely approximated while low-attention regions waste representa-tional capacity. We propose Adaptive Vector-Quantized (AVQ) Atten-tion, which adaptively allocates codebook capacity based on attentionimportance. Starting from a small set of codewords, our method identifiesthe most important codes during the forward pass and refines them withpre-learned child codewords, achieving fine-grained quantization whereit matters most while maintaining coarse quantization elsewhere. We de-velop an implementation using custom Triton kernels that enables the fulladaptive refinement process, including importance scoring, child code-word insertion, and parent contribution replacement, to be carried outwithin the tiled computation paradigm of Flash Attention with minimaloverhead. Our approach maintains O(M N ) complexity while achievingimproved accuracy-efficiency trade-offs compared to fixed-codebook VQ-attention.