MERIT: Maximum-normalized Element-wise Ratio for Language Model Large-batch Training
Authors: Yang Luo
Zangwei Zheng
Ziheng Qin
Zirui Zhu
Yong Liu
Yang You
Venue: ICML 2025 poster (2025)
Keywords: large-batch training
max attention logit
language models
Primary area: optimization

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.

MERIT: Maximum-normalizzed Element-wise Ratio for
Language Model Large-batch Training
Yang Luo
Zangwei Zheng'
Ziheng Qin'
Zirui Zhu'
Yong Liu'
Yang You'
Abstract
1. Introduction
Large-batch training has become cornertonne in accelerating the training of deep neural net- works, yet poses challenges in optimization and generalization. Existing optimizers like AdamW W present performance degradation during language models' large-batch training, due to the informa- tion 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 rea- son is that 1z-norm- 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 over- looks relationships of weight values within rows or columns. Building on these observations, we propose a novel optimizer, MERIT, which lever- ages the max- max-norn norm to calculate the trust ratio to constrain the max attention logit more effectively. Moreover, 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 vari- ous sizes of GPT-2 models demonstrate the su- perior performance of MERIT. Notably, during the training of GPT-2 Medium, MERIT enables 6k batch size without any performance degra- dation compared the standard batch size (480) with 48B training tokens. This work highlights the importance of considering the max attention logit and finer-granularic trust ratio in large-batch training. successfully improves the training stability and paves the way for larger batch us- age, enabling faster development and iteration of large language models. Code is available
The advent of large language models has revolutionized natural language processing, achieving unprecedented per- formance across a wide range of tasks (Touvron et al., 2023; Dubey et al., 2024; OpenAI, 2024). However, the increasing size and complexity of language models always result in a high time cost for the training. With the growing availability of powerful GPU clusters and specialized hardware acceler- ators, large-batch training can dramatically reduce the time required to train state-of-the-an models, making it possible to iterate faster and explore more ambitious architectures by processing more data in parallel.
While large-batch training offers the potential for increased parallelism and faster convergence, it also introduces com- plex optimization dynamies that can impede model perfor- mance and stability (Keskar al., 2017: Goyal et al., 2018; Shallue et al., 2019). Training large language models with large batches typically encounters one main issue: research has shown that training with large batches often leads to models performing poorly on unseen data. When using Adam optimizer with large batch size, Figure shows clear performance degradation, requiring additional training tokens to reach comparable generalization levels.
This paper identifies crucial problem in large-batch train- ing of language models: we observe the sharp increase of max attention logit in attention layers during the train- ing process using Adam W optimizer (Kingma & Ba, 2017; Loshchilov & Hutter, 2019). The inflated max attention logit can result in overly sharp attention distributions, po- tentially causing the model to focus on specific tokens or patterns overly, thus hindering its ability to capture overall information in the data (Zhai et al., 2023).
While LAMB (You et al., 2020) successfully reduces the max attention logit in the first layer of GPT-2 models (Rad- ford et al., 2019; Brown et al., 2020) by applying a weight- wise trust ratio, it faces limitations in further decreasing the value in the medium layer. The limitation arises from the lower efficiency of le norm than max norm (largest abso- lute value) in preventing query/key weights from reaching extremely high values.
School of Computing. National University of Singapore, Singa- pore. Correspondence to: Yang You <youy@ comp.nus.ceoo edu.".
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).
Moreover, our analysis reveals that rows and columns in

---

MERIT: Maximum-normmaizze Element-wise Ratio for Language Model Large-batch Training
Steps Reach 3.28 Loss
GPT2-Small
Steps to Reach 2.98 Val Loss GPT2-Medium
2
Ideal Scaling
214
Ideal Scaling
AdamW
AdamW
LAMB
LAMB
MERIT
MERIT
213
213
Steps
Steps
2
2
2
211
27
28
2
29
2
2
2
Batch Size
Batch Size
Figure I. both GPT-2 models, the connection between hatch size the numher steps required reach a specific validation loss follows similar pattern. At first, the batch size increases, there is phase of ideal scaling (shown dotted line) where doubling the batch size cuts the necessary steps in half. This is followed by period where the benefits start to decrease. Eventually, a point is reached where further increasing the batch size (data parallelism) offers additional advantage. This final stage represents the upper limit of large-batch training effectiveness.
large-batch trained weights often share similarities. The ne- glect of these relationships in the weight-wise ratio method proposed in LAMB leads to training instability as fails to mitigate the negative impact of extreme values from other rows columns. Given rows and columns exhibit high sim- ilarity, allows for calculating element-wisg ratios by con- sidering the weights within the same rows/columns, while eliminating influences from other rows/columns. The pro- posed emer-granularity ratios focus on local weight struc- tures, resulting in more stable large-batch training for lan- guage models.
2. Related Work
2.1. Large-batch Training
Scaling up batch sizes during the training of deep neural networks has been active area of research, as it allows for better parallelization across multiple GPUs and reduces time-to-train. However, naively increasing the batch size often leads to degraded model performance, phenomenon dubbed the "generalization gap"". Several techniques have been proposed to enable large-batch training without com- promising accuracy. Goyal al. (2018) showed that linear scaling of the learning rate with respect to the batch size can maintain model quality for batch sizes up to 8K on Ima- geNet. Other works proposed novel optimization algorithms like LARS (You et al., 2017) and LAMB (You et al., 2020) that dynamically adapt laser-wise learning rates based on pa- rameter norms and momentum. Liu et al. (2022) introduced more efficient SAM (Foret al., 2021) variant for training Vision Transformers (Dosovitskiy et al., 2021) using large batches and Luo et al. (2023) explored memory-efficient optimization techniques for large-batch training of language models. However, LAMB still presents a large max atten- tion logit and shows weight-wise trust ratio calculation containing some error, leading to large-batch training per- formance degradation.
Inspired by these insights, we propose novel optimizer, MERIT, introducing max-norm-based trust ratios to pre- ciscly limit the maximum of query/key matrix and finer- grained ratios for focusing on specific weight structures. We conduct extensive experiments to evaluate the perfor- mance of MERIT compared to existing optimizers across various sizes GPT-2 models. Our findings demonstrate the potential of MERIT to enhance large-batch training by improving convergence properties and generalization per- formance. This work contributes to the ongoing exploration of optimization strategies in large-batch training and high- lights the significance of finer- greiindd ratio calculation in designing effective optimizers.
2

---

