MERIT：用于语言模型大批量训练的最大规范元素比

大批量训练已成为加速深层神经网络训练的基石，但在优化和泛化方面仍面临挑战。现有优化器如 AdamW 在语言模型的大批量训练中表现下降，原因是注意力层中信息瓶颈的加剧，由 max 注意力 logit 的显著增加所致。虽然 LAMB 优化器部分缓解了此问题，但仍有一些注意力层面临此问题。其原因是 $l_2$ 规范基于的信任比例在直接影响查询/键权重的最大值方面效果较差。此外，LAMB 中的权重级信任比例忽略了权重值在行或列中的关系，导致误差。基于上述观察，我们提出了一种新型优化器 MERIT，利用最大规范计算信任比例，以更有效地约束 max 注意力 logit。此外，我们进一步构建元素级信任比例，通过关注局部权重结构来提供更稳健的更新缩放。广泛的 GPT-2 模型大批量训练实验表明，MERIT 性能卓越。值得注意的是，在 GPT-2 中等规模模型的训练中，MERIT 使 6000 批量大小的训练在不进行性能下降的情况下，相当于标准批量大小（480）下的 4800 亿训练标记。本 work 凸显了在大批量训练中考虑 max 注意力 logit 以及更细粒度信任比的重要性。它成功地提高了训练稳定性，为更大批量的使用铺平了道路，加速了大型语言模型的快速开发与迭代。代码已公开于 https://github.com/NUS-HPC-AI-Lab/MERIT。