SageBwd：一种可训练的低比特注意力机制
机器学习
2026-03-03 v1 人工智能
摘要
低比特注意力机制，例如SageAttention，已成为加速模型推理的有效方法，但其在训练中的适用性仍知之甚少。在先前的工作中，我们引入了SageBwd，一种可训练的INT8注意力机制，它对七个注意力矩阵乘法中的六个进行量化，同时保持微调性能。然而，SageBwd在预训练期间表现出与全精度注意力持续存在的性能差距。在这项工作中，我们研究了这种差距产生的原因，并证明了SageBwd在预训练期间能够匹配全精度注意力。通过实验和理论分析，我们得出了一些重要的见解和结论：(i) QK归一化对于在每步大token数下稳定训练是必要的，(ii) 量化误差主要来自反向传播的分数梯度dS，(iii) 减少每步token数使SageBwd能够在预训练中匹配FPA性能，以及(iv) K平滑对于训练稳定性仍然至关重要，而Q平滑在预训练期间提供的益处有限。
引用
@article{arxiv.2603.02170,
title = {SageBwd: A Trainable Low-bit Attention},
author = {Jintao Zhang and Marco Chen and Haoxu Wang and Kai Jiang and Ion Stoica and Joseph E. Gonzalez and Jianfei Chen and Jun Zhu},
journal= {arXiv preprint arXiv:2603.02170},
year = {2026}
}