SageBwd: A Trainable Low-bit Attention
Fuente:
arXiv
Gespeichert in:
| Hauptverfasser: | , , , , , , , |
|---|---|
| Format: | Preprint |
| Veröffentlicht: |
2026
|
| Schlagworte: | |
| Online-Zugang: | |
| Tags: |
Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
|
| _version_ | 1866908860525051904 |
|---|---|
| author | Zhang, Jintao Chen, Marco Wang, Haoxu Jiang, Kai Stoica, Ion Gonzalez, Joseph E. Chen, Jianfei Zhu, Jun |
| author_facet | Zhang, Jintao Chen, Marco Wang, Haoxu Jiang, Kai Stoica, Ion Gonzalez, Joseph E. Chen, Jianfei Zhu, Jun |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2603_02170 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | SageBwd: A Trainable Low-bit Attention Zhang, Jintao Chen, Marco Wang, Haoxu Jiang, Kai Stoica, Ion Gonzalez, Joseph E. Chen, Jianfei Zhu, Jun Machine Learning Artificial Intelligence 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. |
| title | SageBwd: A Trainable Low-bit Attention |
| topic | Machine Learning Artificial Intelligence |
| url | https://arxiv.org/abs/2603.02170 |