Broken Symmetry in BF16 Attention: Why FlashAttention Gradients Blow Up Late in Training
GProj (gauge projection), which restores the zero sum after the cast with two rank-one corrections per row, and cuts the remaining median query/key gradient errors from 219%/13% to 0.34%/0.37%, on par with FP32 attention, for 4.7% more time per training step.
Jun-Lin Chen, Daize Dong, Huan-Wei Di et al.
· 0 citations