动态分配注意力码本容量，提升长序列模型效率与精度。
AVQ-Attention: Adaptive Vector-Quantized Attention
- 根据注意力重要性动态分配码本容量，关键区域精细量化。
- 在保持O(MN)复杂度下，相比固定码本提升准确率与效率平衡。
- 基于Triton实现端到端自适应更新，兼容Flash Attention计算范式。
Transformer模型中注意力机制的复杂度为O(N²)，成为计算瓶颈。向量量化（VQ）注意力通过用M个码字表示键值，将复杂度降低至O(MN)，但其码本容量固定，导致高注意力区域量化粗糙、低注意力区域资源浪费。本文提出自适应向量量化（AVQ）注意力，基于注意力重要性动态分配码本容量。从少量码字出发，前向传播中识别关键码字，并利用预训练子码字进行精细化重构，在重要区域实现细粒度量化，其他区域保持粗粒度。我们开发了基于自定义Triton内核的实现，使重要性评分、子码字插入和父码贡献替换等全过程可在Flash Attention的分块计算框架中高效执行，开销极小。该方法在维持O(MN)复杂度的同时，显著优于固定码本的VQ注意力，在准确率-效率权衡上表现更优。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。