这篇论文提出了一个聪明的点子：让注意力机制像人一样，在关键地方用更多计算资源，其他地方粗略处理。复杂度降低但精度不降，做推理优化的同学可以关注。
AVQ-Attention 通过自适应向量量化将注意力复杂度从 O(N²) 降至 O(MN)，其中 N 为 token 数，M 为码字数量。与固定码本的 VQ-Attention 不同，它根据注意力重要性动态分配码本容量，在高注意区域使用预学习的子码字进行细粒度量化，低注意区域维持粗粒度。方法基于自定义 Triton 内核，可在 Flash Attention 的平铺计算中完成重要性评分、子码字插入等操作，实现最小开销。实验表明 AVQ-Attention 在保持 O(MN) 复杂度的同时，相比固定码本 VQ-Attention 实现了更好的精度-效率权衡。
AVQ-Attention: Adaptive Vector-Quantized Attention
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.