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

📄 arXiv: 2508.20577v1 📥 PDF

作者: Yang Luo, Zangwei Zheng, Ziheng Qin, Zirui Zhu, Yong Liu, Yang You

分类: cs.LG, cs.AI, cs.CL

发布日期: 2025-08-28

备注: ICML 2025

🔗 代码/项目: GITHUB

💡 一句话要点

提出MERIT优化器，通过最大范数归一化元素级比例提升语言模型大批量训练性能。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大批量训练 语言模型 优化器 最大范数 注意力机制

📋 核心要点

  1. 现有优化器如AdamW在大批量训练语言模型时性能下降，主要原因是注意力层中最大注意力logit的急剧增加导致信息瓶颈。
  2. MERIT优化器通过最大范数计算信任比例，更有效地约束最大注意力logit，并构建元素级信任比例，关注局部权重结构，提供更稳健的更新缩放。
  3. 在GPT-2模型上的实验表明，MERIT能够显著提升大批量训练性能，例如在GPT-2 Medium上使用6k批量大小训练，性能与标准批量大小相当。

📝 摘要（中文）

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

🔬 方法详解

问题定义：论文旨在解决大批量训练语言模型时，现有优化器（如AdamW）性能下降的问题。现有方法，特别是基于l2范数的LAMB优化器，在约束注意力层中最大注意力logit方面效果不佳，且权重级信任比例忽略了权重值之间的关系，导致训练不稳定。

核心思路：论文的核心思路是通过最大范数来计算信任比例，从而更有效地约束最大注意力logit。同时，构建元素级信任比例，关注局部权重结构，提供更稳健的更新缩放。这种方法旨在解决现有优化器在大批量训练中遇到的信息瓶颈和权重更新不准确的问题。

技术框架：MERIT优化器的整体框架是在现有的优化器基础上进行改进，主要集中在信任比例的计算方式上。它没有改变整体的训练流程，而是通过更精细的权重更新策略来提高训练的稳定性和效率。主要模块包括：梯度计算、最大范数归一化、元素级信任比例计算、权重更新。

关键创新：MERIT优化器的关键创新在于：1) 使用最大范数来计算信任比例，更直接地约束最大注意力logit；2) 构建元素级信任比例，考虑了局部权重结构，提供了更稳健的更新缩放。这与LAMB优化器中基于l2范数的权重级信任比例有本质区别，后者忽略了权重值之间的关系。

关键设计：MERIT优化器的关键设计包括：1) 最大范数归一化的具体实现方式，如何选择合适的范数类型和归一化策略；2) 元素级信任比例的计算方法，如何定义局部权重结构，以及如何将信任比例应用到权重更新中；3) 与现有优化器（如AdamW）的结合方式，如何调整学习率、动量等超参数以获得最佳性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，MERIT优化器在GPT-2模型的训练中表现出优越的性能。例如，在GPT-2 Medium的训练过程中，MERIT能够使用6k的批量大小，而不会出现任何性能下降，与使用标准批量大小（480）和48B训练tokens相比，性能相当。这表明MERIT能够有效提高大批量训练的稳定性和效率。

🎯 应用场景

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.