标题： MERIT：语言模型大批次训练的最大归一化逐元素比率 显示英文标题

标题： MERIT: Maximum-normalized Element-wise Ratio for Language Model Large-batch Training

摘要： 大规模批量训练已成为加速深度神经网络训练的核心方法，但其在优化和泛化方面带来了挑战。 现有的优化器如AdamW在语言模型的大规模批量训练中会出现性能下降，这是由于注意力层中的信息瓶颈导致最大注意力logit急剧增加。 虽然LAMB优化器部分解决了这个问题，但某些注意力层仍然面临这一问题。 原因是LAMB中的$l_2$-范数信任比例在直接影响查询/键权重的最大值方面效果不佳。 此外，LAMB中的逐权重信任比例容易出错，因为它忽略了行或列内权重值之间的关系。 基于这些观察，我们提出了一种新的优化器MERIT，它利用最大范数来计算信任比例，以更有效地约束最大注意力logit。 此外，我们进一步构建了逐元素的信任比例，通过关注局部权重结构来提供更稳健的更新缩放。 在各种大小的GPT-2模型上的大规模批量训练的广泛实验表明了MERIT的优越性能。 值得注意的是，在GPT-2 Medium的训练中，MERIT使批量大小达到6k时，与标准批量大小（480）相比，使用48B训练标记时没有性能下降。 这项工作强调了在大规模批量训练中考虑最大注意力logit和更细粒度信任比例的重要性。 它成功地提高了训练稳定性，并为使用更大的批量铺平了道路，从而加快了大型语言模型的开发和迭代。 代码可在https://github.com/NUS-HPC-AI-Lab/MERIT获得。 显示英文摘要

摘要： Large-batch training has become a cornerstone in accelerating the training of deep neural networks, yet it poses challenges in optimization and generalization. Existing optimizers like AdamW present performance degradation during language models' large-batch training, due to the information bottleneck in attention layers caused by the sharp increase of max attention logit. While the LAMB optimizer partially addresses this issue, some attention layers still face this issue. The reason is that $l_2$-norm-based trust ratios in LAMB are less effective in directly influencing the max value of query/key weights. Furthermore, the weight-wise trust ratio in LAMB is error-prone as it overlooks relationships of weight values within rows or columns. Building on these observations, we propose a novel optimizer, MERIT, which leverages the max-norm to calculate the trust ratio to constrain the max attention logit more effectively. Moreover, we further construct element-wise trust ratios to provide more robust update scaling by focusing on local weight structures. Extensive experiments of large-batch training across various sizes of GPT-2 models demonstrate the superior performance of MERIT. Notably, during the training of GPT-2 Medium, MERIT enables a 6k batch size without any performance degradation compared to the standard batch size (480) with 48B training tokens. This work highlights the importance of considering the max attention logit and finer-granularity trust ratio in large-batch training. It successfully improves the training stability and paves the way for larger batch usage, enabling faster development and iteration of large language models. Code is available at https://github.com/NUS-HPC-AI-Lab/MERIT.