SageBwd: Обучаемое низкобитное внимание

SageBwd: A Trainable Low-bit Attention

March 2, 2026

Авторы: Jintao Zhang, Marco Chen, Haoxu Wang, Kai Jiang, Ion Stoica, Joseph E. Gonzalez, Jianfei Chen, Jun Zhu

cs.AI

Аннотация

Низкобитная квантизация внимания, такая как SageAttention, стала эффективным подходом для ускорения вывода моделей, однако её применимость к обучению остаётся малоизученной. В предыдущей работе мы представили SageBwd — обучаемую квантизацию внимания в формате INT8, которая квантизирует шесть из семи матричных умножений матрицы внимания, сохраняя производительность при дообучении. Однако SageBwd демонстрировал устойчивый разрыв в производительности по сравнению с вниманием в полной точности (FPA) на этапе предварительного обучения. В данной работе мы исследуем причины этого разрыва и показываем, что SageBwd достигает производительности полной точности при предварительном обучении. Благодаря экспериментам и теоретическому анализу мы получили ряд важных инсайтов и выводов: (i) нормализация QK необходима для стабильности обучения при большом количестве токенов на шаг, (ii) ошибки квантизации в основном возникают из-за градиента оценок dS при обратном проходе, (iii) уменьшение количества токенов на шаг позволяет SageBwd достичь производительности FPA при предварительном обучении, и (iv) сглаживание K остаётся критически важным для стабильности обучения, тогда как сглаживание Q даёт ограниченный выигрыш на этапе предварительного обучения.

English

Low-bit attention, such as SageAttention, has emerged as an effective approach for accelerating model inference, but its applicability to training remains poorly understood. In prior work, we introduced SageBwd, a trainable INT8 attention that quantizes six of seven attention matrix multiplications while preserving fine-tuning performance. However, SageBwd exhibited a persistent performance gap to full-precision attention (FPA) during pre-training. In this work, we investigate why this gap occurs and demonstrate that SageBwd matches full-precision attention during pretraining. Through experiments and theoretical analysis, we reach a few important insights and conclusions: (i) QK-norm is necessary for stable training at large tokens per step, (ii) quantization errors primarily arise from the backward-pass score gradient dS, (iii) reducing tokens per step enables SageBwd to match FPA performance in pre-training, and (iv) K-smoothing remains essential for training stability, while Q-smoothing provides limited benefit during pre-training.