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

## TL;DR

The paper introduces MERIT, a novel optimizer that uses max-norm and element-wise trust ratios to improve large-batch training stability and performance of language models, enabling larger batch sizes without degradation.

## Contribution

MERIT is a new optimizer that effectively constrains max attention logits and utilizes local weight structures, addressing limitations of existing optimizers like AdamW and LAMB in large-batch training.

## Key findings

- MERIT enables training GPT-2 Medium with a batch size of 6,000 without performance loss.
- Extensive experiments show MERIT outperforms existing optimizers in large-batch training.
- MERIT improves training stability and supports larger batch sizes for faster language model development.

## Abstract

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.