MERIT: Maximum-normmaiized Element-wise Ratio for Language Model Large-batch Training
(a) Max Attention Logit- - 1st Layer
(b) Max Attention Logit- 12th Layer
Max Attention Logit Value (Log Scale)
Max Attention Logit Value (Log Scale)
AdamnW MERIT
10
10²
10
AdamiW LAMB
10
10
0
500
1000
1500
2000
500
1000
1500
2000
Steps
Steps
Figure 2. Max attention logit of self-attention layers during the large-hatch training of GPT-2 medium model using three optimizers. (a) Max Attention Logit of first self-attention layer. (h) Max Attention Logit of medium (12th) self-attention layer. A comprehensive visualization is availabe in Appendix I.
2.2. Max Attention Logit
high learning rates, which they termed attention logit growth. Consequently, the attention weights collapse to one-hot vectors and cause unstable training, phenomenon termed attention entropy collapse by Zhai et al. (2023).
The relationship between max attention logit and training stability of transformers has been explored extensively. Re- searchers have previously documented that Transformer training fails when the attention logits become large. De- hghhnn et al. (2023) addressed the challenge of uncontrolled growth in attention logits large-seat transformers by im- plementing query/key normalization, effectively stabilizing the training process and preventing the one-hot atten- tion distributions typical in models with parameters nearing 8 billion. Wortsman et al. (2024) observed the loss diverges and training fails when the max attention logit exceeds ap- proximately 10 Zhai et al. (2023) identified attention entropy collapse as a common issue in Transformer training across various domains and tasks and proposed oReparam to reparameterize the weights of linear layers using spec- tral normalization and learned scalar. Nevertheless, the significant increase in the max attention logit value during large-batch training remains unexplored and leads to the poor performance of existing optimizers.
In the large-batch training of GPT models, larger learning rates are needed compared normal batch sizes, AdamW- based training consistently presents similar max attention logit sharp increase that leads to one-hot-vector attention output analyzed in Appendix C, limiting the expression ability of attention layers. As shown in Figure 2(a), the max attention logit the first self-attention layer during large- batch training with AdamW significantly exceeds the value observed in small-batch training presented Figure 11, leading training instability and performance degradation. As result, AdamW-based brrge-bacchig leads to worse generalization performance than small batches.
3.2. Trust ratio in LAMB
A distinguishing feature of the LAMB optimizer is its im- plementation of the "trust ratio", a mechanism designed to dynamically adjust the learning rates for each neural network layer based on their respective weight norms and update norms. The trust ratio for particular weights w at time is defined as the ratio of the lo-morm of weights to that of updates w:
3. Preliminaries
3.1. Max Attention Logit Growth in Large-batch Training
In the self-attention layer of a Transformer (Vaswani et al., 2017), attention logits are calculated by combining queries Q: and keys K, using the formula 2; = (Q., K2)/ /d., where dy represents the head dimension. These logits are then processed through softmax function to generate at- tention weights aggregate values V,, The max attention logit is defined the max value among the computed atten- tion logits, max ziy. Dehghani et al. (2023) observed that the attention logits became large when using relatively
W2
R
(1)
[ + Xwell
where u, Veete || denotes 12-norm and À the weight decay parameter. Since large-batch training employs high learning rates, LAMB addresses the issue of excessively large or small gradients/updates by adaptively rescaling LRs per layer. This ensures that updates remain aligned with the
172.
3

---

MERIT: Maximum-normallzzed Element-wise Ratio for Language Model Large-batch Training
scale of the initial weights.
Specifically, let w E Rax* and E Ruxa be the weight matrix and update matrix separately, with w(() representing elements at the i-th row and w() representing elements at the j-th column. We first calculate the row-wise ratio r for each row y(0ر = || [mm ///l||/((/// [lm and the column- wise ratio с for each column clo) = [[/ (|/ Lastly, the final glement-wise ratio s determined by s(o-)) max(r(0) cb This improved approach seeks to boost the capacity of the optimizer to adjust to specific weight structures in different parts of the network, leading to improved convergence and generalization performance in language models.
Through the design of trust ratios, LAMB optimizer achieves balance that allows it to exploit the computa- tional benefits of large batch sizes without compromising the robustness of the model training process. However, the increment of max attention logit still exists during large- batch training as presented in Figure 2(b).
4. Algorithm
4.1. Maximum Normalized Ratio
The max attention logit is directly relevant to the max norm (largest absolute value) in key matrix Wx and query ma- trix Wo. as evidenced by the equation: attention logits 2 = XW W XTx//VVa /VT,, where X represents the input sequence to a self-attention layer. Hence, the issue out- lined in Section 3.1 can be addressed by considering max norm when developing the trust ratio for training with large batches. A complete analysis is presented in Appendix D.
4.3. MERIT
Algorithm 1 summarizes our proposed MERIT optimizer. The design of MERIT comes from two parts: maximum- normalized trust ratio and element wise refinement. Finally, we implement an element-wise clipping mechanism that limits the max update magnitude to across all parameter dimensions, which mirrors the update strategy of stochastic Sign Momentum Gradient Descent (Bernstein et al., 2018), serving to enhance the overall stability of the large-batch optimization process. The designs incorporated in MERIT successfully address the issue of rapidly increasing max at- tention logits in the middle self-attention layers of language models, as illustrated in Figure 2(b).
However, there is huge disparity between the max norms of query and key weights and their corresponding l2 norms (Appendix E). In this case, Iz-norm-based trust ratios cannot effectively prevent the huge increase of the maximum of Wo/Wn, K- Thus, LAMB often fails to reduce max attention logit in the medium self-attention layer further as depicted by Figure 2(b). We modify the LAMB optimizer using the max norm instead of the l2 norm when calculating the trust ratio. Therefore, the proposed method gives larger updates to extreme values of weights. This helps prevent extreme values in query and key weights from becoming too large, limiting spikes in the maximum attention logit.
Algorithm 1 MERIT
1: Input: E Rdd learning rate {n 17 parameters 0 < BI, ,2 < 1, 0, = 0, uo = 0
2: for t = to do
9: X EnExx Ve(u,, zz).
4:
m t B1mb-1 + (1 - B1)gt
4.2. Element-wise Trust Ratio
5:
1t- + B20r-1 + (1 - B2)
To further improve the large-batch training performance of layer-wise ratios, we devise an element-wise ratio to capture local weight structures more accurately while maintaining computational efficiency. Due to the multi-headed self- attention mechanism and outlier dimension phenomenon observed throughout the training process of transformers (Kovaleva et al., 2021; Puccetti et al., 2022), weight values exhibit similarities within rows/columns as shown in Figure 3. Given this context, weight-wise trust ratios of LAMB will introduce certain inaccuracies because extreme values in one row/column can adversely impact the training stability of other rows/columns.
6: U1 = Veite
7: Weight-wise Ratio b. = Tu Xara Im
1 Ii
(1) Im
8: Row-wise Ratio r(
we
1| (3) Aau 1 Imm
9: Column-wise Ratio
(
1100 Im..
1a. + -x00 1
10: Element-wise Ratio
s, (f.5) = max (max r
(1) bek
11:
12: W = w: - 7)- clip(s: 8t - (we + Awv)) 1)
13: end for
4.4. Convergence Analysis
To address this limitation, we propose a novel approach that employs an element-wise ratio to leverage the inherent similarity of weights within the same rows or columns of the weight matrix. Our method involves calculating ratios along both rows and columns and then selecting the larger of these two values for each element.
Notation. Let II be the d x d identity matrix, and let = [T,, J2.... ..., In] be its decomposition into column sub- matrices L, = d x dr. For w E Rd let w(t be the block of variables corresponding to the columns of I, i.e., w = w E Rd for = 1,2, ........ h. For any function f Rdd + R. V,f(w) denotes the gradient with respect to
4

---

MERIT: Maximum-normmalize Element-wis Ratio for Language Model Large-batch Training
Query Weights Layer 1
0.35
Key Weights Layer
0.30
100
0.30
100
200
200
Dimersion
Olmera
300
0.20
0.15
0.15
500
600
0.10
600
700
700
0.05
0.05
200
300
400
600
700
100
400 Dimension
500
Dimension
Figure 3. Analysis GPT-2 small's first attention layer reveals patterns in query-key weight magnitudes: weights show high similarities within hoth rows (arising from multi-headed attention architecture) and columns (due the outlier dimension phenomenon).
w+(+) For vectors and v Rdd, we use u2 represent the element-wise square operation and u/o u represent the element-wist division operation. We use || ll. || || and I1-Ilm to denote 12-nomm, 4i-mrmm and max-norm of vector respectively. Consider the following nonconvex stochastic optimization problems of the form
to be at most the training stability improved and we only need to consider the convergence analysis of weight- wise maximum normuulzee ratio that is the lower bound of the proposed MERIT, which is noted as MERIT-W. The following result provides convergence rate for MERIT-W in general nonconvex settings. Following the analysis in (You et al., 2020), we focus on the setting where 31 and
A min (ww = EE- x) ww]l WER
(2)
(1/( ( for all E [7]. Theorem 1. Let 1/: = 1j V 2
where is model parameters to optimize, E is the loss function and IP is probability distribution on the unknown training data X' C I
T, d, = d/h for all [7).. and ay S ||l|m S ax. for all where a, au Then for w; optimized by MERIT-W, we have the following bounds:
Assumption 1. The loss function f(u) is L-smohth with respect to (6), which means there exists a non-negative constant L, Yu, IRd and æ e X:
When B2 0, have
V.l(u,)) - V,l(v, x)1 S Liju - 01)
(3)
1)
E
| [Vf(wa)li
for all E [M]. Let L. (I,,... LTATT represent the vector of Lipschitz constants in dimensions. We use Lang to denote 4.
(olg
(f(wi) f(w") Laag
| 1012 Tdh
0
T
Assumption 2. The variance in stochastic gradients is sub- ject to the following upper bound:
When 82 > 0, we have
E E(V,((w,x)) - (w)(2) s of for all w E I& and [h] E V [V)()]] o for all Rd and i [d].
V ((w)
(4)
2G log(d) V h(1 -B2)
and σ = (oi,... ,,, and = (ō),.. ... R.) are used to denote the vectors standard deviations of stochastic gradient per layer and per dimension separately.
2(f(w)) - (()))
|| vTd
Assumption 3. Gradients are bounded i.e., [VI(w, x)); z G' for all i E [d]. E Rd and E X. Note that such assumptions are typical in the analysis of stochastic first- order methods.
T
where w" represents an optimal solution to the problem outlined in equation 2 and wa an iteration uniformly randomly selected from {w1,".- wr }. For a detailed proof of convergence, please refer to Appendix
Due to the usage of element-wee wise clipping on controlling the worst-case (largest) update size in all parameter dimensions
5

---

MERIT: Maximum-normmaiize Element-wise Ratio for Language Model Large-batch Training
(a) GPT2 Small 2B Tokens
(b) GPT2 Medium BB Tokens
(c) GPT2 Large- 16B Tokens
4.0
4.0
4.0
3.8
3.8
3.8
Vel Loss
3.6
Lion
Lion
Lion
Sophia-G
Sophia-G
Sophia-G
3.4
AdamW
AdamWW
3.2
Adam/W
3.2
LAMB
LAMB
LAMB
3.0
3.2
MERIT
3.0
MERIT
MERIT
500
1000
1500
2000
500
1000
1500
2000
2%8
500
1000
1500
2000
Sleps
Steps
Steps
Figure 4. Final validation loss. (a) GPT-2 Small (125M, batch sizc=1K), AdamW: W: 3,470, LAMB: 3.355, MERIT: 3.280 (b) GPT-2 Medium (355M, batch size=4K), AdamW: 3.172, LAMB: 3.068, MERIT: 2.982. (c) GPT-2 Large (770M, batch size=8K). AdamW: 3.039, LAMB: 2.971, MERIT: 2.897.
5. Experiments
Our learning rate (LR) follows a cosine schedule, with the final L, set to 0.1 of the peak LRR We maintain a constant LR warm-up ratio of 0.02 and apply standard gradient clip- ping (norm) with a threshold of 1.0. In the case of Sophia-G, select 240 examples from each minibatch to compute the diagonal Gauss-Newton and update the diagonal Hessian every 10 steps. We implement the algorithms in Py PyTocch (Paszke et al., 2019) and train all the models in 16 All models are trained on H100 GPUs.
5.1. Setup
Language modeling. We conducted large-batch training experiments on OpenWebText (Gokaslan & Cohen, 2019), training autoregressive models from scratch using settings derived from the Chinchilla scaling law (Hoffmann et al., 2022). Following standard protocol, we set the context length of GPT-2 1024. Our experiments encompassed three model sizes: 125M (small), 355M (medium), and 770M (large). Detailed specifications of the model configu- rations can be found in Appendix A.
Technical details. We mainly evaluate GPT-2 models with their log perplexity and plot the validation loss curves. The results from LAMBADA (Paperno et al., 2016), WikiText (Merity et al., 2017), and SuperGLUE (Wang et al., 2019) evaluations are also included in our experiments.
Baselines. We compare MERIT with LAMB, the dom- inantly used optimizer on large-batch training of lan- guage modeling tasks, Adam with decoupled weight decay (AdamW), Lion (Chen al., 2023), and Sophia-G (Liu et al.,, 2024). For all models, all learning rates are tuned with grid search. The weight decay is set to 0.1 for all optimizers for a fair comparison. We follow Liu et al. (2024) for the settings of 3 values: For AdamW: 31 = 0.9 and $32 = 0.95. For Lion: 31 = 0.95 and B2 = 0.98. For Sophia-G: 81 = 0.92 and 32- 0.99.
5.2. Results
Figure 4 illustrates the validation loss curve on OpenWeb- Text with the same number of steps. MERIT consistently achieves lower validation loss than LAMBB, Adam Lion, and Sophia-G. As the size of the language model increases, the performance gap between MERIT and baselines be- comes larger. Besides, during large-batch training of GPT-2 Large, the performance gap between AdamW and LAMB di- minishes. In contrast, MERIT demonstrates an increased ad- vantage over LAMB under this condition. MERIT achieves a 0.07 lower validation loss on the 123M model (Figure (a)) with the same training tokens, which means a sig- nificant improvement according to training scaling laws in this regime (Kaplan et al., 2020; Hoffmann al., 2022: Liu et al., 2024).
Chinchilla Scaling Law. The Chinchilla scaling law sug- gests an optimal training regime of approximately 20 train- ing tokens per parameter for large language models (Anil et al., 2023; Muennighoff et al., 2023). This principle, de- rived from Hoffmann et al. (2022)'s comprehensive analysis, proposes that model performance maximized when the number of training tokens scales proportionally with the number of parameters under fixed compute budget. Fol- lowing this established benchmark allows for fair compar- ison with other research in the field with limited computing resources.
The scaling law favors MERIT over LAMB. Figure I illustrates the number of steps required for GPT-2 models with varying batch sizes to reach equivalent validation loss on Open WebText. The figure reveals a noticeable decline in generalization performance when training language models with large batch sizes using AdamW W. Importantly, MERIT enables the use of larger batch sizes without compromis-
Implementation. Following the Chinchilla scaling law, we use batch size IK for GPT-2 small with 2B training tokens, 4K for GPT-2 medium with 8B tokens, and 8K for GPT-2 large with 16B tokens for the large-batch training setting.
6

---

MERIT: Maximum-normallzed Element-wise Ratio for Language Model Large-batch Training
GPT2-Small
GPT2-Medium
GPT2-Large
300
160
AdamW
AdamW
120
AdamW
250
LAMB
120
LAMB MERIT
LAMB
Perplexity
200
MERIT
Perplexity
Perplexity
MERIT
100
30
Figure 5. 7ero-shot evaluation LAMBADA and WikiText. With the same number of steps, language models large-hatch pre-trained
with MERIT outperform models pre- pre-trained with AdamW and LAMB on both tasks with lower perplexity scores.
Table Comparison of zero-cvaluation performance for GPT-2 Small and GPT-2 Medium models under Sophia training settings.
Model
ARC
COPA f
HelllaSwag
RACE 1
WIC 1
Avg f
GPT-2 Small -AdamW-Batee Size=480)
43.43
66.00
29.20
29.00
50.16
43.56
GPT-2 Small (MERIT-Batch Size=4k)
45.83
67.00
28.82
27.56
50.16
43.87
GPT-2 Medium (AdamW-Batch Size=480)
49.49
71.00
32.39
30.05
50.00
46.59
GPT-2 Medium (MERIT-Bacee Size=6k)
50.38
70.00
32.32
30.33
50.47
46.70
ing performance (IK for GPT-2 small and 4K for GPT-2 medium) than LAMB. Moreover, the performance gap be- tween MERIT and LAMB, given the same number of train- ing tokens, widens for 355M parameter models compared to 125M parameter models.
5.3. Further Analysis
Performance Gap Between Standard Batch Size and Large Batch Size. We conduct experiments following the training protocol outlined in Liu et al. (2024), using 48 billion tokens for training. Our study compares MERIT's performance with large batches against Adam W's perfor- mance with small batches (batch size=480). As illustrated in Table I, MERIT demonstrates the ability to increase batch sizes to 4K for GPT-2 small pre-training and 6K for GPT-2 medium pre-training without compromising generalization performance. These findings suggest that MERIT enables language models to utilize larger batch sizes as the model scale increases.
The enhanced validation loss per- formance translates to better results in evaluation tasks as shown in Figure 5. We measure the cero-evaluation perfor- mance of trained GPT models LAMBADA and WikiText using perplexity scores. MERIT successfully obtains lower perplexity across both tasks compared with AdamW and LAMB. Our evaluation focuses solely on zero-shot per- formance for pre-trained GPT-2 models, as demonstrating in-context learning typically requires GPT models with at least a billion parameters. Additional evaluation results are available in Appendix G.
Zero-shot Evaluation.
Curvature of Convergence Point. The curvature of the convergence point in the loss landscape differs significantly between small and large batch sizes, impacting model gener- alization and robustness. Large batch sizes often lead con- vergence in sharper minima with higher curvature (Keskar et al., 2017). While these sharp minima may achieve lower training loss, they can result in poorer generalization due to their sensitivity to small changes.
Training Results on Llama. We conducted additional ex- periments using C-Optim to validate the performance of the proposed MERIT optimizer the large-batch training of Llama models (Dubey et al., 2024). In alignment with Chinchilla scaling laws, we trained models on 2.6B tokens from the C4 dataset (Raffel et al., 2023) (batch size= 1 for GPT-2 small and 8B tokens (batch size=4K) for GPT-2 medium. We maintained the same hyperparameter settings, specifically weight decay of 0.1 and beta values of 0.9 and 0.95. Figure 6 demonstrates that MERIT consistently improves performance across different language model ar- chitectures batch training scenarios.
In Figure 7(a), we present the eigenvalue distributions of Hessian matrices at the convergence points of GPT-2 small models pre-trained using AdamW and MERIT algo- rithms. The convergence point achieved by MERIT exhibits smaller top eigenvalue (12.326) and trace (3444.92) than Adam W whose top eigenvalue and trace equal 37.231 1 and 12994.91 respectively, and eigenvalues of MERIT are pre- dominantly confined to the range [-5,5, 5]. This reduced spread of eigenvalues suggests that MERIT converges an overall flatter region in the optimization landscape. Such

---

MERIT: Maximum-normmlizze Element-wise Ratio for Language Model Large-batch Training
(a) Llama 130M 2.6B Tokens
(b) Llama 350M 8B Tokens
4.0
4.0
AdamW
AdamW
3.8
LAMB
3.8
LAMB
MERIT
Val Loss
3.6
3.8
Val
3.4
3.2
3.2
3.0
0
500
1000
1500
2000
2500
500
1000
1500
2000
Steps
Steps
Figure 6. Final validation loss. (a) Llama 130M (batch size=1K), AdamW: 3.277, LAMB: 3.265, MERIT: 3.199 (b) Llama 350M (batch
size=4K). AdamW: 3.014, LAMB: 3.001, MERIT: 2.957.
flat regions, characterized by small eigenvalues and trace, are frrquently associated with improved generalization ca- pabilities in language models.
Table 2. Wall-clock time and TFL.PS.
Optimizer
Model Size
T(step)
TFLOPS
AdamW
770M
242.50s
43.91
AdamW MERIT
LAMB
770M
243.51s
43.73
MERIT
770M
245.466
43.38
10
Density (Log Scale)
AdamW
355M
57.69s
44,05
LAMB
355M
57.93s
43.86
1
MERIT
355M
58.50s
43.43
10
to the attention computation in the transformer architecture. However, illustrated in Figure 8(b), while QK-Norm ables larger learning rates for A in GPT-2 models, it negatively leads to performance degradation in large-batch batch training. A potential explanation for this discrepancy is that QK-Norm aims to stabilize attention computations and inadvertently restricts information flow within the attention layers of language models, which becomes more obvious in large-batch training with much fewer optimization steps.
10
10
10
20
30
Eigenvalue
Figure graphical representation comparing the eigenvalues of Hessian matrices convergence points, contrasting models pre-trained using AdamW versus MERIT..
Comparison of wall-clock time and computational re- sources. In Table 2, we present a comparison of the total computational requirements (measured in TFLOPS) per step and the actual time taken (wall-clock time) A100 GPUs. Following the methodology of Chowdhery al. (2022), report the average time per step (T(step)) and correspond- ing FLOPS. The data in Table reveals that employing maximun-normalized and element-wise trust ratio calcu- lation adds minimal extra computational overhead (1%) compared /2- normom-badeed trust ratio LAMB. Overall, the increase in FLOPS is negligible compared to AdamW and LAMB.
3.600
Adam---Baseine-1e-3
3.575
AdamW-OKKLIN-4--3
AdamW-OKK-IN66-3
3.550
AdamW-QK-LN-8B-3
Loss
3.525
3.500
3.475
3.450
500
1000
1500
2000
Steps
QK-Norm VS MERIT QK-Norm (Dehghani et al., 2023) was developed to mitigate training instabilities encountered when scaling Vision Transformer (ViT) models to unprece- dented sizes with higher learning rates. This technique ap- plies Layer Normalization the query and key vectors prior
Figure 8. QK-Norm leads to performance degradation although improving the feasible learning rates of GPT-2 models pre-training without divergence.
8

---

MERIT: Maximum-normallied Element-wise Ratio for Language Model Large-batch Training
5.4. Ablation Study
that MERIT's gains stem from its combined use of element- wise trust ratio and clipping rather than max-norm alone. Appendix K further supports this by demonstrating the criti- cal role of element-wise operations in early training.
Element-wise Clipping. Figure 9 reveals that MERIT, even without element-wiss clipping, still outper- forms AdamW and LAMB in terms of convergence, albeit with less pronounced improvement. This finding suggests that when we apply the element-wise trust ratio without clipping, certain elements undergo unexpectedly large up- date steps, which can adversely affect the language model performance. These results underscore the importance of element-wise update clipping in large-hatch training scenar- ios, where update magnitudes tend to be larger compared to standard training conditions. For visualization of element- wise clipping ratios during GPT-2 small training, please refer to Appendix H.
Ablation 1:
Table 3. Validation loss comparison across optimizers (AdamW, LAMB, maxLAMB, and MERIT) on GPT-2 small and medium. MERIT achieves the lowest validation loss, demonstrating the ef- fectiveness its clement-wise trust ratio and clipping mechanism.
Optimizer
GPT-2 Small
GPT-2 Medium
Adam
3,470
3.172
LAMB
3.355
3.068
maxLAMB
3,304
3,040
Ablation 2: Weight-wise Ratio Bound. The implemen- tation of a weight-wise trust ratio as a lower bound for element wwie ratios aims to mitigate excessively small up- dates during large-batch training of language models. As illustrated in Figure 9, the application of this lower bound significantly enhances generalization performance, high- lighting the importance of balanced updates across different elements. The bounding ratio is given in Appendix K.
MERIT
3,280
2.982
6. Conclusion
Accelcrating the pre-training of language models heavily relies large batch techniques. In this study, we present the MERIT optimizer, which integrates max norm and local weight information to compute trust ratios. When applied to the large-batch training of GPT models, MERIT enables larger batch size usage than LAMB and Adam while maintaining comparable generalization performance.
Ablation 3: Element-wise Ratio. Element-wise trust ra- tio calculations enhance the generalization capability of language models by providing more robust ratio estimates focusing on local weight structures for individual weight elements. Figure 9 demonstrates the advancement of using an wemmentwwise ratio.
Impact Statement
Sensitivity to learning rates and ablations
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Ablation 3 Ablation 2 Ablation
3,302
3.298
3.323
(+0.022)
(+0.018)
(+0.043)
3.549
3.642
3.891
Acknowledgements
(+0.269)
(+0.382)
(+0.811)
Yang You's research group is being sponsored by NUS startup grant (Presidential Young Professorship), Singapore MOE Tier-1 grant, ByteDance grant, ARCTIC grant, SMI grant and Alibaba grant.
3.356
3.312
3.291
(+0.076)
(+0.032)
(+0.011)
7e-3
9e-3 leaming rate
1e-2
References
Figure 9. Three ablations present obvious performance degrada- tions, which validates the necessity of three design choices.
Anil, R., Dai, A. M., Firat, O., Johnson, M., Lepikhin, D., Passos, A., Shakeri, S., Taropa, E., Bailey, P., Chen, Z.. Chu, E.,, Clark. H., Shafey, E., Huang, Y., Meier- Hellstern, K., Mishra, G. Moreira, E., Omernick, M.. Robinson, K., Ruder, S., Tay, Y., Xiao, K., Xu, Y., Zhang, Y., Abrego, G. H., Ahn, J., Austin, J., Barham, P., Botha, J., Bradbury, J., Brahma, S., Brooks, K., Catasta, M., Cheng, Y., Cherry, C., Choquette- Choo, C. A., Chowd- hery. A., Crepy, C., Dave, S., Dehghani, M., Dev, S., Devlin, J., Díaz, M., Du, N., Dyer, E., Feinberg, V.,
Ablation 4: Norm Choice. To investigate the impact of the norm choice in the trust ratio computation, we conducted an ablation study where we replaced the weight-wite wise nn-norm with the max-norm, resulting in the LAMB optimizer. This variant excludes the clement-wise ratio and clipping components used in MERIT Results in Table 3 on GPT-2 small and medium show that maxLAMB marginally im- proves over LAMB but underperforms MERIT, confirming

---

MERIT: Maximum-normalizze Element-wise Ratio for Language Model Large-batch Training
Feng, F., Fienber, V., Freitag, M., Garcia, X., Gehrmann, S., Gonzalez, L.,, Gur-Ari, G., Hand, S., Hashemi, H., Hou, L., Howland, J., Hu, A., Hui, J., Hurwitz, J., Isard, M.. Ittcheriahh A.. Jagielski, M., Jia, W., Kenealy, K., Krikun, M., Kudugunta, S., Lan, C., Lee, K., Lee, B., Li, E., Li, M., Li, W., Li, Y., Li, J., Lim, H., Lin, H., Liu, Z., Liu, F., Maggioni, M., Mahendru, A., Maynez, J. Misra, V., Moussalem, M., Nado, Z., Nham, J., Ni, E., Nystrom, A., Parrish, A., Pellat, M., Polacek, M., Polozov, A., Pope, R., Qiao, S., Reif, E., Richter, B., Riley, P. Ros, A. C., Roy, A., Saeta, B., Samuel, R., Shelby, R., Slone, A., Smilkov, D., So, D. R., Sohn, D., Tokumine, S., Valter, D., Vasudevan, V., Vodrahalli, K., Wang, X., Wang, P., Wang, Z., Wang, T., Wieting, J., Wu, Y., Xu, K., Xu, Y., Xue. L., Yin, P., Yu, J., Zhang, Q.. Zheng, S., Zheng, C., Zhou, W., Zhou, D., Petrov, S., and Wu, Y. Palm 2 technical report, 2023. URL https: /arxiv.. orgg abs////////////// 2205 10403.
Pellat, M., Lewkowyez, A., Moreira, E., Child, R., Polo-
zov, O., Lee, K., Zhou, Z.,, Wang, X., Saeta, B., Diaz, M., Firat, O., Catasta, M., Wei, J., Meier-Hellstem,, K., Eck, D., Dean, J.. Petrov, S., and Fiedel, N. Palm: Scaling language modeling with pathways, 2022. URL https: /ttp/aiiivv///////////a/////aava org/ abs 2204. 02311.
Dehghani, M., Djolonga, J., Mustafa, B., Padlewski, P., Heek, J., Gilmer, J., Steiner, A. P., Caron, M., Geirhos, R., Alabdulmohsin, I., Jenatton, R., Beyer. L., Tschan- nen, M., Arab, A., Wang, X., Riquelme Ruiz, C., Min- derer, M., Puigcerver, J., Evci, U., Kumar, M., Steenkiste, S. V., Elsayed, G. F., Mahendran, A., Yu, F., Oliver, A., Huot, F., Bastings, J., Collier, M., Gritsenko, A. A., Birodkar, V., Vasconcelos, C. N., Tay, Y., Mensink, T., Kolesnikov, A., Pavetic, F., Tran, D., Kipf, T., Lu- cic, M., Zhai, X., Keysers, D., Harmsen, J., and Houlsby, N. Scaling vision transformers to 22 billion parameters. In Krause, A., Brunskill, E., Cho, K., En- gelhardt, B., Sabato, S., and Scarlett, J. (eds.), Pro- ceedings of the 40th International Conference on Ma- chine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 7480-7512. PMLR, 23-29 Jul 2023. URL /ttp:::///roceeeee ngs. ml press/ v202 2202/dehghani223..25 html.
Ba, J. L.,, Kiros, J. R., and Hinton, G. E. Layer normaliza- tion, 2016. URL https: hww.hhtss/aarrii..rrrabbbb////// 06450.
Bernstein, J., Wang, Y.-X., Azizzadenesheli, K., and Anand- kumar, A. signsgd: Compressed optimisation for non- convex problems, 2018. URL
Dosovitskiy. A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M. Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. URL https: // openreview net/ forum? ddbbbbbyy
abs/1802. 04434.
Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan,
J.. Dhariwal, P.. Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert- A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D. M., Wu,
J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language models are few-shot learners, 2020. URL https:// arxiv. org/ abs/200s 14165.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., Goyal, A., Hartshorn, A., Yang, A., Mitra, A., Sravankumar, A., Korenev, A., Hinsvark, A., Rao, A., Zhang, A., Rodriguez, A., Gregerson, A., Spataru, A., and et al. The llama 3 herd of models, 2024. URL https: http://axvvvvvvo/ee orgg abs 22002240..0 21783.
Chen, X., Liang, C., Huang, D., Real, E., Wang, K., Liu, Y., Pham, H., Dong, X., Luong, T., Hsieh, C.-J,, Lu, Y., and Le. Q V. Symbolic discovery of optimization algorithms, 2023. URL https: arxiv. org/ abs/ 2302. 06675.
Foret, P., Kleiner, A., Mobahi, H., and Neyshabur, B. Sharpness-awars minimization for efficiently improving generalization. International Conference on Learning Representations, 2021. URL https: net fo htt.ttsss///ppeeeeeeee/ /ttp////pee.....eeeee/ee rM.
Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C. Gehrmann, S., Schuh, P., Shi, K., Tsvyashchenko, S., Maynez, J., Rao, A.. Barnes, P., Tay, Y., Shazeer,
Gokaslan, A. and Cohen, V. Openwebtext cor- pus.
http:://Sky1ioo07..gg. OpenMebTextCorpus. 2011, 2019.
N., Prabhakaran, V., Reif, E., Du, N., Hutchinson, B., Pope, R., Bradbury, J., Austin, J., Isard, M., Gur-Ari, G., Yin, P., Duke, T., Levskaya, A., Ghemawat, S., Dev, S., Michalewski, H.. Garcia, X., Misra, V., Robinson, K., Fedus, L., Zhou, D., Ippolito, D., Luan, D., Lim, H., Zoph, B.. Spiridonov, A., Sepassi, R., Dohan, D., Agrawal, S., Omernick, M., Dai, A. M., Pillai, S.,
Goyal, P., Dollár, P., Girshick, R., Noordhuis, P. Wesolowski, L., Kyrola, A., Tulloch, A., Jia, Y., and He, K. Accurate large minibatch sgd: Training imagenet in hour, 2018. URL https: :/ / //axivvv.. org/ 1706. 627.
10

---

MERIT: Maximum-normalized Element-wise Ratio for Language Model Large-batch Training
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E.,, de las Casas, D., Hendricks, L. A., WelbI, J., Clark, A., Hennigan, T., Noland, E., Millican, K., van den Driessche, G., Damoc, B., Guy, A., Osindero, S., Simonyan, K. Flsen, E., Vinyals, O., Rae, J. W., and Sifre, L. An empirical analysis of compute-optimal large language model training. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho. K. (eds.), Advances in Neural Information Processing Systems, 2022. URL https: //openreview wwt.:/eeww.eet/errreeeeeeee wwweee//oooommaat d
Merity, S., Xiong, C., Bradbury, J., and Socher, R. Pointer sentinel mixture models. In International Conference on Learning Representations, 2017. URL https:// openreview net forum? Btt.uummdddByy72ddde.
Muennighoff, N., Rush, A., Barak, B., Le Scao, T., Tazi, N., Piktus, A., Pyysalo, S., Wolf, T., and Raffel, C. A. Scaling data-constrained language models. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Pro- cessing Systems, volume 36, pp. 50358-50376 Curran Associates, 2023.
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language mod- cls, 2020. URL https://///// arxiv. org/ /tt///bb/20000. 2001. 08361.
OpenAL. Gpt-4 technical report, 2024. URL https:/ / owg/oogg//aaooooooooo. arxiv.org/abs////ooooooooo arxiv.
Paperno, D., Kruszewski, G., Lazaridou, A., Pham, N. Q., Bernardi, R., Pezzelle, S., Baroni, M., Boleda, G., and Ferúández, R. The LAMBADA dataset: Word pre- diction requiring a broad discourse context. In Erk, K. and Smith, N. A. (eds.), Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics ( (Volume 1: Lonng Papers), pp. 1525-1534, Berlin, Germany, August 2016. Association for Compu- tational Linguistics. doi: 10.18653/v1/P16-1111 URL https: htt////aaaaattoooee httphhhoowooo..g..o.ooo.e.
Keskar, N. S., Mudigere, D., Nocedal, J., Smelyanskiy, M., and Tang, P. P. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017. URL https:://////....t h/tpp://////ereeeeeeeee wwwpeerwwwwwwwweeeeeeeeeeeeeeeee forum? id-HloyRlygg.
Kingma, D. P. and Ba, J. Adam: A method for stochastic op- timization, 2017. URL https: / /aaxxaivvvv.//// org/ abs/ 1412. 6980.
Kovaleva, O., Kulshreshtha, S., Rogers, A., and Rumshisky, A. BERT busters: Outlier dimensions that disrupt transformers. In Zong, C., Xia, F., Li, W., and Nav- igli, R. (eds.), Findings of the Association for Com- putational Linguistics: ACL-IJCNLP 2021, pp. 3392- 3405, Online, August 2021. Association for Computa- tional Linguistics. doi: 10. 10.18653/v1/20211..o hind://////2...nnnnnnaaaaall
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Köpf, A., Yang. E., Vetoo Z., Rai- son, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S. Pytorch: An imperative style, high-performance deep learning library, 2019. URL http://awwvv.g.aaaaa/a.
https: / /awpiv.vvvv.////aa//..a/a/aaaa org/ abs 19111011211100 01703.
300. URL https::/p//...cc.. / /a//cclaantt aclanthology..... orw.looo..orro//// www.o/gg00002 2021.
Puccetti, G., Rogers, A., Drozd, A., and Dell' Orletta,, F. Outlier dimensions that disrupt transformers are driven by frequency. In Goldberg, Y., Kozareva, Z., and Zhang, Y. (eds.), Findings of the Association for Computational Linguisties: EMNLP 2022, pp- 2886-1344, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 10.18653/v12022.. findings-emnlp.99 URL https: http///aaaaanooooooyyoo/ org/ 2222/2222.222 findings weweddiiinnnmmlllpp.pmmmmmm 93.
findings- acl...ss-ll.... 300.
Liu, H., Li, Z., Hall, D. L. W., Liang, P., and Ma, T. Sophia: A scalable stochastic second-order optimizer for language model pre-training. In The Twelfth International Confer- ence on Learning Representations, 2024. URL https: / //openreview.ro wet.eewwwwwweeeeeeeeee.e...eeeee forum? aeddxxDDeeNNNo...
Liu, Y., Mai, S., Chen, X., Hsieh, C.-J,, and You, Y. To- wards efficient and scalable sharpness-aware minimiza- tion, 2022. URL https: / htt////avvvvvv// www.:xiwii.ooggooooeooereeee 22003. 02714.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. Language models are unsu- pervised multitask learners. 2019. URL https: orrg///ps/Cooo
Loshchilov, I. and Hutter, F. Decoupled weight decay regu- larization, 2019. URL
org/ abs/
/ //api 1 160025533.
1711.05101.
I Luo, Y., Ren, X., Zheng, Z., Jiang, Z., Jiang, X., and You, Y. Came: Confidence-guided adaptive memory efficient optimization. In Proceedings of the ó1st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 4442 4452-4453, 2023.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Explor- ing the limits of transfer learning with a unified text-to- text transformer, 2023. URL https: / //aaxiiv.o.../ /rrg/ abs/1910.106m

---

MERIT: Maximum-norrmalized Element-wise Ratio for Language Model Large-batch Training
Shallue, C. J., Lee, J., Antognini, J., Sohl-Dickstein, J., Frostig, R., and Dahl, G. E. Measuring the effects of
in 76 minutes, 2020. URL https::/aaxiv..ott
abs / 1944. 00962.
data parallelism neural network training, 2019, URL https: arxiv. org/ org/abs/////oooo.............bra /1111.11.. 03600.
Zhai, s., Likhomanenko, T., Littwin, E., Busbridge, D., Ramapuram, J., Zhang, Y., Gu, J., and Susskind, J. M. Stabilizing transformer training by preventing attention entropy collapse. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B.. Sabato, S., and Scarlett, J. (eds.), Pro- ceedings of the 40th International Conference on Ma- chine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 40770-40803. PMLR, 23-29 Jul 2023, URL /tthttpss//prroeeddnng ml press/ v202/zhai22 html.
Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I., and Salakhutdinov, R. Dropout: a simple way to prevent neural networks from overfitting Mach. Learn. Res., 15(1))1929-1958 January 2014. ISSN 1532-4435.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E., and Lample, G. Llama: Open and efficient foundation lan- guage models, 2023. URL https: http:://axxiv.oo///
abs/2302.0 139771
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L. u., and Polosukhin, Attention is all you need. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., neurips. 2017. URL httpss meersiippp.innnnnn........... co/pappe_fflllessper////////
355e223547cce111o00331eaa6368 pdf.
Wang, A., Pruksachatkun, Y., Nangia, N., Singh, A., Michael, J., Hill, F., Levy, O., and Bowman, S. Super- glue: A stickier benchmark for general-puu purpose language understanding systems. In Wallach, H., Larochelle, H., Beygelzimer, A. d'Alché-Buc, F, Fox, E., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https: https::/prrcceeinng.. neurips. cc httppppp__fllllepppprrr/////////
449 1496FFA44eee pdf.
Wortsman, M., Liu, J., Xiao, L., Everett, K. E., Alemi, A. A.. Adlam, B., Co-Reyes, J. D., Gur, I., Kumar, A., Novak, R., Pennington, J., Sohl-Dickstein. J., Xu, K., Lee, J., Gilmer, J., and Kornblith, S. Small-scale proxies for large-scale transformer training instabilities. In The Twelfth International Conference on Learning Represen- tations, 2024. URL https: http:////peeeeeeeee/ec Forum?id=de #0mmxbb2.
You, Y., Gitman, I., and Ginsburg, B. Large batch training of convolutional networks, 2017. URL https: /tttppss///xxiv.../ arg/abs/170m 03888.
You, Y. Li, J., Reddi, S., Hseu, J., Kumar, S., Bhojanapalli, S., Song, X., Demmel, J., Keutzer, K., and Hsich, C.-J. Large batch optimization for deep learning: Training bert
12

---

MERIT: Maximum-normallzed Element-wise Ratio for Language Model Large-batch Training
A. Models and Hyperparamters Configuration
Table 4. Model Configurations and Peak Learning Rate Under Chinchilla Scaling Law.
Model
Size
d model
n head
depth
Lion Ir
Sophia-G
AdamW Ir
LAMB Ir
MERIT
Small
125M
768
12
12
le-4
le-4
le-3
le-2
9e-3
Medium
355M
1024
16
24
8c-5
le-4
4c-3
le-2
9c-3
Large
770M
1280
20
36
8e-5
2e-4
2e-3
8e-3
6e-3
In our study, we examine three GPT-2 variants: small, medium, and large, as described by (Radford al., 2019). The specific configurations for these models are outlined in Table A. We utilize the nanoGPT framework as our codebase. Consistent with nanoGPT's approach, we implement GELU activation functions and omit bias and Dropout (Srivastava et al., 2014) during the pre-training phase.
The GPT-2 models undergo training using the WeenWeeText corpus (Gokaslan Cohen, 2019). We process the text using the GPT-2 tokenizer (Radford et al., 2019). For data organization, we adopt the train-validation split provided by nanoGPT. The training dataset comprises billion tokens, while the validation set contains 4.4 million tokens.
Our training setup employs distributed data-parallel processing with gradient accumulation, allowing for batch sizes of 1K, 4K, and 8K. All model variants are trained using bíloat16 precision. The 125M and 355M parameter models are trained on systems equipped with two H100 GPUs, whereas the 770M parameter models require machines with eight H100 GPUs.
B. Limitations
Comprehensive downstream task assessment. We evaluate large-batch pre-trained models on 7 downstream tasks, which provides valuable but limited insights. A truly comprehensive assessment of language models remains an open research challenge. Our evaluation is further constrained by the modest size of the models studied, which lack advanced capabilities like in-context learning and complex reasoning capability. These limitations indicate the need for caution when extrapolating our findings to larger, more capable models.
Cross-domain applicability and generalization. Our study focuses on large language model optimization. However, a truly versatile optimizer should perform well across various domains such as computer vision, reinforcement learning, and multimodal tasks. Due to computational constraints, have not evaluated the large-batch training performance of our optimizer in these areas. Future work should investigate its efficacy across diverse machine learning paradigms fully assess its generalizability and potential impact.
Scaling up to larger language models and datasets. MERIT has shown promising scalability up to 770M parameter models trained on OpenWebTeext While there are no fundamental barriers to scaling further, our comparison with Adam and LAMB on more extensive models and datasets is constrained by resource limitations. Based on observed improvements in scaling laws and enhanced pre-training stability, we anticipate MERIT to outperform AdamW and LAMB in large-batch training scenarios with larger language models. However, empirical validation of this hypothesis awaits future work with access greater computational resources.
C. Relation between learning rate and attention logits
((++)) = wh W. w(tt1)_ = w((t - V
(((+1) x,W((tt) x,(W:11 - L)) = X.W = V...
xf - = X,(W - w) X,W - VWLL
13

---

MERIT: Maximum-normalized Element-wise Ratio for Language Model Large-batch Training
(f+1) Qi(+1) (K(tt))T
Logits "g
vdl
1 ) vdn
Q10 (X, L)
(X wo L)
K
(X, (x, L) (X L)T
17
vdx
vdz + n
vd
Assuming that the gradients V Vw L and V Ww, L are not dependent (which is typical in gradient descent), the change in the max attention logit is linearly proportional to 1) because r/² terms are negligible compared to terms:
Max Logit s Max Logit! 1 + - xn
D. Influence of Max Norm On Max Attention Logit
Given:
[W = Mq and || Mr
Each element in Wo and W satisfies:
S Weess & Mo and 0 S [WKisl & Mx
Each element in Q and K can be bounded as:
" Σ
"
Q
Σ Xi W Wq S
Ximm|Me Me Σ Xi.ml = MQQ Cx
Similarly,
K5,k S MK Cx
where Cx = Σm= m is constant representing the sum of absolute values the input embeddings for a token and m the hidden size. The attention logit
d Logit = Σ K / < Σ (Mo Cx) Cx)/ Vd V-MMM---
1=1
k=1
Therefore,
Logila. S Vd. MoMMR
Moreover, decoder-only model implementations, each layer includes LayerNorm(Ba al., 2016), which normalizes X to follow Gaussian distribution. This normalization effectively establishes upper bounds Cx on X's values.
Thus,
Logit S Vd. MºMx Odn2 - MoMx).
This proof demonstrates that controlling the max norm of Wq and Wx effectively constrains the upper bound of the max attention logit, which is crucial for large-batch training stability.
14

---

MERIT: Maximum-normmaizzed Element-wise Ratio for Language Model Large-batch Training
E. Huge Difference between l2 Norm and Max Norm
Relative Ratio Between L2 Norm and Maximum Norm
0.996
0.995
0.994
Ratio
0.993
0.992
Query Ratio Layer1
Query Ratio Layer12
0.991
Key Ratio Layer1
0.990
Key Ratio Layer12
500
1000 Steps
1500
2000
Figure 10. Relative ratios between and l2 norm of Wx /Wo E IR 10242x1124 in GPT-2 medium. Ratio is calculated as ([[W]] [[]] [/][]]] in which ||
and || Il denote 12-norm and max norm.
Figure 10 illustrates the relative ratio between le norm and maximum norm for query and key weights across different layers (Layer and Layer 12) during 2K training steps. The ratios consistently remain high, hovering around 0.99-0.996, which reveals a critical insight: the gap between la norm and max norm is huge. When optimizing using only l2 norm, the weight updates affect all parameters globally but fail to constrain the max value of query/key weight matrix.
F. Max attention logit in small-batch training
Figure shows the distribution of max attention logits in small-batch (512) training of the GPT-2 medium model using the same chinchilla scaling law setting. Notably, these max attention logits significantly lower than those observed in large-batch training scenarios. This reduction suggests that the attention outputs are more evenly distributed, which typically leads to improved training convergence and generalization performance.
Max Attention Logit Small Batch
Max Attention Logit Small Batch
AdamW-First Layer
0.085
AdamW-Medium Layer
Max Attention Logit Value
Max Attention Logit Value
2.4
0.080
0.075
2.2
0.070
0.065
2.0
0.060
0.055
0.050
0.045
3k
6k
9k Steps
12k
15k
3k
6k
12k
15k
Steps
Figure 11. Max attention logit of self-attention layers during the small-batch training of GPT-2 medium model using three optimizers. (a) Max Attention Logit first self-an eeff-attention layer. (b) Max Attention Logit of medium self-attention layer.
G. Zero-shot evaluation LAMBADA and WikiText
Zero-shot Evaluation. The improved validation loss leads to better downstream task performance, as demonstrated in Figure 12, When comparing models with equal pre-training steps, the GPT-2 variants trained using MERIT consistently
15

---

MERIT: Maximum-normallzed Element-wise Ratio for Language Model Large-batch Training
outperform those using LAMB and AdamW in zero-shot accuracy across subtasks.
GPT2-Small
GPT2-Medium
GPT2-Large
80
80
80
LAWE
Accuracy (%)
60
Accuracy (%)
60
Accuracy (%)
60
40
40
40
20
20
20
hellaswag
easy
hellaswag
easy
copa
hellaswag
race
Figure 12. Zero-shot evaluation on SuperGLUE benchmark. Given an equivalent number of training steps, models that undergo large-batch
pre-training using MERIT exhibit higher accuracy than those pre-trained with A and LAMB on most tasks.
H. Clipping Ratio in GPT-2 Small Large-batch Training
Figure 13 presents the wlement-wise clipping ratio GPT-2 small training setting (2B tokens) for the 1st, 6th, and 12th layers.
The analysis of clipping effects across different layers reveals distinct patterns in gradient update behavior during training. The input layer (Layer 1) maintains near-zero clipping ratios throughout, suggesting that early-layer gradient updates rarely require adjustment. In contrast, the middle layer (Layer 6) experiences more substantial clipping, with ratios peaking at 12% during later training stages. The output layer (Layer 12) shows minimal clipping, with ratios reaching only 0.25% at maximum.
This layered pattern demonstrates that the clipping mechanism primarily influences the middle layers, leaving input and output layers unaffected. This suggests that the clipping mechanism primarily stabilizes middle layers rather than applying uniform constraint across all layers. Most gradient updates maintain their original direction, with the most significant stabilization occurring in middle layers where feature representations undergo refinement.
Clip Ratio Layer1
Clip Ratio Layer6
0.12
Clip Ratio Layer12
0.0025
0.10
0.0020
Clopping Ratlo
0.0015
Cloping
Cloping
0.0010
0.04
-1
0.02
-2
500
1.5k
2k
0.00
0
500
1k Steps
1.5k
2k
0.0000
500
1k
1.5k
2k
Steps
Stepe
Figure 13. Flement-wise clipping ratio for different layers of GPT-2 small during large-batch training.
I. Visualization of Max-Attention Logits
Figure 14 presents the max attention logits (MAL) across all layers in the training of GPT-2 medium, observing that MERIT consistently stabilizes training by mitigating extreme MAL spikes. In shallow layers (1-8), both MERIT and LAMB effectively reduce the high MAL values seen with AdamW. In mid-depth layers (9-17), MERIT's max-norm control further mitigates the MAL spikes present in LAMB. In deeper layers (18-24), MERIT exhibits comparable or marginally higher MAL values, though these remain significantly lower than in mid-layers and have negligible effects on convergence. These findings highlight MERIT's effectiveness in controlling destabilizing attention logits, especially in crucial mid-layers, facilitating stable large-batch training and better convergence.
16

---

MERIT: Maximum-normalized flement-wise Ratio for Language Model Large-batch Training
MAL of Each Layer for Different Optimizers
Batch Size 4K
Layar1
Layer 2
Layal
Layur 4
Layer
Layut
Layer?
Layer
Loya!!
Layor 10
Layor 11
Layor 12
Layerr
Loyer
Layer
Layer
Layer
Layer
Lugar
Lurae
Figure Max attention logit of self-attentin layers during large-batch training of GPT-2 medium using three optimizers layers
17

---

MERIT: Maximum-normmaiized Element-wise Ratio for Language Model Large-batch Training
J. Convergence Proof of Theorem
Proof. We study how MERIT-W converges across different minibatch sizes. To begin, let's review the update equation of MERIT-W
wi-1 w// 1/.
(4)
(1)
(2)
m
u
(5)
(4)
10
m
for all i E [h].
Since the function f is L-smooth, we obtain the following:
h
f(w(+1) S f(w)) + (V (f
w - - -w(")
+
Li
1 - w/
(00 12
==1
S (wi) -tt
h
d
(3) mll) x
[V:J[we)]> x
U
(t.))
h
L.d
(6)
Ilm
2
The above inequality follows from the Lipschitz continuity the gradient. We bound term Ti the following manner:
h
de
(5.5)
(IVa/(cee)ss x
Ti < 1/t
ml
x
||4 (5) Ilm
di
B2
(Iuz (10 x [V.[(we)]s x 96.)
< 1/6
2 log(ds)
h
de
(tj)
Imm x [if(we) x
ut
[(sig(((V:f(wv)))) # sign(u;"))) (b))))
l 1l
This follows from the fact that ||u ( ml
v V
(d.)
and Voe S G. If 32 = 0, then Ti can be bounded as follows:
I-B2
h de
(1lu8) 1m x (vr((w)))s)
T
logg(d;)
h
d
(
|ml]
(1)
Il
[w:(( x
(sign((V./(m))))) + sign(u (ignt(
(sign
1/ ll
The rest of the proof for 82 0 is similar to argument for the case B2 > 0, which is shown below. Taking expectation, we have the following:
E[Ti]
< 1tt
B2
E
||0 ( Il
(VVi(w:))l sf
2
(og(,)
(6) [(sign([V # sign(g(.))) (3)
u
77+
E
ml
[vif(ue)lz
j=1
h
de
32
(Iw Il x [Vef(we)), )
-
222log (d))
18

---

MERIT: Maximum-normalizze Element-wise Ratio for Language Model Large-batch Training
h
d
E
aullV (u.11////0eo11111sss///////))d (igr sign(s(?))
Using the bound on the probability that the signs differ, we get:
h(1 - B2) E [ < - G log(d)
[w())) + 1000
T vb
Substituting the above bound on Ti in equation 6, we have the following bound:
h(1 B2)
lôll n2
[[s(w++1)] S f(w)) - nay 2G log(d)
[f(w))] + 1
+
vb
2
Summing the above incquality for to T and using telescoping sum, we have the following inequality:
1(1 B2) B[f(wT+1)] & f(w1) may 2G log(d)
Σ Elll ((())(( a
au
|+ll
|
vb
2
Rearranging the terms of the above inequality, and dividing by nTaa we have:
h(1 32) Σ E {()))) = [(x)) [[[((T+1]] V 2G log() T
T
aullāll1
+
[
Tnoi
avvb
20q
f(w1) - f(w0)
aullōl11
|L|l1 nda
Tha
aivb
201
19

---

MERIT: Maximum-normalized Element-wise Ratio for Language Model Large-batch Training
K. Weight-wise Bounding Ratio Distribution
Weight-wise Trust Ratio Bound GPT2 Small
1.0
0.8
Bound Ratio
0.6
0.4
Layer1
0.2
Layer6
Layer12
0.0
0
250
500
750
1000
1250
1500
1750
2000
Steps
Figure 15. The trigger ratio of the weight-wise trust ratio lower bound during the large-batch training of GPT-2 Medium.
Figure 15 further demonstrates that this lower bound becomes particularly crucial in the latter stages of training. This observation indicates that maintaining minimum update magnitude grows increasingly important as the model nears convergence.
20


Reviews:
[{"review_id": "fZPxKZWWZh", "reviewer": "Reviewer_1Wwm", "summary": "This paper proposes a new optimizer MERIT specialized for neural networks with self attention, especially Transformers, in a setting of large batch size training that typically lead low generalization performance. Based on observations that the training instability of Transformer comes from large amount of max norm at self attention weights, this method introduces (1) weight matrix-wise normalization (trust ratio) using max norm, (2) element-wise normalization with approximating with row/column vectors for efficiency, and (3) additional clipping of weight updates. Experiments show its rapidness of convergence on GPT-2 models compared with other optimizers, and yields better downstream performance.\n\nEDIT after rebuttal: the authors answered my questions appropriately despite its laege burden of additional trainIng. Althouth the claim of the paper is limited for training flagship models, I think it is worth sharing the result to the community. I would raise my score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer 1Wwm for the valuable questions and comments. For the concerns and questions, here are our responses:\n\n**Q1: Tendency of long-run training.**\n\n**A1:** Thank you for the comment. Due to computational resource constraints, we are currently unable to apply our proposed optimizer to training runs involving trillions of tokens. As an alternative, we evaluate MERIT using GPT-2 models, which reflect architectural patterns commonly found in modern LLMs, to investigate its scaling behavior.\n\nThe results below show that in large-batch settings, MERIT yields consistent improvements compared with LAMB as model size grows:\n\n| Optimizer | GPT2-Small | GPT2-Medium | GPT2-Large |\n|:---------------:|:----------:|:-----------:|:----------:|\n| AdamW           | 3.470      | 3.172       | 3.039      |\n| LAMB            | 3.355      | 3.068       | 2.971      |\n| MERIT           | **3.280**  | **2.982**   | **2.897**  |\n\nWe observe consistent gains across all GPT-2 sizes under the Chinchilla training setting. As the model size changes, the gap between AdamW and LAMB decreases, while the performance gains of MERIT over LAMB remain stable.\n\nTo further validate these findings, we conducted additional experiments using the open-sourced Llama implementation (https://github.com/kyleliang919/C-Optim). The results are as follows:\n\n| Optimizer | Llama-100M | Llama-250M | Llama-350M |\n|:---------------:|:----------:|:-----------:|:----------:|\n| AdamW           | 3.277      | 3.121       | 3.014      |\n| LAMB            | 3.257      | 3.096       | 3.001      |\n| MERIT           | **3.199**  | **3.046**   | **2.957**  |\n\nThese results further support MERIT's scalability and its effectiveness in large-batch training across widely used LLM architectures.\n\n**Q2: Interaction with treatments out of optimizer** \n\n**A2:** Thank you for the comment. We investigated the effect of QK normalization in the large-batch training of GPT-2 small using 2B tokens. We first applied it to the AdamW optimizer:\n\n| Method | 8e-4 | 1e-3 | 2e-3 | 4e-3 | 6e-3 | 8e-3 | 1e-2 |\n|:-------------:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|\n| AdamW + QK Norm | 3.608 | 3.593 | 3.524 | 3.475 | 3.527 | 3.496 | 3.671 |\n\nCompared to the AdamW baseline (loss = 3.470), QK normalization enables the use of larger learning rates without divergence. However, it consistently results in degraded performance in large-batch training scenarios, suggesting that its benefits in stability do not translate to better optimization outcomes in this setting.\n\nWe then evaluated QK normalization in combination with the MERIT optimizer. Interestingly, MERIT not only maintains stability but also benefits from QK normalization:\n\n| Model       | MERIT (lr = 9e-3) | MERIT + QK Norm (lr = 1e-2) |\n|:-----------:|:-----------------:|:----------------------------:|\n| GPT2-Small | 3.280             | **3.240**                    |\n\n**Analysis:** The introduction of QK Norm on MERIT successfully reduces the validation loss further, and the potential reason is that the use of element-wise trust ratio enables more stable updating when integrated with QK Norm.\n\nTo ensure a fair comparison, we use a vanilla implementation of GPT-2 (nanoGPT) throughout our experiments when evaluating the proposed optimizer. We will consider more combinations of MERIT with other tweaks in the future!\n\n**Conclusion:** While QK normalization alone does not improve performance with AdamW in large-batch training, it complements MERIT effectively, leading to further performance gains. This demonstrates MERIT's flexibility and its potential to integrate well with auxiliary techniques for scaling up training.", "claims_and_evidence": "The claims are straightforward to understand. The proposed method is based on an existing optimizer Adam and LAMB, but introduces more convervative method of regularization that are targeted on a pitfalls that the existing method overlooks.", "methods_and_evaluation_criteria": "This is a paper proposing a drop-in replacement of LLM optimizer. So the outline of the method itself is not special.\nEvaluation was conducted mainly on comparison of tendency during traiing LLMs, comparison of downstream results, and observation of inner weights and magnitude of activations. All of them supports advantage of the proposed method.\nThe experiments was conducted only by comparing existing optimizers, and any other methods than optimizers that tends to be adopted before replacing the optimizer, e.g., normalization techniques mentioned in 2.2, are ignored.", "theoretical_claims": "The paper involves a proof of convergence of the proposed method. As detailed breakdown of the proof is moved to the appendix, I don't judge that the claim is correct. In an empirical perspective, the algorithm 1 seems to converge if Adam converges, as it involves only the components from Adam and additional normalization that suppress updates.", "experimental_designs_or_analyses": "The experiments are conducted on relatively small settings of experiments: GPT-2 with less than 1B params, batch with up to 8M tokens, trained with several billions of training tokens. These settings are significantly smaller than the current main-stream \"small\" pretrained models (e.g., Llama-3 8B, except the batch size: 8M tokens seems a common settings of training current models). I understand that there is a budget limitation to conduct experiments, but it seems results can't be generalized to the main-stream settings of LLM training.", "relation_to_broader_scientific_literature": "If the claim is generalized, it impacts against training schemes of LLMs, that may result reduction of huge amount of training cost.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "(1) tendency of long-run training: current LLMs are typically trained on trillions of tokens regardless of Chinchilla's law and I would like to understand what can be said if the proposed method was applied to such situation.\n(2) interaction with treatments out of optimizer: In the real scenario, replacing optimizer is kind of the last choice due to its impact against the whole model, and other tweaks on the forward pass (e.g., reparametrization and normalization) tend to be applied to suppress instability. I'm also interested in how the proposed method works together with these methods, and whether applying the proposed method (more generally, replacing the optimizer from AdamW) is more important than other tweaks or not.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gtjHvXkL04", "reviewer": "Reviewer_ajhY", "summary": "The authors propose an extension of LAMB, which is an optimizer for large batch training. They propose to use an element-wise trust ratio using max norm and an element-wise clipping, rather than just an l2-norm based trust ratio as is the case for LAMB. The intuition is that the l2 norm is a poor surrogate for max norm, which is directly correlated with the max attention logit. Using max norm enables larger updates for more extreme weight values, thus prevening exploding logit spikes.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer ajhY for the valuable questions and comments. For the concerns and questions, here are our responses:\n\n**Q1: There may be other large batch training methods > 2019.**\n\n**A1:** Thank you for the comment. Most existing large-batch training studies focus on computer vision tasks (e.g., [1], [2], [3]), while discussions specific to autoregressive language models like GPT-2 or LLaMA remain limited.\n\nAlthough some works have attempted to improve LAMB's efficiency (e.g., [4]), they do not show performance gains. Up to now, LAMB remains the most widely used optimizer for large-batch training in language models.\n\n[1] Large Batch Optimization for Deep Learning Using New Complete Layer-Wise Adaptive Rate Scaling.\n\n[2] Achieving small-batch accuracy with large-batch scalability via Hessian-aware learning rate adjustment.\n\n[3] Revisiting LARS for Large Batch Training Generalization of Neural Networks\n\n[4] SLAMB: Accelerated Large Batch Training with Sparse Communication.\n\n**Q2: More qualitative plots showing the max attention logits across all layers.**\n\n**A2:** Thanks for your suggestion. We have drawn the max attention logits (MAL) across all layers in this anonymous link: https://anonymous.4open.science/r/MAL-4K-2K-277A/.\n\n**Analysis:**\n\n- In layers 1-8, where AdamW shows high MAL, MERIT and LAMB effectively suppress it.\n- In layers 9-17, MERIT reduces LAMB's MAL spikes via max-norm control.\n- In deeper layers (18-24), MERIT shows similar or slightly higher MAL, but the values are much smaller than in mid-layers and have minimal impact on convergence.\n\n**Conclusion:** The qualitative plots confirm that MERIT consistently mitigates extreme max attention logits spikes in critical mid-depth layers, where other optimizers struggle. This contributes to more stable training and improved convergence, particularly in large-batch settings.\n\n**Q3: The authors should briefly describe the \"trust ratio\" earlier.**\n\n**A3:** Thanks for your advice. We will put the introduction of the concept \"trust ratio\" earlier for clearer transmission in our revised paper.\n\n**Q4: Why the it is called an \"element-wise\" trust ratio?**\n\n**A4:** Thank you for your question. We initially referred to our proposed trust ratio as \"element-wise\" to distinguish it from the weight-wise trust ratio used in LAMB.\n\nAs shown in Figure 3, the weight-wise trust ratio in LAMB can introduce inaccuracies since extreme values in one row or column may negatively affect the scaling of other rows or columns. To address this, we designed a finer-grained trust ratio that assigns different values to individual elements, reflecting variations in row/column distributions.\n\nTo mitigate possible misunderstanding, we plan to revise it to **row and column-wise trust ratio** in the updated paper.\n\n**Additional Experiments:**\n\nWe conducted additional experiments comparing MERIT and LAMB under the same settings in Table 1. The corresponding results are shown below:\n\n**GPT-2 Small**:\n\n| **Model**| **ARC ↑** | **COPA ↑** | **HellaSwag ↑** | **RACE ↑** | **WIC ↑** | **Avg ↑** |\n|-|-|-|-|-|-|-|\n| GPT-2 Small (AdamW-Batch Size=480)             | 43.43     | 66.00      | 29.20            | 29.00      | 50.16     | 43.56     |\n| GPT-2 Small (LAMB-Batch Size=4k)              | 44.40     | 64.00      | 28.46            | 27.46      | 50.16     | 42.90 |\n| GPT-2 Small (MERIT-Batch Size=4k)              | 45.83     | 67.00      | 28.82            | 27.56      | 50.16     | **43.87** |\n\n**GPT-2 Medium**:\n\n| **Model**                                      | **ARC ↑** | **COPA ↑** | **HellaSwag ↑** | **RACE ↑** | **WIC ↑** | **Avg ↑** |\n|-|-|-|-|-|-|-|\n| GPT-2 Medium (AdamW-Batch Size=480)            | 49.49     | 71.00      | 32.39            | 30.05      | 50.00     | 46.59     |\n| GPT-2 Medium (LAMB-Batch Size=6k)            | 47.47     | 72.00      | 30.90            | 28.52      | 50.31     | 45.84 |\n| GPT-2 Medium (MERIT-Batch Size=6k)            | 50.38     | 70.00      | 32.33            | 30.33      | 50.47     | **46.70** |\n\n**Analysis:**\n\n- LAMB's performance degrades on most tasks as batch size increases, confirming its sensitivity to scale.\n- MERIT consistently improves or maintains performance across tasks compared to AdamW in small-batch training, even with large batch sizes, demonstrating better generalization and stability.\n- On average, MERIT outperforms LAMB by ~1 point across downstream tasks, showing its practical advantage in both zero-shot and fine-tuning scenarios.\n\n**Conclusion:**  \nThese results reinforce MERIT's strength in scaling up batch size without compromising downstream performance. This further confirms that MERIT enables the use of larger batch sizes without sacrificing performance, in contrast to LAMB.", "claims_and_evidence": "The main claim is that the max attention logits are the main difficulty when doing large batch training and that LAMB poorly addresses this by using a surrogate l2-norm based scaling. They provide evidence by showing that MERIT does reduce the max attention logits qualitatively in later layers and complimentary quantitive results showing good performance with very large batch sizes. The qualitative results are not too thorough in confirming that this does hold in practice, however, the theoretical results are sufficient.", "methods_and_evaluation_criteria": "The authors conduct experiments training on the WikiText dataset with modestly sized GPT models and show good/comparable zero-shot perfance to AdamW despite using a much larger batch size. This seems to be sufficient, although there may be other large batch training methods > 2019 (LAMB) that need to compared with and that I am not too familiar with. Other reviewers might have some comments here.", "theoretical_claims": "The authors provide some more concrete derivations connecting the learning rate and attention logit values. This is important since MERIT scales the learning rate in such a way to avoid spiking the attention logits and subsequently improve training convergence and generalisation.\n\nAdditionally, the authors provide convergence derivations, which looks correct to me.", "experimental_designs_or_analyses": "They appear to be sound and valid.", "relation_to_broader_scientific_literature": "Optimizers and large batch training is a very important topic for scaling training of moderately sized models.", "essential_references_not_discussed": "None that I am aware of.", "other_strengths_and_weaknesses": "No other strengths or weaknesses.", "other_comments_or_suggestions": "For the readers less familiar with this topic, the authors should briefly describe the \"trust ratio\" earlier, or even in the abstract the conceptual idea. It will make those not very closely following this field more easily follow the introduction and motivation before going into the method.", "questions_for_authors": "I am not finding it too clear why the it is called an \"element-wise\" trust ratio? From my understanding, the ratio s_{i,j} is the maximum along the ith row and jth column (of the weight matrix normalised by the updates). Can a different term be used to describe this modification? or can the authors more clearly explain why this naming is used - since the computation is not just dependant on w_{i,j} and u_{i,j}, but also the rows and columns w_{i, :}, w_{:, j).. etc.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BZ44p5TYLl", "reviewer": "Reviewer_QXND", "summary": "This paper proposes a new optimization algorithm MERIT for tackling the large-batch training of language models. The method normalizes the update scale based on the maximum norm of both column and row statistics of the update. Empirical study shows that MERIT with large batch size achieves similar performance as Adam with standard batch size, and exhibits slightly faster convergence than baseline algorithms including Adam, Lion, Sophia-G, etc.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer QXND for the valuable questions and comments. For the concerns and questions, here are our responses:\n\n**Q1: Why MERIT is beneficial to large-batch training?**\n\n**A1:** Thanks for the point. We have provided a theoretical analysis of why the max-norm is efficient in limiting the increasing of max norm of query/key weights:\n\nGiven the query weight $W \\in m \\times n$ and the updating of MERIT:\n\n$$\nW_{t+1} = W_t - \\eta_t \\cdot \\frac{||W_t||_m}{||U_t||_m} \\cdot U_t.\n$$\n\nThen, \n\n$$\n||W_{t+1}||_m = ||W_t - \\eta_t \\cdot \\frac{||W_t||_m}{||U_t||_m} \\cdot U_t||_m \\leq ||W_t||_m + \\eta_t \\cdot \\frac{||W_t||_m}{||U_t||_m} \\cdot ||U_t||_m = ||W_t||\\_m + \\eta_t \\cdot ||W_t||_m.\n$$\n\nAs for LAMB, the upper bound is:\n\n$$\n||W_{t+1}||_m = ||W_t - \\eta_t \\cdot \\frac{||W_t||_2}{||U_t||_2} \\cdot U_t||_m \\leq ||W_t||_m + \\eta_t \\cdot \\frac{||W_t||_2}{||U_t||_2} \\cdot ||U_t||_m \\leq ||W_t||\\_m + \\eta_t \\cdot ||W_t||_2.\n$$\n\nIt is obvious that  $||W_t||_m \\leq ||W_t||_2$ in most cases and usually $||W_t||_m << ||W_t||_2$. Thus, MERIT provides a tighter bound for max norm of query/key weights compared with LAMB.\n\nIn Appendix D, we analyzed how the max norms of query and key weights affect max attention logits (MAL). In the visualization(https://anonymous.4open.science/r/MAL-4K-2K-277A/), MERIT effectively prevents MAL from growing too large in layers 9–17, where LAMB struggles to maintain control.\n\nAs a result, our proposed optimizer can effectively limit the spike of MAL better in language models' large-batch training from both perspectives of theoretical analysis and empirical results.\n\n**Q2: Justification on scaling to large batch size.** \n\n**A2:** Thanks for your suggestion. We have made extensive experiments to provide a comprehensive comparison between AdamW, LAMB and MERIT. Here are the corresponding results.\n\nGPT2-Small with 2B training tokens:\n\nValidation loss:\n\n| Optimizer | bsz = 128 | bsz = 256 | bsz = 512 | bsz = 1K |\n|:-:|:-:|:-:|:-:|:-:|\n| AdamW| 3.172| 3.194| 3.288| 3.470|\n| LAMB | 3.189| 3.176| 3.247| 3.355|\n| MERIT| 3.205| 3.182| 3.214  | 3.280| \n\nLAMBADA Perplexity:\n\n| Optimizer | bsz = 128 | bsz = 256 | bsz = 512 | bsz = 1K |\n|:-:|:-:|:-:|:-:|:-:|\n| AdamW| 109.556| 100.836| 163.994| 279.166|\n| LAMB| 110.658| 96.300| 134.550| 169.678|\n| MERIT| 111.320| 93.789|113.665|154.482| \n\nGPT2-Medium with 8B training tokens:\n\nValidation loss:\n\n| Optimizer | bsz = 512 | bsz = 1K | bsz = 2K | bsz = 4K |\n|:-:|:-:|:-:|:-:|:-:|\n|AdamW| 2.830|2.871| 2.976| 3.172|\n|LAMB| 2.870|2.864| 2.919| 3.068|\n|MERIT| 2.841|2.853| 2.890| 2.982| \n\nLAMBADA Perplexity:\n\n| Optimizer | bsz = 512 | bsz = 1K | bsz = 2K | bsz = 4K |\n|:-:|:-:|:-:|:-:|:-:|\n| AdamW | 32.304 | 33.290| 57.657| 148.211|\n| LAMB | 45.900 | 38.484 | 52.780| 93.699 |\n| MERIT  |32.468 | 34.608| 39.174| 64.160 | \n\n**Analysis:**\n\n- **AdamW** performs well in small-batch settings but degrades rapidly as batch size increases, both in validation loss and LAMBADA perplexity.\n- **LAMB** stabilizes training at larger batch sizes but suffers from degraded generalization, especially visible in perplexity on LAMBADA.\n- **MERIT** consistently achieves lower validation loss and better perplexity across batch sizes, particularly in large-batch settings (e.g., bsz ≥ 2K).\n- The advantage of MERIT becomes more significant as model and batch size scale, demonstrating its effectiveness and robustness in large-batch optimization.\n\n**Q3: The ablation study on the effect of the norm choice is missing.**\n\n**A3:** Thanks for the advice. We have added an ablation study where we replace the weight-wise $l_2$-norm with the max-norm to compute the trust ratio, without using element-wise ratio and clipping, referred to as *maxLAMB*. The results are summarized below:\n\n**GPT-2 Small:**\n| Optimizer | Validation Loss |\n|:-:|:-:|\n|AdamW|3.470|\n|LAMB|3.355 |\n|maxLAMB| 3.304|\n|MERIT| **3.280**|\n\n**GPT-2 Medium:**\n| Optimizer | Validation Loss |\n|:-:|:-:|\n| AdamW | 3.172 |\n| LAMB | 3.068 |\n| maxLAMB | 3.040 |\n| MERIT | **2.982** |\n\n**Analysis:** The results show that simply replacing the $l_2$-norm with max-norm (maxLAMB) provides marginal gains over LAMB, but it still falls short of the performance achieved by MERIT.\n\nAdditionally, as shown in Appendix K, we plot the trigger ratio for the lower bound of the weight-wise trust ratio during GPT-2 Medium training. The result reveals that the element-wise ratio and clipping play a crucial role in the early training phase, supporting the design choice behind MERIT.\n\n**Conclusion:**  The ablation confirms that MERIT's improvement does not stem from max-norm alone. Instead, its combination between element-wise trust ratio and clipping is key to achieving better performance.\n\n**Q4: it would be more clear to present the ablation results seperately in a table, rather than a figure of grid.**\n\n**A4:** Thanks for the advice. We will make corresponding revisions to present the ablation study separately in the updated paper.", "claims_and_evidence": "* **Why MERIT is beneficial to large-batch training.** The paper argues that MERIT benefits large-batch training by reducing the maximum attention logit, which is achieved by utilizing the maximum norm rather than $l_2$ norm used in LAMB. However, the magnitude of attention logit is determined by the weight norm of query and key matrices. It is unclear why adjusting the update scale will constrain the norm of these matrices. The link between the algorithmic design and the reduction of maximum attention logits is missing.\n\n* **Justification on scaling to large batch size.** Since MERIT targets for large-batch training, it is crucial to plot the validation loss of MERIT across different batch sizes, and show that its convergence and downstream performance indeed scales with the batch size. The comparison between the baseline methods such as LAMB and Adam should be presented as well (e.g. the Table 2 of the [1]).\n\n[1] Large Batch Optimization for Deep Learning: Training BERT in 76 minutes", "methods_and_evaluation_criteria": "Yes. The algorithm is based on a seminal work on large-batch training named LAMB. The setting of GPT-2 pre-training is standard.", "theoretical_claims": "The proof is based on strong assumptions such as bounded gradient.", "experimental_designs_or_analyses": "**Effect of maximum norm.** Based on Figure 7, the performance improvement given by clip operation and element-wise ratio are marginal. The main benefit comes from the weight-wise ratio, which is a feature from the LAMB optimizer. Hence, it seems that the main difference between MERIT and LAMB is to substitute of $l_2$  norm into maximum norm. However, the ablation study on the effect of the norm choice is missing.", "relation_to_broader_scientific_literature": "The paper contributes machine learning field by introducing a large-batch training method.", "essential_references_not_discussed": "The paper has discussed the most related works.", "other_strengths_and_weaknesses": "Please see my comments above.", "other_comments_or_suggestions": "The paper's presentation can be improved. For instance, it would be more clear to present the ablation results seperately in a table, rather than a figure of grid.", "questions_for_authors": "Please see my comments above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "d92tZXSPgl", "reviewer": "Reviewer_ybUy", "summary": "This paper addresses the issue of designing an effective optimizer for large-batch training of LLMs.  Normally, as batch size, B, increases, we also increase the LR (often linearly with B, sometimes with the square root of B).  But at some point, optimizers struggle with such high LRs, and become unstable or fail to train well.  Techniques to address this have been proposed previously, in particular LAMB from You et al. (2020).  Meanwhile, people training LLMs have noticed that very unbalanced QK attention logits have contributed to instability in Transformers, and also hyperparameter sensitivity.  Putting these two issues together, this paper proposes that perhaps it’s this “attention entropy collapse” issue that is preventing large-batch training from working well.  They propose MERIT to prevent this.  It’s basically a version of LAMB that scales and normalizes by the maximum weight (max-norm), rather than the L2 norm of the layerwise weights as in LAMB.  MERIT does seem to reduce the max attention logit in the first layer (where it seems to grow particularly high), but so does LAMB, actually.  But MERIT does seem to reduce it more in Layer 12 in a GPT2 medium model.  In experiments at different GPT2 model scales, trained to around 20 TPP, where every optimizer is tuned to their optimum max LR with grid search, MERIT shows small advantages over LAMB, and LAMB over AdamW.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer ybUy for the valuable questions and comments. All suggestions on claims and layout revisions will be revised in the updated paper. For the concerns and questions, here are our responses:\n\n**Q1: Need for the figure plotting how the max attention logit changes as we increase the batch size.**\n\n**A1:** Thank you for the suggestion. We ran extensive experiments to find the best setting for each optimizer and batch size, and then plotted the max attention logit (MAL) changes accordingly.\n\n- **Peak MAL** means the maximum MAL across all layers and training steps.\n- **Layers with MAL > 100** counts how many layers exceed a MAL of 100.\n\nPls check full plots at: https://anonymous.4open.science/r/MAL-F3FE.\n\nPeak MAL (Val Loss in parentheses):\n\n|Optimizer|512|1K|2K|4K|\n|:-:|:-:|:-:|:-:|:-:|\n|AdamW| 326 (2.830)|664 (2.871)|604 (2.976)|1152 (3.172)|\n|LAMB|464 (2.872)|352 (2.864)|480 (2.919)|804 (3.068)|\n|MERIT|172 (2.841)|286 (2.853)|262 (2.890)|458 (2.982)|\n\n**Analysis:** \n\n- Peak MAL generally correlates with higher validation loss.\n- All optimizers show an increasing trend of Peak MAL as batch size grows.\n- MERIT exhibits the slowest growth in Peak MAL, indicating better stability.\n\nLayers with MAL > 100:\n\n|Optimizer|512|1K|2K|4K|\n|:-:|:-:|:-:|:-:|:-:|\n|AdamW|8|9|12|14|\n|LAMB|3|8|10|10|\n|MERIT|3|4|6|9|\n\n**Analysis:** The increasing batch size leads to more layers that have high MAL, which affects the training stability.\n\n**Conlusion:** MERIT reduces both Peak MAL and the number of layers MAL > 100 as batch size increases, enabling more stable and effective large-batch training.\n\n**Q2: Need for the figure plotting how the max attention logit changes across all layers.**\n\n**A2:** Thanks for your comment. We plotted the full MAL distribution across all layers for GPT-2 Medium with batch size 4K. See details at: https://anonymous.4open.science/r/MAL-4K-2K-277A/.\n\n**Analysis:**\n\n- In layers 1-8, where AdamW shows high MAL, MERIT and LAMB effectively suppress it.\n- In layers 9-17, MERIT reduces LAMB's MAL spikes via max-norm control.\n- In deeper layers (18-24), MERIT shows similar or slightly higher MAL, but the values are much smaller than in layers 1-17 and have minimal impact on convergence.\n\nCrucially, MERIT shows significantly lower peak MAL across all layers than LAMB:\n\n|Optimizer|Peak MAL|Val loss|\n|:-:|:-:|:-:|\n|AdamW|1152 (layer1)|3.172|\n|LAMB|804 (layer11)|3.068| \n|MERIT|458 (layer13)|2.982| \n\n**Conclusion**: MERIT more effectively reduces excessive MAL peaks in critical layers (9-17), leading to better model stability and improved validation loss under large-batch training.\n\n**Q3: Figure 1 is suspiciously perfect.**\n\n**A3:** Thank you for your question. In Figure 1, we approximated that an optimizer \"achieves\" the target loss if its final validation loss falls within $\\pm$0.05 of the expected range (e.g., 3.23-3.33 in the left plot). This simplification was made to reflect training efficiency trends under limited computing resources without enough grid search, which is why the figure may appear perfect -  apologies for any confusion.\n\nDue to time constraints, we could not determine the exact step counts for all configurations. To provide a more rigorous comparison, we conducted extensive experiments and report the best validation loss achieved at each batch size:\n\n**GPT-2 Small:**\n\n| Optimizer | bsz = 128 | bsz = 256 | bsz = 512 | bsz = 1K |\n|:-:|:-:|:-:|:-:|:-:|\n|AdamW|3.172|3.194|3.288|3.470|\n|LAMB|3.189|3.176|3.247|3.355|\n|MERIT|3.205|3.182|3.214|3.280|\n\n**GPT-2 Medium:**\n\n|Optimizer|bsz=512|bsz=1K|bsz=2K|bsz=4K|\n|:-:|:-:|:-:|:-:|:-:|\n|AdamW|2.830|2.871|2.970|3.172|\n|LAMB|2.870|2.864| 2.919| 3.068|\n|MERIT|2.841|2.853|2.890|2.982|\n\n**Conclusion:** AdamW shows clear performance degradation as batch size increases. In contrast, MERIT maintains strong performance and its gains over LAMB grow with batch size-highlighting MERIT's scalability and robustness in large-batch training.\n\n**Q4: Figure 6(b) results.**\n\n**A4:** Thank you for your question. We used GPT2-Small for the experiments in Figure 6(b). To further examine the effect of QK-LayerNormalization, we conducted more extensive tests:\n\n**Large-Batch Setting (bsz = 1K):**\n\n|Method|8e-4|1e-3|2e-3|4e-3|6e-3|8e-3|1e-2|\n|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|\n|AdamW + QK Norm|3.608|3.593|3.524|3.475|3.527|3.496|3.671|\n\nCompared to the AdamW baseline (loss = 3.470), QK Norm enables higher learning rates without divergence. However, it consistently results in worse validation loss in large-batch training, indicating that its stability benefits do not translate into better optimization outcomes under large-batch training.\n\n**Small-Batch Setting (bsz = 128):**\n\n|Method|2e-3|4e-3|6e-3|\n|:-:|:-:|:-:|:-:|\n|AdamW + QK Norm|3.171|**3.167**|3.182|\n\nQK Norm slightly improves small-batch performance (baseline loss = 3.172), suggesting its weaker effect in large-batch settings may stem from fewer steps or the change in training dynamics.", "claims_and_evidence": "[013, left]: Claim: “Large-batch training ... poses challenges in generalization”\n- I think there is a misunderstanding about what is meant by the “generalization gap.”  Larger batches eventually perform worse (on train and test error) for the same total number of training examples, because the gradient approximation can only improve so much as B increases – even full-batch GD wouldn’t work much better than just using a very large batch size (but would consume a lot more total examples).  This phenomenon is illustrated in your Figure 1.  Note this has nothing to do with “unseen data” – you would see this same result if all settings instead had a common target *training* loss.  The idea that larger batches can hurt generalization (performance on unseen data) was an idea that was prevalent during the 2010s decade, and I believe has been well and truly refuted, e.g., see Shallue et al., McCandlish et al.  E.g., Shallue et al report, “We find no evidence that larger batch sizes degrade out-of-sample performance.”\n- unless you have a specific reason to think your models are failing to generalize, I would probably re-write your paper not to mention “generalization” at all.  E.g., you can just say “performance” rather than “generalization performance” or “validation loss” as needed.\n- I also find that this paper does not seem to understand that fundamentally, as you increase batch sizes, you get redundant information, so very large batches are just inherently going to not work as well as smaller ones given the same total data (although they will do better with the same number of STEPS/batches).  Sentences like “As a result [of attention entropy collapse], AdamW-based large-batching leads to a worse generalization performance than small batches” are attributing all the degradation to training dynamics, but fundamentally, large batches will just not work as well.\n\n[016, left]: Claim: the reason prior optimizers fail to work well with large batches is due to “sharp increase of max attention logit”.\n- after reading this paper, I am not convinced of this.  Perhaps MERIT is improving on LAMB and AdamW for some other reason?\n- Shouldn’t there be a plot in the main paper showing how the max attention logit changes as we increase the batch size?  If you could show that this was specifically a problem with large-batches, and that MERIT does better with larger batches, and MERIT is kind of designed to address this problem, then I think this claim would be more valid.\n- Note: I would REALLY like the paper to show this.\n\n[029, left]: “MERIT ... leverages the max-norm to calculate the trust ratio to constrain the max attention logit more effectively.”\n- I also found the evidence for this lacking.  We see results in Figure 2 for two specific layers, of one model.  Can’t we get these statistics over all the layers of all the models?\n- If you could show that MERIT reduces the max attention logits more over ALL the layers, that would go a long way to showing that this really was the problem.\n- Note: I would REALLY REALLY like to see this too.", "methods_and_evaluation_criteria": "For Figure 1, since you use a LR decay schedule to 10% of the peak LR, how do you train to exactly the same validation loss?  This relates to m concern: Figure 1 is suspiciously perfect.\n\n[188, left] “the proposed method gives larger updates to extreme values of weights”\n- I suppose LAMB would do this too.  I don’t really get why this is a good idea, or why it prevents anything.\n- Like, I understand why you might want to normalize the update by the max-norm (denominator of ratio), but why you want the updates to be on the scale of the max-norm is less obvious to me.\n- It seems to me you want to prevent having extreme max-norms to begin with, not behave a certain way once they are there?\n- But anyways, it does seem like this helps a little bit (Figure 2b), but it’s not clear at this point that this is really significant for results.\n\nFor 5.1 and 5.2, is the idea that we are already picking a very large batch size, to the point where we expect MERIT > AdamW (i.e., unlike the batch sizes in Figure 1?).  Might be good to say this.  Like, based on Figure 1, it seems you picked the exact smallest batch size such that AdamW, LAMB, and MERIT behave differently at all.\n\nCurvature of convergence point:\n- I’m pretty sure we’re not training GPT2 models to a convergence point.  Do you mean the point where you stopped training?\n- I don’t believe that either AdamW or MERIT has any trouble generalizing from training to test data.  You could check this quite easily, right?  E.g., scatter plot validation loss vs. train loss for all your training runs with each optimizer.\n\nRegarding the Figure 6(b) results:\n- Which GPT2 model is this?\n- Why not try QK-LN-1e-3????  And if that’s better, why not QK-LN-8e-4?  I’m not convinced the LR is tuned here.\n- How can we conclude it’s a problem with large batches?  Did we try this for small batches and QK-LN does work fine?", "theoretical_claims": "4.4 Convergence analysis:\n- You need to explain what you are intending to show, and then explain exactly what you found: is the convergence rate of MERIT better than AdamW or LAMB?  Under what specific assumptions?  With the current discussion, I have trouble understanding why this analysis is useful.  You don’t mention this analysis as a contribution in your introduction, or refer to it in your conclusion.  So what’s the point?  Is it something that should be in the appendix?", "experimental_designs_or_analyses": "Figure 1 is ultimately quite surprising to me: AdamW, MERIT, and LAMB take EXACTLY the same number of steps to reach the target validation loss with the same batch size, across three batch sizes?  I mean... wouldn’t it depend on their hyperparameters to some extent?  It’s just very surprising.  And also, it is suprising that MERIT shows absolutely no degradation at the largest batch size, but AdamW and LAMB do (while AdamW and LAMB show no degradation at all for the smaller batch sizes).  This figure is so perfect that it’s suspicious, is what I’m saying.  Even just cutting off the plot before larger batch sizes... it’s okay if perfect scaling doesn’t continue!  There is no optimizer that can scale perfectly forever!\n- also, note my above point: we need to know the number of steps in advance in order to set the decay schedule... so how can you even determine how many steps it takes to reach a target validation loss?", "relation_to_broader_scientific_literature": "This paper builds fairly directly on LAMB, but incorporating recent findings regarding attention entropy collapse.\n\nI think the comparison with QK-layernorm is very important, as it’s an independent mechanism designed to address this problem.  But unfortuntely, I do not think the comparison is very convincing.\n\n[035, right]: This paper “identifies” the problem of large batches: Sharp increase in the max attention logit in attention layers using AdamW\n- I am a little confused.  Since the problem is only now identified, are you saying that LAMB partially addresses this problem *without identifying it as a problem*.  At this point in the paper, it might be better to say that LAMB was designed to do X, but, although it was not the specific intention, it does help with attention entropy collapse, although less so than MERIT.", "essential_references_not_discussed": "[027, right]: “Training large language models with large batches typically encounters one main issue: research has shown that training with large batches often leads to models performing poorly on unseen data” [citation needed]\n\n[087, right]: ‘dubbed the “generalization gap”’ [citation needed]", "other_strengths_and_weaknesses": "The paper was very well written, and the experiments were conducted in a very convincing way, and training to 20 tokens-per-parameter following Chinchilla was an excellent decision, and grid-searching the LR was very important and well done.", "other_comments_or_suggestions": "Intro:\n\nOverall this was good.\n\nMight have been helpful to point out in the intro that “high max entropy logit” is also known as “attention entropy collapse” and techniques like QK-layernorm have been proposed to address it, but, I assume, not in the context of large batches?\n\nIn the intro, I could have used some help understand what MERIT will *do* exactly.  It uses max norms, it calculates trust ratios, but how does it actually play a role as an optimizer?  Does it precondition, or scale, or normalize weight updates?  Does it reguarlize training in some way?  Assume the reader understands how optimizers work, but may not be familiar with LAMB, for example.\n\nFigure 1: caption: “further increasing the batch sizes offers no additional advantage”\n- just to clarify, this plot does NOT show that phase, right?  Even for vanilla AdamW, increasing the batch size does decrease the number of steps, it looks like this holds even to 2^13.  Having this statement in the caption might lead one to think the figure shows this no-additional-advantage phase.\n- just to emphasize this point: in some cases, if we want a good LLM ASAP, it might be worth using more total training examples, provided we can get our results faster (fewer steps).  So there's a time \"advantage\" even beyond perfect scaling.\n\n[129, right]: that’s a run-on sentence, not grammatically correct, you should probably divide into multiple sentences, and I think you mean “expressive” ability.\n\nPreliminary → Preliminaries ?\n\nEquation (1):\n- is that λ the weight decay?  This does not seem to be defined anywhere in the paper.\n- I’m confused, you’re saying this whole term is the trust ratio?  Or is the ratio not just the ratio of norms, and then we apply this ratio to the weight update???  And this effectively scales the LR?  Because otherwise, as written, this value R looks like a vector to me.\n- Could we preface this whole section by saying, “since large-batch training uses large LRs, LAMB proposes to re-scale these LRs at each layer to mitigate the effects of very large or very small gradients/updates, and to ensure that the updates that are made are on the scale of the original weights” or something like that.  The connection to large-batch training is just stated, but not *explained* in the current version.\n\nFigure 4 caption says “final validation loss” but the figure has intermediate validation losses as well.  I understand you report the final losses in the caption.  Maybe you need to say, “Final validation losses are:” and then continue from there.\n\nPerhaps Figure 6 should just be two separate figures?", "questions_for_authors": "See other comments", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: In this paper, the authors propose a modification to the LAMB optimizer by replacing the L2 norm with the $l_\infty$ norm for trust ratio estimation. They provide a theoretical convergence analysis to support this modification. However, I agree with reviewer QXND to some extent that the argument concerning controlling the maximum attention logits appears limited and somewhat trivial. The paper would significantly benefit from a more comprehensive analysis of the impact of different update scales, and why this benefits the large batch training, as suggested by the reviewer.

Nonetheless, this discovery remains noteworthy for the community. It is also interesting to see that selecting different norms can influence specific training dynamics properties.
Decision: Accept (poster)