低精度Transformer推理中的随机舍入:小型 GPT-2 的可变精度仿真研究
Stochastic Rounding in Low-Precision Transformer Inference: A Variable-Precision Emulation Study of a Small GPT-2
摘要
低精度Transformer推理应该使用随机舍入 (SR) 还是舍入到最近值 (RN)?答案取决于您在网络中查找的位置。我们通过固定数字格式并仅改变各个操作站点的舍入规则来隔离这种影响。为了能够以自由选择的精度进行实验,我们通过可变精度随机舍入 (VPSR) 算法将 PRISM 矢量化舍入库扩展到任意虚拟精度,证明舍入决策是在硬件浮点中精确评估的。我们开发了两项分析,为这种站点级别的权衡提供了补充见解。首先,线性投影的概率前向误差界限表明,SR 的误差包络在缩减长度 $n$ 中随着 $O(\sqrt{n} u)$ 增长,而 RN 的误差包络为 $O(n u)$,这一差距在低精度下迅速扩大,并且在长多层感知器 (MLP) 下投影中最为明显。其次,将输出 softmax 处的预期交叉熵损失变化二阶分解为有符号漂移、漂移曲率和 Fisher 加权方差罚分,揭示了为什么两个站点表现相反:MLP 噪声主要是均匀 logit 偏移,而 softmax 是不变的,因此 SR 的方差很大程度上被折扣;头部噪音在整个词汇表中是不一致的,事实并非如此。在 $t=6$ 有效位的 DistilGPT-2 上,观察结果与理论相符:MLP 中的 SR 将复杂度提高到全精度参考的 1.15 倍,而 RN 的复杂度为 2.21 倍。在语言模型头部,顺序相反,因为 SR 引入了非均匀方差,而确定性 RN 则不引入任何方差。在混合精度配置($t=6$ 的 MLP 输出)中,将 SR 分配给 MLP,将 RN 分配给头部,带来的困惑度在全精度参考的 1.10 倍以内,比匹配位 RN 减少了 28%。