将注意力计算压缩到1比特,速度超闪光注意力2倍以上
BinaryAttention: One-Bit QK-Attention for Vision and Diffusion Transformers
- 仅保留查询与键的符号位,用位运算替代浮点乘法
- 在视觉与扩散模型上达到甚至超越全精度性能
- 适合追求极致推理速度的部署场景
Transformer在视觉任务中取得显著成功,但其注意力模块的计算复杂度仍是主要瓶颈。现有方法多采用8比特或4比特量化以平衡效率与精度。本文基于理论分析指出,注意力的二值化可保留关键相似性关系,提出BinaryAttention,一种高效且准确的1比特qk-attention方法。具体而言,仅保留查询和键的符号位,并用位运算替代浮点点积,大幅降低计算开销。通过引入可学习偏置缓解1比特量化带来的信息损失,并实现端到端加速。为保持注意力精度,采用量化感知训练与自蒸馏技术,有效抑制量化误差并保证符号对齐的相似性。在A100 GPU上,BinaryAttention比FlashAttention2快2倍以上。大量实验表明,在视觉与扩散Transformer基准上,BinaryAttention性能与全精度相当甚至更优。本工作为低比特视觉与扩散Transformer提供了高效可靠的替代方案。代码与模型见https://github.com/EdwardChasel/BinaryAttention。
原文摘要 · Abstract (English)
Transformers have achieved widespread and remarkable success, while the computational complexity of their attention modules remains a major bottleneck for vision tasks. Existing methods mainly employ 8-bit or 4-bit quantization to balance efficiency and accuracy. In this paper, with theoretical justification, we indicate that binarization of attention preserves the essential similarity relationships, and propose BinaryAttention, an effective method for fast and accurate 1-bit qk-attention. Specifically, we retain only the sign of queries and keys in computing the attention, and replace the floating dot products with bit-wise operations, significantly reducing the computational cost. We mitigate the inherent information loss under 1-bit quantization by incorporating a learnable bias, and enable end-to-end acceleration. To maintain the accuracy of attention, we adopt quantization-aware training and self-distillation techniques, mitigating quantization errors while ensuring sign-aligned similarity. BinaryAttention is more than 2x faster than FlashAttention2 on A100 GPUs. Extensive experiments on vision transformer and diffusion transformer benchmarks demonstrate that BinaryAttention matches or even exceeds full-precision attention, validating its effectiveness. Our work provides a highly efficient and effective alternative to full-precision attention, pushing the frontier of low-bit vision and diffusion transformers. The codes and models can be found at https://github.com/EdwardChasel/BinaryAttention.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。