TLDR: MERIT is a new optimizer designed for large-batch training of language models, addressing performance degradation caused by increasing max attention logits. It uses a max-norm-based trust ratio to effectively constrain attention logits and introduces element-wise trust ratios for more robust update scaling by focusing on local weight structures. Experiments on GPT-2 and Llama models show MERIT outperforms existing optimizers like AdamW and LAMB, enabling larger batch sizes without performance loss and achieving better generalization with minimal computational overhead.
Large language models have transformed natural language processing, but their increasing size and complexity demand significant training time. Large-batch training offers a solution by processing more data in parallel, dramatically reducing training duration. However, this approach often introduces challenges in optimization and can lead to performance degradation, particularly with widely used optimizers like AdamW.
A key issue identified in large-batch training of language models is the sharp increase in the ‘max attention logit’ within attention layers. This can cause attention distributions to become overly sharp, making the model focus too narrowly and hindering its ability to capture broader information. While the LAMB optimizer partially addresses this, it still faces limitations because its l2-norm-based trust ratios are less effective at directly influencing the maximum values of query/key weights. Additionally, LAMB’s weight-wise trust ratio can be error-prone, as it overlooks important relationships between weight values within rows or columns.
Addressing these challenges, researchers have introduced a novel optimizer called MERIT: Maximum-normalized Element-wise Ratio for Language Model Large-batch Training. MERIT is designed to enhance the stability and performance of large-batch training for language models. You can read the full research paper here: MERIT: Maximum-normalized Element-wise Ratio for Language Model Large-batch Training.
How MERIT Works
MERIT incorporates two main innovations:
First, it leverages the **max-norm** to calculate the trust ratio. The max attention logit is directly related to the maximum absolute values (max norms) in the key and query matrices of attention layers. By using max-norm instead of the l2-norm (as in LAMB), MERIT can more effectively constrain these extreme values, thereby limiting the spikes in the max attention logit. This approach provides larger updates to extreme weight values, preventing them from becoming excessively large.
Second, MERIT constructs **element-wise trust ratios**. Recognizing that weights often exhibit similarities within rows or columns due to multi-headed self-attention and outlier dimensions, MERIT calculates ratios along both rows and columns. For each element, it selects the larger of these two ratios. This finer-grained approach focuses on local weight structures, providing more robust update scaling and improving training stability by mitigating the negative impact of extreme values from other rows or columns.
The MERIT algorithm also includes an element-wise clipping mechanism, which limits the maximum update magnitude to 1 across all parameter dimensions. This further enhances the overall stability of the large-batch optimization process.
Experimental Validation and Results
Extensive experiments were conducted using various sizes of GPT-2 models (125M, 355M, and 770M parameters) trained with large batches on the OpenWebText corpus. MERIT was compared against existing optimizers such as LAMB, AdamW, Lion, and Sophia-G.
- MERIT consistently achieved lower validation loss across all GPT-2 model sizes, demonstrating superior performance. The performance gap between MERIT and baselines widened as the language model size increased.
- Notably, MERIT enabled the use of significantly larger batch sizes without performance degradation. For instance, GPT-2 Medium could be trained with a 6k batch size using MERIT, comparable to a standard batch size (480) with AdamW, but with 48 billion training tokens.
- The optimizer also showed improved zero-shot evaluation performance on tasks like LAMBADA, WikiText, and SuperGLUE, indicating better generalization capabilities.
- Further experiments on Llama models also demonstrated MERIT’s consistent performance improvements in large-batch training scenarios.
- Analysis of the loss landscape revealed that MERIT converges to flatter minima, characterized by smaller top eigenvalues and trace of Hessian matrices, which is often associated with improved generalization.
- Crucially, the computational overhead introduced by MERIT’s maximum-normalized and element-wise trust ratio calculations was found to be minimal, adding only about 1% extra computational cost compared to LAMB.
Ablation studies confirmed the necessity of each design choice in MERIT—element-wise clipping, the weight-wise ratio bound, and the element-wise ratio—all contributing significantly to its enhanced performance.
Also Read:
- Bi-LoRA: A Dual-Module Approach to Efficiently Improve Generalization in Large AI Models
- Accelerating Large Language Models with Arbitrary Precision Computing
Conclusion
The MERIT optimizer represents a significant advancement in accelerating the pre-training of large language models. By integrating max-norm-based trust ratios and element-wise refinement, MERIT effectively addresses the challenges of max attention logit growth and provides more robust update scaling. This leads to improved training stability, better generalization performance, and the ability to utilize larger batch sizes, paving the way for faster development and iteration of next-generation language models.
