何处、何时以及多少比特?采用随机舍入的 SGD 的缩放定律
Where, When, and How Many Bits? Scaling Laws for SGD with Stochastic Rounding
精选理由
为低精度训练提供精度、模型规模与训练时长的权衡规律,助力在算力预算内优化预测效果。
AI 摘要
arXiv:2610.41074v1 公告类型:新论文 摘要:低精度训练可减少计算量,但会引入额外的舍入误差。我们从理论上研究了随着模型规模、训练时间和精度的变化,随机舍入(SR)如何影响预测风险。我们考虑一个第一层固定且随机的双层线性网络,使用具有幂律谱的新采样高斯数据,通过小批量 SGD 进行训练。在每个 SGD 步骤中,前向和反向计算所用的特征分别独立进行舍入。对于训练前选定的一般学习率与精度调度方案,我们证明了最终模型和均匀参数平均模型的预测误差具有相匹配的上下界。在精度固定、学习率恒定且采用参数平均的情况下,我们推导了最小预测误差随计算量增加而下降的规律,并确定了实现这一计算最优缩放规律的模型规模和训练时间。仅对前向特征进行舍入会保留精确算术下的计算量缩放规律,而对反向特征进行舍入则会改变这一规律。参数平均可能改变哪些计算需要更高精度,以及更高精度在何时有用。我们的结果还给出了满足给定准确度容差所需的比特数,并揭示了反向舍入如何改变学习率调度方案的设计。在固定计算预算下,如果增大模型或延长训练带来的收益超过 SR 引入的额外误差,降低精度就能改善预测表现。模型仿真展示了这些权衡,基于 MLP 和小型 transformer 的实验则考察了理论模型之外的舍入选择。