[论文解读] Small-scale proxies for large-scale Transformer training instabilities

Mitchell Wortsman, Peter J. Liu+14 more|arXiv (Cornell University)|Sep 25, 2023

Explainable Artificial Intelligence (XAI)被引用 4

一句话总结

本文通过小规模代理实验研究大规模Transformer模型中的训练不稳定性，证明关键不稳定性——如注意力层中的logit增长和对数概率的发散——可在小模型中复现并借助qk-layernorm和z-loss等技术得到缓解。作者表明学习率敏感性是一种可靠的诊断指标，且激活和梯度范数的缩放行为可预测不稳定性在出现前的征兆。

ABSTRACT

Teams that have trained large Transformer-based models have reported training instabilities at large scale that did not appear when training with the same hyperparameters at smaller scales. Although the causes of such instabilities are of scientific interest, the amount of resources required to reproduce them has made investigation difficult. In this work, we seek ways to reproduce and study training stability and instability at smaller scales. First, we focus on two sources of training instability described in previous work: the growth of logits in attention layers (Dehghani et al., 2023) and divergence of the output logits from the log probabilities (Chowdhery et al., 2022). By measuring the relationship between learning rate and loss across scales, we show that these instabilities also appear in small models when training at high learning rates, and that mitigations previously employed at large scales are equally effective in this regime. This prompts us to investigate the extent to which other known optimizer and model interventions influence the sensitivity of the final loss to changes in the learning rate. To this end, we study methods such as warm-up, weight decay, and the $μ$Param (Yang et al., 2022), and combine techniques to train small models that achieve similar losses across orders of magnitude of learning rate variation. Finally, to conclude our exploration we study two cases where instabilities can be predicted before they emerge by examining the scaling behavior of model activation and gradient norms.

研究动机与目标

• 研究在相同超参数下于小规模下不显现但在大规模Transformer中出现的训练不稳定性。
• 开发小规模代理模型，在无需大量计算资源的情况下研究这些不稳定性。
• 识别并验证有效干预措施（如qk-layernorm和z-loss）在学习率变化范围内稳定训练的效果。
• 分析优化器和模型修改（如学习率预热、权重衰减、μ-Param）对小模型中学习率敏感性的影响。
• 探索激活和梯度范数的缩放趋势是否可在训练中不稳定出现前进行预测。

提出的方法

• 使用Flax和AdamW优化器以GPT-2风格架构训练小规模仅解码器Transformer模型，采用默认超参数。
• 在三个数量级范围内系统性地改变学习率，并测量最终评估损失，将学习率（LR）敏感性作为汇总统计量。
• 引入并评估qk-layernorm和z-loss正则化，以降低学习率敏感性，并在宽泛的学习率范围内稳定训练。
• 研究预热时长、权重衰减和μ-Param对小模型中学习率敏感性和训练稳定性的影响。
• 分析模型规模和学习率下激活范数与梯度范数的缩放行为，以检测不稳定的早期迹象。
• 利用观察结果：默认AdamW的epsilon值过大，导致更新过小，从而与不稳定性机制建立关联。

实验结果

研究问题

• RQ1 在高学习率下训练时，大规模Transformer中观察到的训练不稳定性（如注意力logit增长和输出logit发散）是否可在小规模模型中复现？
• RQ2 已知的大规模缓解技术（如qk-layernorm、z-loss）在小规模代理设置中是否仍有效？
• RQ3 标准优化器和模型干预措施（如预热、权重衰减、μ-Param）在小模型中如何影响学习率敏感性？
• RQ4 激活和梯度范数的缩放趋势是否可在不稳定性显现前进行预测？
• RQ5 AdamW的epsilon超参数在不稳定性中扮演何种角色，特别是与梯度范数缩放的关系？

主要发现

• 在高学习率下训练时，小模型中可复现注意力logit增长和输出logit发散等训练不稳定性。
• qk-layernorm和z-loss正则化可降低学习率敏感性，使小模型在三个数量级的学习率变化范围内实现稳定训练。
• 随着模型深度的增加，学习率敏感性比宽度增加得更快，表明深层模型对学习率选择更敏感。
• 梯度范数随模型规模和学习率的增加而减小，表明默认的AdamW epsilon（1e-8）过大，导致参数更新过小。
• 渐进式锐化与自适应优化器动态（如使用β₂的Adam）之间的相互作用解释了周期性损失尖峰现象，减小β₂可通过加快预条件器的适应速度来防止此类现象。
• 激活和梯度范数的缩放行为可在不稳定性出现前进行预测，为训练稳定性提供无需大规模训练的诊断工具。