Rank-Aware Spectral Bounds on Attention Logits for Stable Low-Precision Training

变压器注意力得分是双线性形式$S_{ij} = x_i^\top M x_j / \sqrt{d_h}$，其最大模长支配着低精度训练中的溢出风险。我们推导了一种\emph{基于秩的浓度不等式}：当交互矩阵$M = W^Q W^{K\top}$的秩为$r \ll d$时，$\max_{i,j}|S_{ij}|$的尾概率衰减为$\exp(-d^{2}\alpha^{2}/(\gamma r))$，而非$\exp(-d\alpha^{2})$，其中$\gamma > 1$为典型性参数。对于变压器注意力而言，$r = d_h$，这导致比现代体系结构中基于秩无关的界限更紧约10-28倍。我们将此结果应用于FP8训练，推导了\emph{几何感知比例因子}，其提供了无需观察激活值的原则性溢出保证。该方法通过隐式幂迭代计算从谱范数$\|W^Q W^{K\top}\|_2$得到每层比例因子，包含一种避免键扩展的分组查询注意力公式，同时与融合注意力内核兼容。我们在GPT-2 XL至Llama-2-70B范围内的模型上测试，几何感知缩放消除了延迟缩放在瞬态场景中的溢出，同时在可比下游MMLU准确率方面取得相当的结果。