提出几何感知缩放方法，解决低精度训练中注意力溢出问题。
Rank-Aware Spectral Bounds on Attention Logits for Stable Low-Precision Training
- 基于秩感知浓度不等式，推导出更紧的注意力数值边界。
- 在FP8训练中实现零溢出，且下游任务准确率与传统方法相当。
- 适用于大模型，兼容融合注意力核，无需观测激活值。
Transformer中的注意力分数是双线性形式 $S_{ij} = x_i^ op M x_j / \ oot d_h \of$，其最大幅值决定低精度训练中的溢出风险。我们推导出一个 extit{秩感知浓度不等式}：当交互矩阵 $M = W^Q W^{K op}$ 的秩 $r \ll d$ 时，$\\max_{i,j}|S_{ij}|$ 的尾概率以 $\\exp(-d^{2}α^{2}/(γr))$ 衰减，而非传统的 $\\exp(-dα^{2})$，其中 $γ>1$ 为典型性参数。对于秩为 $d_h$ 的Transformer注意力，该结果比无秩假设的界紧密8–28倍。我们将此应用于FP8训练，推导出 extit{几何感知缩放因子}，可在不观测激活值的情况下提供严格的溢出保障。该方法通过隐式幂迭代计算每层缩放因子，使用分组查询注意力避免键扩展，并兼容融合注意力核。在GPT-2 XL至Llama-2-70B上，该方法在瞬态场景中消除溢出，而下游MMLU准确率与延迟缩放方法相当。
原文摘要 · Abstract (English)
Attention scores in transformers are bilinear forms $S_{ij} = x_i^\top M x_j / \sqrt{d_h}$ whose maximum magnitude governs overflow risk in low-precision training. We derive a \emph{rank-aware concentration inequality}: when the interaction matrix $M = W^Q W^{K\top}$ has rank $r \ll d$, tail probabilities for $\max_{i,j}|S_{ij}|$ decay as $\exp(-d^{2}α^{2}/(γr))$ rather than $\exp(-dα^{2})$, where $γ> 1$ is a typicality parameter. For transformer attention where $r = d_h$, this yields $8$--$28\times$ tighter concentration than rank-agnostic bounds in modern architectures. We apply this result to FP8 training, deriving \emph{geometry-aware scale factors} that provide principled overflow guarantees without observing activations. The method computes per-layer scales from the spectral norm $\|W^Q W^{K\top}\|_2$ via implicit power iteration, includes a grouped query attention formulation that avoids key expansion, and remains compatible with fused attention kernels. Across GPT-2 XL to Llama-2-70B, geometry-aware scaling eliminates overflows in transient scenarios where delayed scaling fails, while achieving comparable downstream MMLU accuracy.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。