AVQ-Attention: Adaptive Vector-Quantized Attention

2026-07-14

AI Summary

This work addresses the O(N²) computational bottleneck of Transformer attention by proposing an adaptive vector quantized attention mechanism. Unlike conventional vector quantization methods that rely on fixed codebooks—leading to coarse approximations in high-attention regions and redundant representations in low-attention areas—our approach dynamically identifies salient regions during the forward pass starting from a small set of initial codewords. It refines these critical regions with pre-learned sub-codewords while maintaining coarse representations elsewhere. This is the first method to enable attention-aware, dynamic codebook allocation, overcoming the limitations of static structures. By integrating importance scoring, sub-codeword insertion, and parent codeword replacement via a custom Triton kernel, our design seamlessly aligns with the FlashAttention tiling paradigm, achieving a significantly improved trade-off between accuracy and efficiency while preserving O(MN) complexity.

Abstract

The $\mathcal{O}(N^2)$ complexity of attention over $N$ tokens remains a computational bottleneck in transformer models. Vector-Quantized (VQ) attention reduces this to $\mathcal{O}(MN)$ by representing keys with $M$ 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 $\mathcal{O}(MN)$ complexity while achieving improved accuracy-efficiency trade-offs compared to fixed-codebook VQ-attention.

Problem

• attention complexity
• vector quantization
• codebook allocation
• transformer efficiency
• adaptive representation

Innovation

• Adaptive Vector Quantization
• Efficient Attention
• Codebook Allocation
• Flash Attention
• Triton Kernels