Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation is Wasteful
Authors: Martin Marek
Sanae Lotfi
Aditya Somasundaram
Andrew Gordon Wilson
Micah Goldblum
Venue: NeurIPS 2025 poster ()
Keywords: batch size
language model
LLM
SGD
Adam
Primary area: deep_learning

Abstract: Conventional wisdom dictates that small batch sizes make language model pretraining and fine-tuning unstable, motivating gradient accumulation, which trades off the number of optimizer steps for a proportional increase in batch size. While it is common to decrease the learning rate for smaller batch sizes, other hyperparameters are often held fixed. In this work, we revisit small batch sizes all the way down to batch size one, and we propose a rule for scaling Adam hyperparameters to small batch sizes.  In particular, rather than holding the decay rate of the second moment fixed across batch sizes, we propose to hold its half-life fixed in terms of tokens. We find that small batch sizes (1) train stably, (2) are consistently more robust to hyperparameter choices, (3) achieve equal or better per-FLOP performance than larger batch sizes, and (4) notably enable stable language model training with vanilla SGD, even without momentum, despite storing no optimizer state. Building on these results, we provide practical recommendations for selecting a batch size and setting optimizer hyperparameters. We further recommend against gradient accumulation unless training on multiple devices with multiple model replicas. Finally, we show that a small batch size combined with an optimizer with a small state size can provide the performance benefits of full fine-tuning while maintaining a similar memory footprint to LoRA.

Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation Is Wasteful
Martin Marek New York University martin.menyu edu
Sanae Lotfi New York University
Aditya Somasundaram Columbia University
Andrew Gordon Wilson New York University
Micah Goldblum Columbia University
Abstract
Conventional wisdom dictates that small batch sizes make language model pretrain- ing and fine-tuning unstable, motivating gradient accumulation, which trades off the number of optimizer steps for proportional increase in batch size. While it is common to decrease the learning rate for smaller batch sizes, other hyperparameters are often held fixed. In this work, revisit small batch sizes all the way down batch size one, and we propose rule for scaling Adam hy perpaarameterss to small batch sizes. In particular, rather than holding the decay rate of the second moment fixed across batch sizes, we propose to hold its half-life fixed in terms of tokens. We find that small batch sizes (1) train stably, (2) are consistently more robust hyperparameter choices, (3) achieve equal or better per-FLOP performance than larger batch sizes, and (4) notably enable stable language model training with vanilla SGD, even without momentum, despite storing no optimizer state. Building on these results, we provide practical recommendations for selecting batch size and setting optimizer hyperparameters. We further recommend against gradient accumulation unless training on multiple devices with multiple model replicas. Finally, show that a small batch size combined with optimizer with small state size can provide the performance benefits full fine-tuning while maintaining similar memory footprint to LoRA.
Introduction
Large batch sizes are widely believed to improve the stability of language model training [1, 2, 3, 4]. As consequence, sophisticated optimizers that perform well large-batch training are increasingly standard practice [5, 6, 7, 8, 9, 10]. In fact, it is common to simulate batch sizes even larger than the maximum batch size that fits into device memory through gradient accumulation [4, 3, 11, 12].
In small batch size pretraining experiments, one may observe loss spikes and heavy instability [13, 14, 15, 16). While it common decrease the learning rate for smaller batch sizes, other hyperparameters, such as the decay rates for the first and second moments in Adam (3, and B2), are often held fixed across batch sizes. We show that if instead of holding 32 fixed, hold the half-life of 32 fixed (measured in number of tokens), stable training is possible all the way down to batch size one, as illustrated in Figure la.
When scaling hyperparameters in this way, we find that small batch sizes can confer computational advantages and greater robustness. In the small batch regime, observe that the speed of conver- gence is less sensitive to optimizer hyperparameters like the learning rate, and momentum becomes
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Batch size
Batch size 512
0.999
3.60
5,0
Adalactor
8; (momentum)
0.99
3.55
Mupn
4.4
0.97
Fneweb
4.2
0.9
3.8
3.40
0.0
1
16
64
256
1024
4096
2 2-12-1122222 112-102-22--- 112--12---2-12---2-
2-1 2
2-112-102"
2-
2-0
2-
2-5
Batch
Learning rate
Learning rate
(a) Large batches require more sophisticated optimizers
(b) Small batches are more robust to hyperparameters
Figure I: Small batch sizes are more robust to optimizer design and hyperparameter values. (a) Loss achieved by transformer decoder-only language model with 30M active parameters trained on 600M tokens of FineWeb-Ede data using SGD, Adafactor, Adam, and Muon. At small batch sizes, all optimizers achieve similar loss; large batch sizes, the gap between optimizers grows. (b) Loss achieved by GPT-2 (124M) trained on 2.5B tokens of Fine Web using Adam when tuning the learning rate and 31 hyperparameters for batch sizes 1 and 512. While both batch sizes achieve similar lowest loss, the smaller batch size is much more robust to the hyperparameter values, reducing need for hyperparameter tuning. For batch size 512, we use the default hyperparameters from Brown et al. [17] and the nanoGPT public repository [18]. For batch size 1, we turn off weight decay and rescale 82 to preserve the token half-life, resulting $2 = 0.9999.
less necessary, as we find Figure 1b. In contrast, large batch sizes require large learning rates to compensate for taking fewer steps at given compute budget. Large step sizes in turn require that the optimizer make predictions about the loss surface far away from the current iterate. We find that predicting such far-away parameter updates requires a sophisticated and carefully tuned optimizer. For very small batch sizes, we find that even vanilla stochastic gradient descent (SGD), with momentum weight decay, becomes competitive for training language models, in contrast to findings from recent work [19]. The ability to perform stable training without momentum can confer significant memory advantages, since do not need to store the corresponding optimizer state, which can have significant memory footprint. Based on our observations, we also revisit Adafactor, memory-efficien variant of Adam, and show that it can be compelling alternative to Adam in the small batch regime, enabling training of larger models in a memory memory-constrrinne constraiind setting.
These findings could have significant bearing on practice. Rather than perform gradient accumula- tion or use the largest batch size that fits into device memory, could be preferable use smaller batch sizes that enable simpler procedures, computational and memory advantages, and require less tuning. Our findings suggest a best practice of using the smallest batch size that maximizes model throughput. We find that these prescriptions hold in both LLM pretraining and fine-tuning.
is common in convex optimization use higher batch sizes as the loss converges a minimum [1, 20]. Intuitively, the higher the gradient to gradient noise magnitude ratio, the smaller the batch size we should take. For example, the gradient noise magnitude is close to zero, then we obtain roughly the same gradient estimate with small batch size as a large batch size but pay far fewer floating-poit point operations (FLOPs), so small batch sizes are efficient. Smaller batch sizes allow to take more steps for fixed FLOP budget. As we converge to a minimum, that same ratio instead often tends to zero since the gradient norm vanishes while the gradient noise may not, in which case we may want larger batch size. We hypothesize that language model training lives in the far-from-convergent regime where small batch sizes are efficient. Following compute-optimal scaling laws, we would not train language models into the convergence regime since could achieve better performance by instead training a larger model on less data [21, 22].
The paper is structured follows: in Sections 2 and 3, we cover the background and related work for language model training. Section we demonstrate that small batch sizes not only perform on par with larger batch sizes for several optimizers that we carefully tune, but they also exhibit more robustness optimizer and hyperparameter choices. We also show that the gap between vanilla SGD and more sophisticated optimizers shrinks in the small batch regime. Moreover, we derive scaling heuristics for Adam's hyperparameters and demonstrate that they yield improved performance when transferred to new, larger-scale settings. Finally, apply these results memory-efficien

---

fine-tuning using small batch size and Adafactor, which achieves the best trade-off between memory footprint and performance. Motivated by our findings, we provide practical prescriptions for training language models in Section 5. Namely, we recommend using the smallest batch size that maximizes model throughput and advise against gradient accumulation for most practitioners.
2 Background: Optimization for Language Models
In this section, we briefly review relevant optimization algorithms for language model training and the role of their hyperparameters.
Stochastic gradient descent (SGD) [23, 24]. SGD updates model parameters by computing gradients
on mini-batches of data as follows: B841 0₁ 19t,, where 0 are the model parameters at step t, the learning rate, and Sr Ve is the gradient of the loss function with respect 0g averaged over B samples. We refer to B as the batch size and consider the extreme case of B 1 in our experiments.
We show that although SGD extremely simple, it can perform competitively when correctly tuned at small batch sizes. Momentum often added to accumulate a moving average of past gradients, smoothing updates, and accelerating convergence. However, we focus on vanilla SGD in our experiments make the point that, with a small batch size, even momentum is unnecessary.
Adam [25]. Adam is an adaptive optimizer that maintains an exponential moving average (EMA) of the gradient and squared gradient, referred to as the first and second moments, and denoted my and Dr. The timescales of these moving averages are controlled by decay rates 31 and :3:
my = 81m4 + -9
U 1201-1 + (1 82)
The model parameters are then updated as follows: 0z-1 = 0g -or++ where is small constant to prevent division by zero.
Higher values of 31 and 32 place more weight on past gradients, meaning that the moments are averaged over longer timescales and evolve slowly. When training language models, values like 3, = 0.9 and 32 [0.95, 0.98) are often used in practice [17, 26]. While most researchers and practitioners do not scale 31 and 32 with the batch size, we show in work that scaling B2 in particular is crucial for achieving good performance with Adam at a small batch size.
Moment half-life. We found it helpful across many of our experiments measure the typical timescales (measured in number of tokens) that the first and second moments in Adam are aver- aged over, instead of directly working with 31 and (3..
8
B1
B²
Decay
At each Adam update step [eq. (1)], the contri- bution of previous gradients the first moment gets reduced by a factor of Bn and the contri- bution the second moment gets reduced by factor of 32. Hence, by definition, there exists certain number of steps n, after which the con- tribution of any mini-batch gradient is halved: 3" 3. Since every update step corresponds to observing (B T) tokens, where B is the batch size and T the sequence length, we can con- vert the number of optimizer steps into number of observed tokens. As such, we can express the number of tokens it takes to decay mini-batch gradient by a factor 2 as the half-life of the
B3
8a
half-life
B
g
B7
0
3
5
Optimizer step
Figure 2: Moment half-life. The first and second moments in the Adam optimizer are exponential moving averages of past mini-batch gradients. At each step, the contribution of past mini-batch gradi- ents decays by factor of B. After given number of optimizer steps (or equivalently, after given number of training tokens), the contribution of any batc- gradient will decay by factor of We call this number of tokens the decay half-life.
1
decay rate, denoted by l1/2, where 3 4/2 The half-life provides measure of the "typical" timescale that gradients are averaged over. For
3

---

example, if we use a half-life of b1/2 = 10M tokens, this means that the second moment is averaged over roughly 10M tokens. Throughout this paper, we will refer to the half-life of the first moment as 1 and the half-life of the second moment as /2, analogously to the more conventional parameterization B1 and 12..
Adafactor [27]. The second moment estimate ey used in Adam requires storing a number of floating point values equal to the parameter count of the model itself. Instead, Adafactor only stores the per-row and per-column sums of the moving average for the neural network weight matrices and estimates the per-perrmaaa par-matreeer second moment using these sums. Adafactor also does not store a first moment estimate at all. Therefore, for a weight matrix with shape d1 x d2, the memory requirements for the moving averages are reduced from dj x d2 for Adam to d1 + d2 for Adafactor. We demonstrate in Section 4.5 that thanks to this sublinear memory cost, Adafactor achieves the best performance-memete trade-off.
Muon [9]. Another optimizer that was recently demonstrated to be competitive for language model training MomentUm Orthogonalized by Newton-Schula (Muon). Muon is specifically designed for 2 2-dimensional parameter tensors, such as weight matrices in linear layers, where the Newton- Schulz iterative method is applied to the first moment estimate my before using it to update the parameters. Muon is often used in conjunction with other optimizers like Adam. In our experiments using Muon, we only apply hidden layers and use Adam for input embedding and final layers.
Gradient accumulation. When hardware memory constraints prevent us from using large batch sizes directly, we can simulate large batch size by summing gradients across multiple successive micro-batches and only periodically performing an optimizer step. Under gradient accumulation, the effective batch size is equal to the micro- batch size multiplied by the number of gradient accumulation steps. While gradient accumulation allows for larger effective batch sizes, it still increases memory usage compared to using a small batch size, because requires storing the accumulated gradients.
3 Related Work
Batch size has long been studied in both theory and practice: its impact on the optimization path
[28, 29, 30, 31], its interaction with different optimizers [32, 33, 34, 35], and its effect on the optimal learning rates and momentum hyperparameters [36, 37, 38, 39, 34, 40, 41, 42, 41, 43, 44]. Prior work has also explored related themes, such as the flatness argument favoring SGD [45, 46, 47, 48), SGD's implicit biases [49, 50, 51, 52]. and the comparison between stochastic and full-batch optimization [53, 54, 55, 56]. However, modern language model training is neither convex nor characterized by multiple epoch training like vision settings. In fact, we often train for a limited token budget and stop far before convergence, since otherwise scaling laws suggest that we could train much better language model for the same compute budget by training a larger model on fewer tokens [22]. Hence, prior intuitions, such as the need for larger batch size in the convergence phase, may not apply. For this reason, in the rest of this section, we cover works that focus on modern language model training specifically.
Zhang et al. [41] and Shallue et al. [34] examine the critical batch size, which represents the threshold beyond which greater data parallelism would lead to diminishing returns, and tune all optimizer hyperparameters. They argue that any batch size below the critical batch size should perform equally well when training for the same number of tokens. Differently from these works, we empirically demonstrate that small batch sizes, including the extreme value of batch size one, perform on par with or even better than larger batch sizes, and we also show that small batches are more robust to optimizer and hyperparameter choices and enable SGD to perform competitively. We validate these findings on language models with modern design choices and up to 1.3 billion parameters.
On the other hand, Vyas et al. [57] argue that gradient noise from small batch sizes does not yield similar performance advantages in the single epoch setting as it does in multiple-ppoch training. Experimental results from Filatov al. [58] also suggest that small batches do not work as well as large batches for language models trained on C4 [59], but notably, the authors do not tune the decay rates 31 and 32 of the Adam optimizer. In the same vein, Xiao [14] claim that both excessively small and large batch sizes underperform for language model training with Adam, leading to U-shaped curve of the loss as a function of the batch size. We reproduce the exact experiments performed by Xiao [14] and show that their conclusions are an artifact of incomplete hyperparameter tuning;

---

find that simply using slower decay rate of the second moment estimate B2 makes small batch sizes competitive.
Zhao et al. [19] train language models on C4 using different diagonal preconditioning optimizers, namely Adam, Lion, Adafactor with momentum, SGD with momentum, and signed SGD with momentum. They demonstrate that all these optimizers, aside from SGD which lags behind, achieve comparable optimal performance and robustness to hyperparameterss. We challenge this conclusion in our work and argue that in the small batch size regime, SGD can perform on par with Adam.
Shallue et al. [34], Kidambi et al. [40] and Zhang et al. [41] find that momentum is not necessary for SGD when the batch size is small, which aligns with our findings. However, Kidambi et al. [40] focus on small-scale vision experiments, Shallue et al. [34] and Kidambi et al. [40] do not compare
the performance of SGD vs. Adam, and Zhang et al. [41] claim that SGD performs significantly
worse than Adam. We hypothesize this observation could be a result of Zhang et al. [41] using 64 as the smallest batch size we test batch sizes all the way down to 1. Zhang et al. [35] highlight that the benefits of more sophisticated optimizers diminish as batch size decreases, but their theoretical results assume a convex quadratic objective and their empirical results only cover a single language modeling experiment that is limited a two-layer transformer.
Porian et al. [42] and Zhang et al. [41] find that increasing the value of 32 is important for small-batch training with Adam. However, their approach is based on discrete hyperparameter sweeps across 32. In contrast, we introduce and validate a principled scaling rule that holds the second -omentt half-life constant in terms of tokens. Bushridge et al. [60] propose scaling rule for the decay coefficient of exponential moving averages similar to ours, but applied in the context of model weight averaging. Zhang et al. [44] show that there exists threshold of (32 such that when 32 185. Adam converges (this threshold decreases with batch size, i.c. the threshold is higher for smaller batch sizes). Chowdhery al. [43] note that increasing 82 during training leads better performance for rare token embeddings. While this is an important insight, we provide a simpler batch-size-aware heuristic that applies even when using a constant 3g throughout training.
4 Experimental Results
Recent works find that larger batch sizes and relatively sophisticated optimizers are necessary for stable language model training. For instance, Zhao et al. [19], Kunstner et al. [61], and Zhang et al. [62] observe that SGD has an extremely slow convergence speed for training language models in comparison to adaptive optimizers like Adam. Other works find that Adam performs poorly with small batch sizes [57, 14, 63, 64].
In this section, we revisit and challenge these beliefs by noting that such findings arise due to poor hyperparameter choices and specific experimental setups, for example, not adjusting the decay parameters 31 and 32 for Adam at small batch sizes, or using too large a batch size for SGD. We put these beliefs to the test by running thorough grid search over optimizer hyperparameters and batch sizes for SGD, Adam, Adafactor, and Muon. As the cost of a grid search grows exponentially with the number of hyperparameters, we only perform exhaustive grid searches on small transformer decoder-only model with 30 million active parameters following Liu et al. [65] and Xiao [14]. We use these results to propose heuristics for scaling hyperparameters across batch sizes, and we validate findings at a larger scale by pretraining GPT-2 (1 (124M) [66] and GPT-3 (1.3B) [17] and fine-tuning Gemma (4B) [67]. In our experiments, we study batch sizes spanning more than four orders of magnitude: {1, 4, 16, 64, 256, 1024, and 4096).
4.1 Small Batch Sizes Render Sophisticated Optimizers Unnecessary
We train the 30M model 600 million tokens from the FineWeb-Edu dataset [68], following Chinchilla scaling laws [22]. Our model adopts modern design choice, including rotary embeddings [69], QK- nK-noaaalization [70], RMSNorm [71], GELU activations [72], and separate input and output embeddings. We use a context length of 512 and tokenize the dataset using the GPT-2 tokenizer. We provide additional experimental details in Appendix A.
To verify the effect of the batch size and optimizer choice on the performance of our model, we run a dense grid search over all the hyperparameters of SGD, Adam, Adafactor, and Muon. In particular,
5

---

we tume the learning rate and the decay factors (31, 52) jointly for each batch size to obtain the hyperparameter configuration that provides the lowest validation loss.
The results in Figure 1a highlight that all optimizers perform best at the smallest batch size, and as the batch size increases, not only does the performance of each optimizer degrade but also the gap between different optimizers widens.
Why do large batches require more sophisticated optimizers? Every time our optimizer takes step, makes a prediction about the loss function away from the current parameter vector. When train with higher learning rate and take larger steps, for example in large batch training, optimizer must make predictions about the loss function farther away from the current parameter vector. We hypothesize that taking larger steps leads to harder prediction problem and that this harder prediction problem requires a more sophisticated or better-tuned optimizer.
Why is momentum less necessary for small batch sizes? Momentum can be seen as means of dampening oscillations that occur on ill-conditioned loss functions [73]. When we take large steps, may overshoot the minimum across short (i.e. high curvature) axes, leading to oscillations. By averaging gradients across iterations, these oscillations cancel out the momentum term. However, when the optimizer takes small steps, as is the case in small batch training paired with a small learning rate, parameter updates do not overshoot the minimum across short axes, so there are no oscillations to damp. We illustrate this effect in Appendix B through toy example.
Summary: The difference in performance between optimizers shrinks under small batch sizes, In fact, vanilla SGD without momentum performs competitively for small batch sizes.
4.2 Large Batch Training is Sensitive to Hyperparameters
To evaluate robustness to hyperparameters each batch size, we perform an ablation around the optimal hyperparameters achieving the lowest validation loss for the Adam optimizer. We use the same 30M parameter model trained on the FineWeb-Edu dataset as in Section 4.1.
We report in Figure 3 the change in the loss with respect to the lowest value of the loss achieved every batch size, ranging from to 4096. On the x-axis, we report the scaled value of the hyperparameters relative to the value of the hyperparameter that achieved the best loss. We perform this rescaling to make the plots for different batch sizes easier to compare, since the scale of their optimal hyperparameters varies.
1024
0.3
Change in 1065
256
0.2
16
0.1
0.0
2
2-2
2-1
29
21
22
2
2
29
2.
2
2
20
27
Scaled learning
Scaled Eg
Figure 3: Small batch sizes are robust to hyperparameter misspecification. Loss sensitivity to learning rate, 31, and 82 using Adam at batch sizes from 1 to 4096. Each curve sweeps single hyperparameter while holding the others fixed at their optimal values. We observe that smaller batch sizes exhibit broader low-loss regions across all hyperparameters, indicating greater robustness to misspecification. The x-axis shows a scaling of each hyperparameter relative to its optimum value. We parameterize 81 and 83 in terms of their decay half-lives, described in Section 2. Results shown for the 30M parameter model trained on 600M tokens of Fine Web-Edu.
We observe a striking result: small batch sizes are significantly more robust to all of Adam's hyperparameters than large batch sizes. More importantly, batch size one achieves nearly optimal loss for the entire range of learning rates, 31, and 32 hyperparameters that we consider in our search, whereas the loss for larger batches increases sharply vary the hyperparameter values.

---

To highlight the compounding value of this robustness displayed by small batch sizes as the number of hyperparameters increases, we perform joint 2-dimensional ablation with respect both the learning rate and decay rate 31 on GPT-2 model in Figure lb. Consistent with the previous result, batch size one yields much broader low-loss regions across both hyperparameters compared batch size 512. Therefore, if practitioners were to consider not only the budget to run single training run but also the computational budget for tuning the corresponding optimizer hyperparameters, a smaller batch size might be preferable.
Why are smaller batch sizes more robust to hyperparameter choices? Similar to the previous subsection, we hypothesize that large batch sizes require careful hyperparameter tuning since they require large step sizes and therefore have make predictions about the loss surface further away from the current parameter vector.
Summary: Small batch sizes are more robust to hyperparameter misspecification and might be preferable after accounting for the hyperparameter tuning budget.
4.3 Adam Hyperparameters: How to Scale Them with Batch Size and Why That Is Necessary
We use our grid search results on the 30M model to test existing scaling laws for Adam hyperparame- ters and also devise new scaling heuristics for these hyperparameters. Figure shows the validation loss on FineWeb-Bao for different Adam hyperparameters, namely the learning rate, and decay parameters 81 and 82. We similarly plot the loss for different values of the second moment half-life tg. which refers to the number of tokens it takes for gradient's contribution to the momentum to be reduced half of its initial value. We provide a more detailed description of to in Section 2.
For the learning rate, we observe that the square root scaling recommended by several works and commonly used in practice does not hold [74, 7, 37, 75, 76]. In fact, Figure 4 (left) shows that the learning rate scales more slowly than a square root factor with the batch size. For instance, as we scale the batch size from to 1024, the square root rule would indicate that the corresponding learning rate should be scaled by factor of 32, whereas we find that a factor of only about 3 empirically works better.
0.99993
0.15
rate
0.01
0.98
0.9999
t2 10M
100M
0.12
0,36
0.999
0.09
0.92
G 10M
0.001
0.84
0.99
1M
0.06
0.7
-0.03
0.0003
6Λ 205
1996
6
1024
2096
0
y6
256
1024
2096
0.1M
256
1024
0.00
size
Batch
Figure 4: Fixing the half-life of the second moment estimate in terms of tokens 12 scales better than fixing 3z. We plot the validation loss FineWeb-Edu for the 30M parameter model when varying the Adam learning rate, 31. B2. and the second moment half-life t2, across different batch sizes. The learning rate does not follow the commonly recommended square root scaling with batch size. We also observe that the default 31 = 0.9 performs well across batch sizes. In contrast, keeping 32 fixed leads to suboptimal performance at small batch sizes. Instead, fixing the second moment half-life to across batch sizes provides a simple and effective scaling rule.
We also validate that the default value of the first moment decay 81 = 0.9 used by practitioners for language model training does indeed achieve good performance across batch sizes. In contrast, keeping 33 fixed to its default value, typically in the range [0.95, 0.98] for language models, does not scale effectively to smaller batches, which require a significantly larger value of B2. On the other hand, measuring the decay of the second moment in tokens through the half-life 12 offers strong scaling heuristic: keeping l2 fixed its optimal value as we scale up or down the batch size provides good performance without re-tuning this hyperparameter at new batch sizes. This finding translates to the following scaling heuristic: if we were to scale the batch size from B to B" , the new 83 would depend on the old 3 associated with B as follows:
7

---

32 22 B BB)
Next, we test our scaling heuristics in different setting proposed in Xiao [14]. The authors train a decoder-only transformer with 19M non-embedding parameters using Adam on the C4 dataset [59], and only tune the learning rate with 81 and {8 fixed to their default values of 0.9 and 0.95, respectively. We reproduce their Figure 10(a) results in Figure 5 (left), where small batch sizes of 16 and 32 perform worse than larger batch sizes of 128 and 256.
0095
1-000
1024
3.9
loss
256
3.7
128
64
3.5
16
10
10
10 10
10
10
10-
10
Learning
Learning rate
Figure 5: Scaling 32 according to our heuristics significantly improves the performance of small batch sizes. Left: Results from Xiao [14], showing suboptimal performance for small batch sizes due to fixed Adam hyperparameters. Right: Our results after scaling 82 such that the decay half-life to is constant across batch sizes. The smallest batch sizes now perform on par with larger batch sizes.
While our scaling heuristics imply that 31 = 0.9 could work well across batch sizes, we recommend scaling 82 such that l2 is fixed. By following this simple one-shot scaling recommendation and without any additional hyperparameter tuning, we achieve significantly different results in Figure (right). Strikingly, small batch sizes 16 and 32 achieve on par performance with larger batch sizes 128 and 256, therefore implying that: (i) the previous results from Xiao [14] are largely a consequence of incomplete hyperparameter tuning, (ii) our simple one-shot heuristic generalizes to other settings.
Summary: While keeping 31 fixed to the default value of 0.9 scales well with the batch size, 32 should be scaled such that the token half-life is fixed. With this prescription, small batch sizes can perform on par with or better than larger batch sizes.
4.4 SGD Performs Competitively with Adaptive Optimizers the Billion+ Parameter Scale
Through our experiments with 30M parameter models, we conclude that: (I) small batch sizes can perform on par with better than larger batch sizes, (2) small batch sizes require less sophisticated optimizers and are more robust to hyperparameter misspecificatioon (3) in particular, vanilla SGD can be competitive with adaptive optimizers the small batch regime, and (4) we propose a new scaling heuristic for 32 such that the half-life to is fixed, and validate that it leads to improved performance. In Figure 6, we test these findings larger scale by training GPT-2 (124M) [66] and GPT-3 (1.3B) [17] on the Fine Web dataset [68].
for each optimizer was computationally prohibitive; therefore, we used AdamW with the default hyperparameter settings from Brown al. [17] (batch size 512, 31 0.9, B2 = 0.95, w 0.1) as our baseline. When using batch size 1 with Adam, we do not perform any hyperparameter tuning we turn off weight decay and scale the learning rate following our empirical results from Section 4.3. We compare two scaling heuristics for 32: keeping 3g fixed and keeping the half-life t2 fixed. We do, however, sweep over four learning rates for Adafactor and SGD, since we have no reference values for these optimizers. Interestingly, we find that SGD performs on par with the AdamW baseline, and Adam and Adafactor with batch size both outperform the W baseline.
For the larger 1.3B model, hyperparameter tuning of the learning
For the smaller (1 (124M) model, we follow the same procedure, except we tune the learning rate of each optimizer, which improves the performance of the AdamW baseline.
8

---

GPT-3 (1.38)
GPT-2 (124M)
4.0
4.4
3.6
Validation
4.2
3.4
4.0
3.2
3.8
3.6
2.8
3.4
10B
2B
Training tokens
Training tokens
Adam
AdamW
Adafactor 85-1
Adam B5-1
SGD BS
BS-512
8, - 0.95
Pz - 0.95
B2 0.9999
Bg 0.9999
Figure 6: Vanilla SGD performs competitively at larger model scales with minimal tuning. We train GPT-3 (1.3B) and GPT-2 (124M) on FineWeb and compare validation loss across optimizers. Left: For GPT-3 (1.3B). vanilla SGD with batch size and no momentum performs on par with the Adam W configuration from Brown et al. [17], despite having no optimizer state, Our proposed Se scaling rule improves Adam's performance small batch sizes and outperforms the baseline. Right: After tuning the learning rate of each optimizer on GPT-2 (124M), Adam and Adafactor at batch size 1 match the performance of AdamW batch size 512, while SGD performs slightly worse.
4.5 Full Fine-tuning on a Budget
In Sections 4.1 to 4.4, we studied pretraining across different model scales and discovered that with small batch sizes, optimizers with small or even state (like Adafactor and SGD) can perform competitively with Adam. We now ex- ploit this observation for memory-efficien fine-
Other
Accum, grads.
2 moment
Memory (GB)
1 moment
60
Activations
tuning.
Model
40
While pretraining is typically bottlenecked by compute, fine-tuning is instead often bottle- necked by memory. Indeed, fine-tuning large model on a small dataset does not require many FLOPs, but it does require storing the (large) model and optimizer state in memory. In Fig- ure 7, we show an example where the GPU has enough memory to store the model, but not enough memory to also store an expensive op- timizer like Adam.
Adam (acc.)
Adam
Muon
Adafactor
SGD
Figure 7: Minimum memory requirements to train GPT-3 (13B), assuming all parameters and activations stored in 16-bit precision, gradi- ent checkpointing is used after every transformer block layer, and the backward pass is fused with the optimizer update step (so that the full model gradient never needs to be materialized in mem- ory). The dashed line shows the available memory NVIDIA A A100 40GB GPU, which is not large enough to store expensive optimizer states accumulated gradients. We measured the true memory usage of cach optimizer by performing fused backward passes in Py RyTccch on NVIDIA B200 GPU.
common approach to reduce memory usage during fine-tuning is to freeze the model weights and only train Low-Rank Adaptation (LoRA) [77] modules on top of the frozen model, which substantially reduces the number of trainable parameters and also proportionately reduces the size of the optimizer state. However, LoRA often underperforms full parameter fine-tuning (78, 79].
We propose to use small batch size combined with small optimizer like Adafactor get the performance benefits of full fine-tuning while having similar memory requirements to LoRA.

---

We provide an example in Figure 8 by fine- tuning non-instruction-tunee
Gemma (4B) on the MATH dataset for 5 epochs. As base- line, we fine-tune the model with Adam and batch size 16 (requiring gradient accumulation), storing all parameters in float32 precision, which requires storing total of 16 bytes in memory per model parameter. Another baseline is pro- vided by LoRA, which requires only around 2 bytes per model parameter since the pretrained model checkpoint is stored in bfloat16 preci- sion, and the number of trainable parameters for LoRA is negligible compared the model size.
19%
MATH
17%
LoRA Adafactor BS=1
Adam BS=1
Adam BS=16
B5=1
Figure 8: Adafactor enables memory-efficient fine-tuning. We fine-tune and evaluate a non- instruction-tunee Gemma (4B) model [67] on the MATH dataset [80]. We compare the use of LoRA (r 16) against full model fine-tuning, and small vs. large batch size. For each training method, the width of the plotted bar is proportional to the memory footprint. LoRA uses bfloat 6 model weights and float32 adaptors, and Adafactor similarly uses bfloat16 weights but float32 opti- mizer state. We fix the second moment half-life to be consistent across batch sizes (32 0.95 for batch size 16 and 32 = 0.997 for batch size 1). We each optimizer for 5 epochs and sweep over learning rates. We train LoRA using Adam.
We make two observations based on our results. First, when using Adam, we reduce the batch size all the way down to one while follow- ing the 82 scaling heuristic from Equation (2), eliminating the need to store accumulated gra- dients in memory. Second, when small batch size is used, we can replace Adam with Adafac- tor, which has much smaller state size. To further reduce the memory use of Adafactor, we store its (tiny) optimizer state in float32, but store model weights in bfloat using stochas- tic rounding update the weights after every optimizer step (we discuss the importance of stochastic rounding Appendix A.3).
We recommend that practitioners training large models in memory constri-ndd settings exploit the benefits of small batch sizes (e.g. using optimizers with a small state size) rather than trying to emulate the large batch size setting typically used in industry.
Summary:
Practical Recommendations and Discussion
We show that small batch sizes, even low as 1, consistently match outperform large batches across model scales when tuned properly, both pretraining and fine-tuning. particular, SGD with batch size performs on par with the AdamW baseline for GPT-3 (1.3B) with minimal tuning. We also propose a simple scaling rule for Bg based on fixed second-momen half-life in tokens, which generalizes across model scales. Our results contradict common assumptions about large batch training and gradient accumulation.
Our general recommendation for setting the batch size is to use the smallest batch size that maximizes model throughput (measured in tokens / second) or equivalently the model FLOPs utilization (MFU) [43]. To achieve this objective, the accelerator must spend more time performing computation than moving data between HBM and register memory, which in practice means using batch size of at least several hundred tokens per device [81]. When training frontier models on large clusters, this can correspond to batch size many millions of tokens. But when training on a single device, the smallest batch size required to achieve high arithmetic intensity could be as low as a few thousand tokens. We therefore recommend avoiding gradient accumulation unless training on multiple devices with multiple model replicas, where the bandwidth between model replicas is bottleneck. Lastly, when working in a memory-consttrined setting, we recommend exploiting the benefits of simple optimizers combined with small batch size.
Several open questions remain. How do these findings interact with batch size schedules? Can we further reduce weight precision in fine-tuning from 16 to 8 or even 4 bits? Can we design more effective optimizers with a small state, tuned the small batch size regime? And can we theoretically understand why fixing the sccond-moment half-life in tokens generalizes across scales?
10

---

Acknowledgements
We thank Shikai Qiu and Alexandra Souly for helpful discussions. This work was supported by NSF CAREER IIS-2145492, NSF CDS&E-MSS 2134216, NSF HDR-2118810, BigHat Biosciences, and Google's TPU Research Cloud (TRC) program:
research. /og.g/seeeeee.rrr...ooooo//////
References
[1] Nidham Gazagnadou, Robert Gower, and Joseph Salmon. Optimal mini-batch and step sizes for saga. In International conference machine learning, pages 2142-2150. PMLR, 2019.
[2] Atsushi Nitanda, Tomoya Murata, and Taiji Suzuki. Sharp characterization of optimal minibatch size for stochastic finite sum convex optimization. Knowledge and Information Systems, 63(9): 2513-2539, 2021.
[3] Sungkyung Kim, Adam Lee, Junyoung Park, Sounho Chung, Jusang Oh, and Yoon Lee. Paramettr-effficient fine-tuning of instructblip for visual reasoning tasks. In Effi- cient Natural Language and Speech Processing Workshop at NeurlPS 2023, 2023. URL
ittps://neurrps22223- enlsp. github. 1o wwwppppppseppppr.eeeeee pdf.
[4] Deepanway Ghosal, Yew Ken Chia, Navonil Majumder, and Soujanya Poria. Flacuna: Unleash- ing the problem solving power of vicuna using flan fine-tuning. arXiv preprint arXiv:2307.02053, 2023.
[5] Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arXiv: 1708. 03888, 2017.
[6] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference Machine Learning, pages 1842-1850. PMLR. 2018.
[7] Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint arXiv:1904.00900 00962, 2019.
[8] Hong Liu, Zhiyuan Li, David Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. arXiv preprint arXiv:2305. 14342, 2023.
[9] Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng Yo, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https://kellerjordann giu........ ww/posts/mmmttieeeeeeseee
[10] Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. Soap: Improving and stabilizing shampoo using adam. arXiv preprint 21X.::24099.12//.1 2024.
[11] Aohan Zeng, Xiao Liu, Zhengxiao Du Zihan Wang. Hanyu Lai. Ming Ding, Zhuoyi Yang, Yifan Xu, Wendi Zheng, Xiao Xia, et al. Glm-1306 An open bilingual pre-trained model. In Proceedings of the International Conference on Learning Representations (ICLR), 2023.
[12] Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick. Suzana Ilié, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, et al. Bloom: A 176b-parameter open-access multilingual language model. arXiv preprint arXiv:221 1.05 100, 2022.
[13] Peng Xu, Dhruv Kumar, Wei Yang, Wenjie Zi, Keyi Tang, Chenyang Huang, Jackie Chi Kit. Cheung, Simon J.D. Prince, and Yanshuai Cao. Optimizing deeper transformers on small datasets. arXiv preprint arXiv:2012. /5355 2020.
[14] Lechao Xiao. Rethinking conventional wisdom in machine learning: From generalization to scaling, arXiv preprint 2411:24099.11566, 2024.

---

[15] Maxime Labonne. Fine-tune llama 3.1 ultra-efficindly with unsloth. https://hugginggacc.
ht/blog/mmaboome/fftt--tii 11ama3, 2024. Accessed: 2025-05-15.
[16] Weights & Biases. Comprehensive guide to Ilm pre-training steps. https:// /
articles/traininn www.mm/preetrainiee stepss/ 2025. Accessed: 2025-05-15.
[17] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[18] Andrej Karpathy, nanoGPT: The simplest, fastest repository for training/finetuning medium- sized gpts., 2023. URL https::/ http://gttttbbbt htw.tttu..oooomeeeeeeemmeeeeeeet www//oaaaaattt//laa /anoGGTT.
[19] Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, and Sham M Kakade. Deconstructing what makes a good optimizer for autoregressive language models. In The Thirteenth International Conference Learning Representations, 2025.
[20] Scott Sievert and Shrey Shah. Improving the convergence of sgd through adaptive batch sizes. arXiv preprint arXiv: 191.00000 2019.
[21] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08366 2020.
[22] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbi, Aidan Clark, al. Training compute-optimal large language models. Advances in neural information processing systems, 2022.
[23] Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pages 400-407, 000-407, 1951.
[24] Léon Bottou. Large-scale machine learning with stochastic gradient descent. In Proceedings of COMPSTAT 19th International Conference on Computational StatisticsParis France, August 22-27, 2010 Keynote, Invited and Contributed Papers, pages 177-186. Springer, 2010.
[25] Diederik Kingma. Adam:
method for stochastic optimization. arXiv preprint
arXiv: 141.6904 2014.
[26] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13997 2023.
[27] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pages 4596-4604. PMLR, 2018.
[28] Stanislaw Jastrzębski, Zachary Kenton, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. On the relation between the sharpest directions of dnn loss and the sgd step length. arXiv preprint 180700000... 2018.
[29] Stanislaw Jastrzebski, Maciej Szymczak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras. The break-even point on optimization trajectories of deep neural networks. arXiv preprint 20Xiv:2002.09572, 2020.
[30] Kiwon Lee, Andrew Cheng, Elliot Paquette, and Courtney Paquette. Trajectory of mini-batch momentum: Batch size saturation and convergence in high dimensions. Advances in Neural Information Processing Systems, 35:36944 36957, 2022.
[31] Alexander Atanasov, Alexandru Meterez, James Simon, and Cengiz Pehlevan. The optimiza- tion landscape of sgd across the feature learning strength. arXiv preprint arXiv:2410.04642, 2024.
[32] Randall Wilson and Tony R Martinez. The general inefficiency of batch training for gradient descent learning. Neural networks, 16(10):1429-1451, 2003.
12

---

[33] Jianmin Chen, Xinghao Pan, Rajat Monga, Samy Bengio, and Rafal Jozefowicz. Revisiting distributed synchronous sgd. arXiv preprint arXiv: 1604.066040008/ 2016.
[34] Christopher Shallue, Jachoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E Dahl. Measuring the effects of data parallelism on neural network training. Journal of Machine Learning Research, 20(1 12):1 2019.
[35] Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger B Grosse. Which algorithmic choices matter which batch sizes? insight from noisy quadratic model. Advances in neural information processing systems, 32, 2019,
[36] Samuel L Smith and Quoc V Le. A bayesian perspective generalization and stochastic gradient descent. arXiv preprini arXiv:
2017.
[37] Diego Granziol, Stefan Zohren, and Stephen Roberts. Learning rates as a function of batch size: random matrix theory approach to neural network training. Journal of Machine Learning Research, 23(173): 2022.
[38] Runzhe Wang. Sadhika Malladi, Tianhao Wang, Kaifeng Lyu, and Zhiyuan Li. The marginal value of momentum for small learning rate sgd. arXiv preprint arXiv:2307.1511 15166, 2023.
[39] Jingwen Fu, Bohan Wang, Huishuai Zhang, Zhizheng Zhang, Wei Chen, and Nanning Zheng. When and why momentum accelerates sgd: An empirical study. arXiv preprint arXiv:2306.0900 2023.
[40] Rahul Kidambi, Praneeth Netrapalli, Prateek Jain, and Sham Kakade. On the insufficiency of existing momentum schemes for stochastic optimization. In 2018 Information Theory and Applications Workshop (ITA), pages 1-9, 2018. doi: 10. 1.09/TTA.211..501177..
[41] Hanlin Zhang, Depen Morwani, Nikhil Vyas, Jingfeng Wu, Difan Zou, Udaya Ghai, Dean Foster, and Sham Kakade. How does critical batch size scale in pre-training? arXiv preprint arXiv:2410.21666, 2024.
[42] Tomer Porian, Mitchell Wortsman, Jenia Jitsev, Ludwig Schmidt, and Yair Carmon. Resolving discrepancies in compute-optimal scaling language models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https: http//ppeeeeeeeeeeeeeeeeeeeeeeee net/fforuaiic==4f$$k11e..
[43] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240): 1-113, 2023.
[44] Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi thi-Quan Luo. Adam can converge without any modification on update rules. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and Oh, editors, Advances in Neural Infor- mation Processing Systems, volume 35, pages 28399, Curran Associates, Inc., 2022. URL https: htttps://proceedinggg (6250a55644221005eeaaaa6l93 -arpraaaaaaaaattooaaaaeaeee--rrr Coaper-Coofeeennne pdf.
meurips.cc/paper wtt.fffeeeeeeeeaeeeeeeee /appr//0222/fill/
[45] Yann LeCun, Léon Bottou, Genevieve B Orr, and Klaus-Robert Müller. Efficient backprop. In Neural networks: Tricks of the trade, pages 9-50. Springer, 2002.
[46] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint 15fi::66000048666 2016.
[47] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in I hour. arXiv preprint arXiv: 100iv:700.0000 02677, 2017.
[48] Dominic Masters and Carlo I Luschi. Revisiting small batch training for deep neural networks. arXiv preprint 1xXiv:1804.06112 2018.
13

---

[49] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization, Advances in neural information processing systems, 32, 2019.
[50] David GT Barrett and Benoit Dherin. Implicit gradient regularization. arXiv preprint arXiv:2009.11162 2020.
[51] Jeff Z HaoChen, Colin Wei, Jason Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. In Conference on Learning Theory, pages 2315-2357. PMLR, 2021.
[52] Samuel L Smith, Benoit Dherin, David GT Barrett, and Soham De. On the origin of implicit regularization in stochastic gradient descent. arXiv preprint vrXiv:2101.12170 2021.
[53] Rong Ge, Furong Huang, Chi Jin, and Yang Yuan. Escaping from saddle points online stochastic gradient for tensor decomposition. In Conference on learning theory, pages 797-842. PMLR, 2015.
[54] Simon S Du, Chi Jin, Jason D Lee, Michael Jordan, Aarti Singh, and Barnabas Poczos. Gradient descent can take exponential time escape saddle points. Advances in neural information processing systems, 30, 2017.
[55] Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.0m 022 18, 2020.
[56] Jonas Geiping, Micah Goldblum, Phillip E Pope, Michael Moeller, and Tom Goldstein. Stochas- tic training not necessary for generalization. arXiv preprint arXiv:2109. 191199 2021.
[57] Nikhil Vyas, Depen Morwani, Rosie Zhao, Gal Kaplun, Sham Kakade, and Boaz Barak. Beyond implicit bias: The insignificance of sgd noise in online learning. arXiv preprint arXiv:2306.0m 08590, 2023.
[58] Oleg Filatov, Jan Ebert, Jiangtao Wang, and Stefan Kesselheim. Time transfer: On optimal learning rate and batch size in the infinite data limit, arXiv preprint arXiv:2410.05838, 2024.
[59] Jesse Dodge, Maarten Sap, Ana Marasović, William Agnew, Gabriel Ilharco, Dirk Groeneveld, Margaret Mitchell, and Matt Gardner. Documenting large webtext corpora: A case study on the colossal clean crawled corpus. arXiv preprint arXiv:2/044 08758, 2021.
[60] Dan Busbridge, Jason Ramapuram, Pierre Ablin, Tatiana Likhomanenko, Eeshan Gunesh Dhekane, Xavier Suau, and Russell Webb. How to scale your EMA. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL bettps://operreeee wet.ffrrmmiddDxeeeeeeee
[61] Frederik Kunstner, Alan Milligan, Robin Yadav, Mark Schmidt, and Alberto Bietti. Heavy- tailed class imbalance and why adam outperforms gradient descent on language models. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang, editors, Advances in Neural Information Processing Systems, volume 37, pages 30106-30144 Curran Associates, Inc., 2024. URL https: http:////oooceeenneeeereee neurips.. Con.apprrr-ooeeeeeeee.ee...e
/paprr/
[62] Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhi-Quan Luo. Why transformers need adam: A hessian perspective. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https: htttts:://ppeeeeieeeeeeee X6rqEpbnj3.
met/forummidd
[63] Qi Liu and Wanjing Ma. The epochal sawtooth effect: Unveiling training loss oscillations in adam and other optimizers. arXiv preprint arXiv:2410. 10056, 2024.
[64] Shuaipeng Li, Penghao Zhao, Hailin Zhang, Xingwu Sun, Hao Wu, Dian Jiao, Weiyan Wang, Chengjun Liu, Zheng Fang, Jinhao Xue, Yangyu Tao, Bin Cui, and Di Wang. Surge phenomenon in optimal learning rate and batch size scaling, arXiv preprint 1405.1 2024.
14

---

[65] Peter J. Liu, Roman Novak, Jaehoon Lee, Mitchell Wortsman, Lechao Xiao, Katie Everett, Alexander A. Alemi, Mark Kurzeja, Pierre Marcenac, Izzeddin Gur, Simon Korbblith, Kelvin Xu, Gamaleldin Elsayed, lan Fischer, Jeffrey Pennington, Ben Adlam, and Jascha-Sohl Dickstein. Nanodo: minimal transformer decoder-only language model implementation JAX., 2024. URL. htt.: / titt......... com/ coom/gogle--oo delpmiidedddddd www.dddiddddooooo./ nanodo.
[66] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 18))9, 2019.
[67] Gemma Team, Aishwarya Kamath, Johan Ferret, Shreya Pathak, Nino Vieillard, Ramona Merhej, Sarah Perrin, Tatiana Matejovicova, Alexandre Ramé, Morgane Rivière, et al. Gemma 3 technical report. arXiv preprint arXiv:2503. 19786, 2025.
[68] Guilherme Penedo, Hynek Kydliček, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, Thomas Wolf, et al. The fineweb datasets: Decanting the web for the finest text data at scale. Advances in Neural Information Processing Systems, 37:3081 1-30849, 2024.
[69] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 16861,, 2024.
[70] Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. In International Conference on Machine Learning, pages 7480-7512. PMLR, 2023.
[71] Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
[72] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.0888 08415, 2016.
[73] Gabriel Goh. Why momentum really works. Distill, 2017. doi: 10.23915/dissil.00000 URL http://distill, http:///ddstilllppp//ooo/// pwbpppppb/1/1/ haabb2217/momettum /momentum.
[74] Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team, An empirical model of large-batch training. arXiv preprint 1rXiv:1812.06162. 2018.
[75] Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the sdes and scaling rules for adaptive gradient algorithms. Advances in Neural Information Processing Systems, 35: 7697-7711, 2022.
[76] Xi Wang and Laurence Aitchison. Batch size invariant adam. arXiv preprint 18844 2024.
[77] Edward Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022. URL pettps://openreeeee net/forune /tt///oruunndideeeeeeeeeee.
[78] Dan Biderman, Jacob Portes, Jose Javier Gonzalez Ortiz, Mansheej Paul, Philip Greengard, Connor Jennings, Daniel King, Sam Havens, Vitaliy Chiley, Jonathan Frankle, et al. Lora learns less and forgets less. arXiv preprint 2XXx:24405006673. 2024.
[79] Shibo
Wang and Pankaj Kanwar. hakhamaneshi and rehaan ahmad, 2023. URL https: //wp::://www... anyscale com/blog/ tanin-tunngg- 1lms-lora-
Artur miederfahrenhorst and kourosh
lora-or----- or- -o--oollllllll full- mall-parameter---. an- in- depth- wwwahhasaaaaysssssaaaaaasa wettlysssssihhhh--se with-withh-- 10a--1ama-222
[80] Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang. Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. J. Vanschoren and Yeung, editors, Proceedings of the Neural Infor- mation Processing Systems Truck on Datasets and Benchmurks, volume 1, 2021. URL hlat::///attssttteecccmee------e http://gg....rriipp.c..ppppeeerr www.ff/epeppeeeeeee 200/://iieesssseeecdo/////aaaaa2 Paperr round2.paff pdf.
neurips hatpps.cc/ppppae
15

---

[81] Jacob Austin, Sholto Douglas, Roy Frostig, Anselm Levskaya, Charlie Chen, Sharad Vikram, Federico Lebron, Peter Choy. Vinay Ramasesh, Albert Webson, and Reiner Pope. How to scale your model. 2025. Retrieved from
[82] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[83] Norm Jouppi, George Kurian, Sheng Li, Peter Ma, Rahul Nagarajan, Lifeng Nai, Nishant Patil, Suvinay Subramanian, Andy Swing, Brian Towles, et al. Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings. In Proceedings of the 50th annual international symposium on computer architecture, pages 1-14, 2023.
16

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We have performed extensive experiments to back our claims. We outline objectives and follow through with corresponding methods, results, and analysis that substantiate those claims. The abstract offers a concise and truthful summary the work.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We discuss the scale limitations of our grid search as we need perform on a smaller model compared to Section 4.4 where we perform larger scale experiments but without grid search. We also clearly state open questions in the discussion section.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
17

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: The results we present are primarily experimental. Although we have looked into results from the theoretical fields, we have not built upon them.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We clearly state the models, datasets, and hyperparameters used during training in the paper. We also mention the design choices adopt, namely, rotatory embeddings, QK -K-ooaaalizttion RMSNorm. GELU activations, separate input and output embeddings, among others. This allows others to replicate the results and validate our main claims.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often good way to accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
18

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Yes, we provide open access our code, along with clear and sufficient instruc- tions to reproduce our results: https://gs /:g::://ghttbt com wwtcmoototiiieeeeeeee martin-maaaatteaaaaa batch- size.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips..
httpe//uuides/oeeeubbbiiieee selssssssiiiiyyyy for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / /ttpss.ccppullliie gewpi//g////////eeooooocoo (ecodeeee/eoeummmiiii SondeSuumissiioPoone
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: We clearly state the models, datasets, and hyperparameters used during training in the paper. We also mention the design choices we adopt, namely, rotatory embeddings, QK-normalization, RMSNorm, GELU activations, separate input and output embeddings, among others.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
19

---

Justification:
We perform a thorough grid search and report results that reflect the statistical significance of the loss the optimal hyperparameter settings.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi-
dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative егrоr rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Yes, the paper will provide sufficient information the computational resources used for the experiments incluuuin hardware type, memory, and execution time- the supplementary materials, which will be included in the appendix. This will ensure reprodueibility and support the paper's claims regarding efficiency.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: //tne///neesss/ee
http:///cerrsppp..pppiiiieeeeeee httpp/uull/ttthhseeeeeeeeeeeee csGuidellidlii ines?
Answer: [Yes]
Justification: Our work fully adheres to the NeurIPS Code of Ethics. We have carefully considered the broader societal and environmental implications of our research and ensured alignment with the conference's ethical standards. Additionally, we affirm compliance with the guidelines on research integrity, plagiarism, fraud, and reproducibility as outlined in the Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
20

---

If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Yes, the paper discusses the potential positive societal impacts of reducing the resource requirements for training large language models, such as making LLM research more accessible, and enabling training not explicitly address potential negative societal impacts, such as the broader consequences of democratizing powerful language technologies without sufficient safeguards.
shource-constrainee settings. However, it does
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymentt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI...
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (c.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We do intend release new datasets or models.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
21

---

12. Licenses for existing assets
Question:
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We have appropriately cited all publicly available and open-source research used in our work. The original authors are explicitly acknowledged and given full credit for their contributions.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode. thcode. com datasets
has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors encouraged to reach out the asset's ss creators.
13. New assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: We do not intend release any models datasets. However, intend to make the code publicly available after acceptance and provide proper documentation to the public.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
Does not apply.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
22

---

According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based on the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Does not apply.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: Does not apply.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our L.M policy (https:/ cattps://neurd: wwtp/ccc///acaaaeaeeee http://p.......oeeeeeeeee/////// for what should or should not be described.
23

---

A Experimental Details
A.1 Model architecture
Across all of our pretraining experiments, we used four decoder-only transformer models [82] that share the same architecture and only differ layer dimensions and tokenizer. We specify the exact layer dimensions and tokenizers used in Table 1.
We used model architecture based GPT-2 [66] with several small changes based modern practices to improve model FLOPs utilization (MFU) [43], simplify implementation, and marginally improve model performance:
Head dimension 128
We increased the head dimension on the smaller models from 64 to 128 to achieve a higher MFU. We trained each model on TPU v4, which has 128 x 128 multiply- accumulators. Using a head dimension any lower than 128 would prevent from filling the
thereby lowering MFU. This decision follows the Gemma family of models optimized for TPU training [67].
RMSNorm [71]
We wanted to include the Muon optimizer in our benchmarks, but Muon does not support training 1D layers. To get around this limitation, simply omitted the use of any ID layers by using RMSNorm with no bias.
Untied embeddings
We do not tie the weights of the first and last layer. This increases the number of trainable parameters but does not increase the number of active parameters or the cost of training.
Rotary positional embeddings (RoPE) [69]
QK-norm [70]
GELU [72]
We pretrained most models to follow Chinchilla scaling laws, i.e. using 20 tokens per active parameter. The only exception to this rule is the GPT-3 (1.3B) model, which we only trained for 10B tokens, similar to the original GPT-2 (1.5B) model.
Model
30M
19M[14]
GPT-2 (124M)
GPT-3 (1.3B)
Dataset
Fineweb-Edd
Fineweb
Fineweb
Tokenizer
GPT-2
T5
GPT-2
GPT-2
Vocabulary size
50257
32101
50257
50257
Model/ embedding dimension
384
512
768
2048
Hidden dimension
1536
2048
3072
8192
Head dimension
128
128
128
128
Number of layers
6
12
24
Sequence length
512
512
1024
2048
Non-embedding parameters
11M
19M
85M
1.2B
Embedding parameters
2x 19M
2 x 16M
2 39M
2 103M
Active parameters
30M
35M
124M
1.3B
Total trainable parameters
49M
52M
162M
1.4B
Training tokens
600M
705M
2.5B
10B
Table 1: Model dimensions and configurations across different architectures.
A.2 Learning rate schedule
For pretraining, we used a linear warmup for the first 5% of the training steps (from zero to peak learning rate), followed by cosine decay (from peak learning rate to zero). We the same schedule for all batch sizes. For fine-tuning, we used a constant learning rate.

---

A.3 Stochastic rounding
Across all experiments, we compute activations and store the optimizer state in float32 precision. We also store all trainable model parameters in floar32 precision except for Adafactor in Fig- ure 8, which uses bfloat16 weights order to match the memory footprint of LoRA. While it a common practice to use 6 weights for computation, the master copy of model weights typically still stored in float32 to enable ac- cumulating small updates to the weights. When the master weights are stored in b6loat16 (which only has ~244 decimal points of precision), stan- dard deterministic rounding can bias the weight updates, meaning that accumulating gradients across many small batches no longer approxi- mates computing gradients from single larger batch. Because of this bias, reducing the batch size might actually be detrimental with stan- dard deterministic rounding. One approach to remove this bias is compute updated weights in float32, then stochastically round them to 6 for storage. In a pretraining experiment in Figure 9, we show that stochastic rounding can bring the performance of 6 weights close to the performance of float32 weights, even for the smallest batch sizes. In the fine-tuning experiment in Figure 8, we did not observe statistically significant difference between float32 weights and 6 weights with stochastic rounding for Adafactor.
5.0
BF16 (closest)
BF16 (stochastic)
FineWeb Edu loss
FP32
4.6
4.4
4.2
4.0
3.8
3.6
256
1024
Batch
Figure 9: Stochastic rounding enables bfloat16 training. We repeat the training runs from Fig- ure 1a using the optimum hyperparameters found for Adafactor. We ablate the precision of the model weights, using: (1) float32 - same as in Figure la, (2) 6 with closest rounding. (3) 6 with stochastic rounding. The optimizer state is always stored in float32 precision since has negligible size.
A.4 Computational cost
We used 32 TPU v4 chips train all models [83]. Table 2 shows the time for single training run of each model as well as the total computational cost for each model. Using batch sizes and 2 resulted in roughly 30-70% drop in MFU compared to using a larger batch size, depending on the model size and sequence length. In practice, we do not recommend using batch size so small that it severely degrades MFU. We only used the smallest batch sizes scientifically study the effect of batch size.
Model
30M
19M[14]
GPT-2 (124M)
GPT-3 (1.3B)
Gemma3 (4B)
Accelerator
TPU v4 chip
TPU v4 chip
TPP v4 chip
TPP v4 chips
4 TPU v4 chips
Average run time
24 min
25 min
7.5 hours
100 hours
44 min
Num. runs: Figure
3110
312
11
Num, runs: Figure
1911
Num, runs: Figure 4
6500
Num. runs: Figure
1260
Num. runs: Figure
40
Numn runs: Figure
256
Total num. of runs
11521
1260
352
256
Total TPU v4 chip hours
4600
525
2640
4400
750
Table 2: Computational cost of each experiment.
B Additional Results
In Figure 10 we illustrate on a toy problem why momentum is required when the batch size is large but not required when the batch size is small. We run SGD on two variables (x, y) to minimize the value of a loss function defined as x + 10y2. Notice that this loss function much steeper in the vertical direction compared to the horizontal direction.
We simulate the use of small and large batch size by adding noise the gradient estimate at every step. In particular, we define the minibatch gradient estimate at each step to be the true gradient
25

---

multiplied by a random variable sampled from a normal distribution M(1, o²), where the noise scale controls the signal-to-noise no-se ratio of the minibatch gradient estimate. We simulate large batch size by using a high signal-to-ratio of 5, and small batch size by using small signal-to-noise ratio of only 0.3.
In our experiments training language models, we fixed the computational budget across all batch sizes by only training for single epoch. Analogously in this toy example, we only run 10 optimization steps using the simulated large batch size, but we run 100 optimization steps using the small batch size. The idea that if we use small batch size, the minibatch gradient estimates at each step become more noisy, but in turn we get to take more optimizer steps.
SGD
SGD momentum
Large batch size
Small batch size
Figure 10: Toy optimization problem. We compare the of SGD and SGD with momentum loss function defined as x - 10y². We simulate a small and a large batch size by adding Gaussian noise the gradient estimate each optimizer step. The error bars show 50% interquartile range sampled across different random seeds for the minibatch gradient noise.
In the top row of Figure 10, we see that when a large batch size used, we only get to take small number of steps, and so in turn every step has to be large. Because the loss function is much steeper in the vertical direction compared to horizontal direction, the large step size results in oscillations along the vertical direction. In this case, using momentum helps dampen the oscillations along the vertical direction and speeds up convergence along the horizontal direction. fact, in the extreme case of full-batch gradient descent (which has gradient noise), using momentum probably improves the convergence rate of gradient descent for this toy problem [73].
We show the small batch size setting in the bottom row of Figure 10. Notice that since the batch size is small, every step becomes more noisy, but in turn we can take more steps and reduce the size of each step. As result of decreasing the step size, the optimizer no longer overshoots along the vertical direction, obviating the need momentum. There are no longer any oscillations to dampen, so using SGD with and without momentum results in similar performance.
26


Reviews:
[{"review_id": "J3DEGYTd73", "reviewer": "Reviewer_3pDd", "summary": "This paper revisits small-batch training for language models, targeting the prevailing belief that large batches are essential for stability and efficiency. The authors show that, with proper hyperparameter tuning, in particular adjusting the second-moment decay in Adam based on token half-life, small batch sizes, down to batch size one, can achieve stable training, robustness to hyperparameter choices, and competitive or superior per-FLOP performance compared to larger batches. Notably, vanilla SGD without momentum performs on par with sophisticated optimizers in the small-batch regime, enabling significant memory savings. The paper argues against unnecessary gradient accumulation and advocates using the smallest batch size that maximizes hardware throughput.", "questions": "- You briefly mention Muon and Adafactor, but did you consider newer second-order or quasi-second-order optimizers (e.g., Sophia, Shampoo)? How might these perform in the small-batch regime relative to vanilla SGD?\n- While small batches reduce memory and improve robustness, have you quantified the impact on wall-clock training time, given the reduced hardware throughput on modern accelerators? How should practitioners balance these trade-offs in practice?\n- Did you evaluate whether models trained with small batches differ in downstream task performance (e.g., question answering, summarization) relative to large-batch-trained models?\n- While small batches reduce memory and improve robustness, have you quantified the impact on wall-clock training time, given the reduced hardware throughput on modern accelerators? How should practitioners balance these trade-offs in practice?", "limitations": "Scale of large-model experiments, limited domain diversity, hardware efficiency trade-offs, absence of theoretical guarantees, incomplete optimizer coverage, and potential risks of underfitting.", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths:\n- Challenges widespread assumptions about the necessity of large batch sizes in LM training, contributing valuable insights to optimization practice.\n- Demonstrates that vanilla SGD without momentum can perform competitively, offering significant reductions in memory and compute requirements.\n\nWeaknesses:\n- The paper largely provides empirical observations without rigorous theoretical explanations for why small-batch regimes enable competitive performance, which could limit the depth of understanding and generalization (more on that on the two points below).\n- While results extend to 1.3B-parameter models, the largest experiments remain less exhaustive (e.g., no full hyperparameter grid search), potentially leaving open questions about scalability to the largest frontier models.\n- The primary experiments focus on FineWeb-Edu and related benchmarks; it is unclear whether findings generalize equally well to other domains or highly diverse corpora (e.g., code, multilingual data).\n-  While advocating for small batch sizes, the paper briefly acknowledges throughput constraints but does not deeply quantify the performance penalties of very small batches on modern hardware accelerators.", "was_revised": true, "final_justification": "The author did a very good effort in the rebuttal to answer all questions (thanks for that). However, and despite answer many questions, **I think two main fundamental issues still stand out (and both do not give confidence in the robustness of the proposed favorable effects of small batch sizes):**\n\n- The paper largely provides empirical observations without rigorous theoretical explanations for why small-batch regimes enable competitive performance, which could limit the depth of understanding and generalization. The authors argue that experimentation alone is still a good indicator for effectiveness, however it is widely known that models can be sensitive to particular configurations/hyperparameters, in addition that not knowing why small batch sizes work also makes it unclear if this is a general effect, or just an artifact of the particular experimentation setup used.\n\n- No exploring of the setting at which models are pretrained with different batch sizes and then evaluated using a fixed downstream fine-tuning setup. Without that, this is lack of concrete evidence that small batch size does indeed have the desired effect (after controlling for everything else).\n\n**I respect the effort the authors put in the paper, and the rebuttal, yet I will maintain my score.**", "rebuttal": "Thank you for your detailed review. We address your feedback below. \n\n**Novel results on other optimizers, datasets, and models:** Inspired by your feedback, we extend our experiments to the mC4 dataset, a multilingual variant of the C4 dataset comprised of natural text in 101 languages [1]. We also run additional fine-tuning experiments where we fine-tune and evaluate a non-instruction-tuned Gemma 3 (4B) [2] on the MATH dataset [3]. Finally, we test additional optimizers at scale. We describe our results below.\n\n**Scalability of our findings:** To demonstrate that our findings and recommendations extend to larger scale, we re-ran our grid-search experiments reported in Figure 1b using a larger GPT-2 (124M) model trained on 2.5B tokens of FineWeb, and we observe an even clearer trend of batch size 1 being more robust to hyperparameters. Here we show a binarized version of the plot where the dark regions correspond to validation loss <3.5 and the white regions correspond to validation loss >3.5. The smaller batch size is much more robust to both the momentum (y-axis) and the learning rate (x-axis) hyperparameters.\n\n             Batch size 1          Batch size 512\n          +----------------+     +----------------+\n    0.999 |████████████████|     |                |\n    0.99  |████████████████|     |                |\n    0.97  |████████████████|     |                |\n    0.9   |████████████████|     |     ███████    |\n    0.7   |████████████████|     |                |\n    0     |████████████████|     |                |\n          +----------------+     +----------------+\n          2⁻¹⁴ 2⁻¹² 2⁻¹⁰ 2⁻⁸     2⁻¹¹ 2⁻⁹  2⁻⁷  2⁻⁵\n            Learning rate          Learning rate\n\nWe also tested Adafactor at 124M and 1.3B parameter scale for pretraining on FineWeb and in both cases it performed on par or even slightly better than Adam and AdamW with batch sizes 1 and 512 as shown in the table below. \n\n|               |  AdamW, BS=512 (baseline) |  Adam, BS=1, ${\\beta_2}=0.95$ |  Adam, BS=1, ${\\beta_2}=0.9999$ |    Adafactor, BS=1 |  SGD, BS=1 | \n| ---                    | ---                  | ---                       |  ---                | ---                    | ---                  | \n| GPT-2 (124M)    | 3.43           | 4.40                |  3.38               | 3.36⭐️                   | 3.52            | \n| GPT-3 (1.3B)      | 2.95               | 3.45              |  2.79            | 2.78⭐️              | 2.97                  | \n\nWe would like to kindly note that running our grid-search experiments at the 1.3B scale would require around 2.7M TPU-v4-chip hours, at an approximate cost of $8.7M, which is not an accessible resource for us. However, several of our experiments with different models, datasets, and optimizers demonstrate the validity of our findings and recommendations at scale, including our fine-tuning experiments with the 4B-parameter Gemma 3 model. \n\n**Pretraining on the mC4 dataset:** Inspired by your recommendation to extend our experiments to a diverse multilingual dataset, we used Adafactor with batch size 1 and AdamW with batch size 512 to train GPT-2 (124M) for 2.5B tokens on the m4C multilingual dataset tokenized using the T5 tokenizer. We used the same hyperparameter values as in the table above (where the learning rate was tuned for each optimizer) and obtained a validation loss of **0.96** for Adafactor with batch size 1 and **0.99** for AdamW with batch size 512. We specifically chose Adafactor for this comparison since it is a simple optimizer with an extremely small memory footprint. We find it interesting that when a small batch size is used, such a simple optimizer performs competitively with the gold-standard AdamW, even in the multilingual setting.\n\n**Finetuning on the MATH dataset:** We fine-tune a non-instruction-tuned Gemma 3 (4B) model on the MATH dataset using various optimizers and batch sizes, allowing us to extend our experiments to a larger model size, a reasoning task, and a different regime where we train for multiple epochs. The evaluation metric is the percentage of problems correctly solved by the model. We compare full-parameter fine-tuning to Low-Rank Adaptation (LoRA), which freezes the pretrained weights and only updates small trainable low-rank matrices. LoRA is commonly used for fine-tuning because it reduces memory costs.\n\nHowever, our results reported in the table below show that we can instead address the memory constraints using a small batch size, and do full fine-tuning. Specifically, Adafactor (which has the same memory footprint as LoRA) and Adam with batch size one both outperform Adam with a larger batch size and LoRA with batch size one. This result suggests that constraining updates to a low-rank subspace, as in LoRA, is not necessary and may in fact deteriorate the performance. Full fine-tuning with small batches avoids gradient accumulation, favors memory-efficient optimizers, and achieves better results. \n\n| Adam, BS=16 | Adam, BS=1 | Adafactor, BS=1 | LoRA, BS=1 | \n| ---                    | ---                  | ---                       |  ---                |\n| 18.2 ± 0.4%    | 18.7 ± 0.4%   | 18.7 ± 0.2%        | 16.7 ± 0.4% |  \n\n**Rigor and generalizability of our work:** While our work is primarily empirical rather than theoretical, we follow a rigorous scientific approach. We clearly specify our hypotheses about small batch size training, carefully remove confounding factors such as weight decay and variable token budget, and test one hyperparameter at a time thanks to our grid-search. This rigorous approach enables us to derive insights and formulate theories that are testable and falsifiable. Additionally, our experiments span multiple model scales, optimizers, and datasets. With the newly added fine-tuning experiments, we extend our analysis to both single-epoch and multiple-epoch training regimes, therefore broadening the scope and generalizability of our findings. To provide intuition behind our findings, we included in Appendix B (added to the supplementary material) a toy problem illustrating why momentum is necessary for stable training with large batch sizes but not required when training with small batches.\n\n**Throughput and trade-offs:** Our general recommendation is to use the _smallest_ batch size that results in good hardware utilization. To saturate their matrix multiplication units, modern ML accelerators typically require an arithmetic intensity of several hundred FLOPs / byte, i.e., they need to perform at least several hundred GEMM FLOPs per one byte of data loaded from global memory. In our case, this means that every model parameter loaded from global memory needs to be used for at least several hundred tokens in a batch.\n\n**Reporting wall-clock time:** We now report wall-clock training times for different batch sizes using our codebase on a TPU v6e-1 accelerator. We find that very small batches have low arithmetic intensity, while very large batches incur overhead from gradient checkpointing. In our setup, BS=4 achieves the best throughput. To re-iterate, we recommend using the smallest batch size that achieves good hardware utilization, and this threshold varies across hardware platforms.\n\n|                    |    BS=1        |  BS=4        |  BS=16      | BS=32       | \n| ---             | ---             | ---            |  ---           |     ---        |\n| Time         |  3h:32min  |    ⭐️2h:43min |   3h:14min  |   4h:16min | \n\n\n**Considering newer second-order or quasi-second-order optimizers:** Our choice to use Muon and not other second-order or quasi-second-order optimizers is motivated by the fact that Muon (2024) is the newest and most performant optimizer for large language model training. Muon has been used in several recent state-of-the-art settings, including the NanoGPT speedrun, CIFAR-10 speedrun, and training the Kimi K2 models. In contrast, Shampoo (2018) has significant memory overhead, and Sophia (2023), while promising, underperforms compared to Muon for LLM training. Therefore, we believe that Muon is the most relevant quasi-second-order method for our work. Our choice of Adafactor as another baseline stems from its extremely low memory footprint.\n\n**Impact on downstream performance:** Our additional fine-tuning experiments show that small-batch training performs competitively in downstream tasks under memory constraints. However, we did not explore the setting where models are pretrained with different batch sizes and then evaluated using a fixed downstream fine-tuning setup. We agree this is an interesting direction, and consider it a promising avenue for future work. \n\nThank you again for your detailed feedback, we believe it has positively impacted our paper and extended the scope of our experiments. We made a significant effort to address your points and hope that you would consider raising your score. Please let us know if you have any other questions that we can answer.\n____\n\nReferences: \n\n- [1] Xue, L., Constant, N., Roberts, A., Kale, M., Al-Rfou, R., Siddhant, A., Barua, A. and Raffel, C., 2020. mT5: A massively multilingual pre-trained text-to-text transformer. arXiv preprint arXiv:2010.11934.\n- [2] Team, G., Kamath, A., Ferret, J., Pathak, S., Vieillard, N., Merhej, R., Perrin, S., Matejovicova, T., Ramé, A., Rivière, M. and Rouillard, L., 2025. Gemma 3 technical report. arXiv preprint arXiv:2503.19786.\n- [3] Hendrycks, D., Burns, C., Kadavath, S., Arora, A., Basart, S., Tang, E., Song, D. and Steinhardt, J., 2021. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.03874."}, {"review_id": "Fu155YExGo", "reviewer": "Reviewer_yNXQ", "summary": "This paper investigates the training performance of LLM in small batch sizes.", "questions": "The paper cites in [40] that a \"surge\" phenomenon has been reported in the Adam optimizer. Could this issue potentially undermine the validity of the conclusions presented in this work?", "limitations": "See Weakness.", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\nSome of the points raised in the paper are inconsistent with existing knowledge and offer new perspectives. For example, SGD can achieve performance similar to AdamW, and small batch training is more robust.\n\n**Weaknesses**\n\n1. The model used in the experiment is too small. The authors drew many of their conclusions based on a 42M model, and although these conclusions were verified using a 1.3B GPT-3 model, they were not verified using larger models. This raises concerns about the generalizability of these conclusions.\n2. While Figure 1(b) compares the robustness to other hyper-parameters at batch sizes of 1 and 256, the continuous behavioral changes from small to large batch sizes are not thoroughly discussed across the entire range.\n3. The paper appears structurally incomplete, and crucially, key appendixes are missing.", "was_revised": true, "final_justification": "The authors' newly presented experimental results and evidence have substantially alleviated my concerns regarding generalizability.", "rebuttal": "Thank you for your feedback. We address your questions and concerns below. \n\n**Additional experiments at a larger scale:** In the submitted version of our paper, we pretrained models ranging from 30M active parameters to 1.3B (corresponding to $10^{17}$ – $10^{20}$ FLOPs). We trained all models (except the largest 1.3B model) to be Chinchilla optimal. We trained the 1.3B model for 10B tokens following the GPT-2 paper as well as the modded-nanoGPT repository. Since then, we have expanded our results in two ways: we tested more optimizers at the 124M and 1.3B scale and we tuned the learning rate of each optimizer at the 124M scale. And perhaps even more importantly, we added a fine-tuning experiment with a Gemma 3 4B, which was (over)trained on 4T tokens. We show that even this \"late\" in the training stage of a model, batch size 1 performs just as well as batch 16, even when fine-tuning for multiple epochs. We report these results in more detail below and in the main response. \n\n**Finetuning on the MATH dataset:** We fine-tune a non-instruction-tuned Gemma 3 (4B) model on the MATH dataset using various optimizers and batch sizes, allowing us to extend our experiments to a larger model size, a reasoning task, and a different regime where we train for multiple epochs. The evaluation metric is the percentage of problems correctly solved by the model. We compare full-parameter fine-tuning to Low-Rank Adaptation (LoRA), which freezes the pretrained weights and only updates small trainable low-rank matrices. LoRA is commonly used for fine-tuning because it reduces memory and compute costs while maintaining reasonable performance. \n\nHowever, our results reported in the table below show that we can instead address the memory constraints using a small batch size, and do full fine-tuning. Specifically, Adafactor (which has the same memory footprint as LoRA) and Adam with batch size one both outperform Adam with a larger batch size and LoRA with batch size one. This result suggests that constraining updates to a low-rank subspace, as in LoRA, is not necessary and may in fact deteriorate the performance. Full fine-tuning with small batches avoids gradient accumulation, favors memory-efficient optimizers, and achieves better results. \n\n| Adam, BS=16 | Adam, BS=1 | Adafactor, BS=1 | LoRA, BS=1 | \n| ---      | ---                  | ---                       |  ---                |\n| 18.2 ± 0.4%    | 18.7 ± 0.4%   | 18.7 ± 0.2%        | 16.7 ± 0.4% |  \n\n\n**Grid-search for pretraining GPT-2 (124M):** We re-ran our grid-search experiments reported in Figure 1b using a larger GPT-2 (124M) model trained on 2.5B tokens of FineWeb, and we observe an even clearer trend of batch size 1 being more robust to hyperparameters. Here we show a binarized version of the plot where the dark regions correspond to validation loss <3.5 and the white regions correspond to validation loss >3.5. The smaller batch size is much more robust to both the momentum (y-axis) and the learning rate (x-axis) hyperparameters.\n\n             Batch size 1          Batch size 512\n          +----------------+     +----------------+\n    0.999 |████████████████|     |                |\n    0.99  |████████████████|     |                |\n    0.97  |████████████████|     |                |\n    0.9   |████████████████|     |     ███████    |\n    0.7   |████████████████|     |                |\n    0     |████████████████|     |                |\n          +----------------+     +----------------+\n          2⁻¹⁴ 2⁻¹² 2⁻¹⁰ 2⁻⁸     2⁻¹¹ 2⁻⁹  2⁻⁷  2⁻⁵\n            Learning rate          Learning rate\n\nWe also tested Adafactor at 124M and 1.3B parameter scale for pretraining on FineWeb and in both cases it performed on par or even slightly better than Adam and AdamW with batch sizes 1 and 512 as shown in the table below. \n\n|               |  AdamW, BS=512 (baseline) |  Adam, BS=1, ${\\beta_2}=0.95$ |  Adam, BS=1, ${\\beta_2}=0.9999$ |    Adafactor, BS=1 |  SGD, BS=1 | \n| ---                    | ---                  | ---                       |  ---                | ---                    | ---                  | \n| GPT-2 (124M)    | 3.43           | 4.40                |  3.38               | 3.36⭐️                   | 3.52            | \n| GPT-3 (1.3B)      | 2.95               | 3.45              |  2.79            | 2.78⭐️              | 2.97                  | \n\n\nWe would like to kindly note that running our grid-search experiments at the 1.3B scale would require around 2.7M TPU-v4-chip hours, at an approximate cost of $8.7M, which is not an accessible resource for us. However, several of our experiments with different models, datasets, and optimizers demonstrate the validity of our findings and recommendations at scale, including our fine-tuning experiments with the 4B-parameter Gemma 3 model.\n\n**Continuous behavioral changes from small to large batch sizes:** It is true that in Figure 1b, we only compare two batch sizes (1 and 256). However, we also measured the sensitivity of Adam to its hyperparameters in Figure 2. We tested batch sizes 1, 4, 16, 64, 256, 1024, and 4096, and saw a gradual increase in hyperparameter sensitivity with each increment of batch size.\n\n**Surge phenomenon:** The surge phenomenon described in [40] demonstrates that optimal learning rate for Adam increases with batch size until a critical point, beyond which it decreases and then rises again (thus, a “surge”). We believe this result is complementary to our findings and not in any conflict. For example, we observed in Figure 3 that the optimal learning rate for batch size 4096 was lower than for batch size 1024.\n\n**Appendix:** We included our Appendix in the supplementary materials, together with our source code. You can download it by clicking on the “Supplementary Material:  zip” button in our submission’s page on Openreview.\n\nThank you again for your thoughtful feedback. We believe our revisions and additional experiments meaningfully improve our paper, and we hope you will consider raising your score. Please let us know if you have additional questions we can address."}, {"review_id": "URU6NmcTHp", "reviewer": "Reviewer_uKGs", "summary": "The paper argues that neural networks can be trained in a stable fashion at small batch sizes, with non-preconditioned optimizers like SGD, even without the use of momentum. It does experiments on various scales showing that the gap between preconditioned and non-preconditioned optimizers decreases at small batch sizes. It also shows on small scale experiments that training at small batch sizes is much more robust to hyperparameters as compared to large batch sizes.", "questions": "1. One major thing missing from the current manuscript is the description of the learning rate schedule. Can the authors please add a proper description and how it is varied with batch size.\n\n\n2. In Section 5.2, about hardware utilization, the authors state that we need a minimum batch size of a few hundred tokens to not be bottlenecked by memory bandwidth. Are the authors referring to the memory bandwidth of the GPU? Is this about transferring the data to the GPU or between HBM and SRAM of the GPU? Can the authors make it more clear?", "limitations": "The main limitation of the work is missing a proper discussion of previous theoretical and empirical works with similar claims. If the authors can improve on this, I would be happy to revise my rating.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "The main strength of the work is the rigorous exploration of the question of training at small batch sizes. By proper sweeps, it shows that training at small batch sizes is even more robust to hyperparameters, and generally performs at par with large batch sizes. It also provides a general scaling rule for $\\beta_2$ for Adam, based on keeping its half-life constant in terms of number of tokens.\n\nIn my opinion, the main weakness of the work is in not citing some previous theoretical works [1,2] which have already shown that preconditioners benefits decrease at small batch sizes. Other theoretical works [3] also establish that the benefit of momentum decreases at small batch sizes. Apart from this, scaling beta2 with small batch sizes has also been shown by previous works [4]. I do think the work is important as it highlights a crucial point - we can always train comparably at small batch sizes, even with more robustness to hyperparameters. However, I certainly think that the paper needs to properly discuss these related works and state that they empirically verify in language modeling setup that these things hold. Their practical recommendations are very useful. I would be willing to revise my rating once proper discussion of previous related work is done.\n\n[1] - Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model - 2019\n\n[2] - New Insights and Perspectives on the Natural Gradient Method - 2020\n\n[3] - On the insufficiency of existing momentum schemes for Stochastic Optimization - 2018\n\n[4] - Resolving Discrepancies in Compute-Optimal Scaling of Language Models 2024", "was_revised": true, "final_justification": "I have looked at the rebuttals and other reviews, and with the discussion of previous related work, I consider this paper as technically solid. I am therefore revising my rating to 5.", "rebuttal": "Thank you for your detailed and constructive feedback. We address your feedback below.\n\n**Related work discussion:** Thank you for highlighting important related work that we missed in our submitted manuscript. We have now properly cited these works and updated our related work section to include the following discussion of how our work complements and differs from these works. \n\nWhile similarly to our work, Zhang et al. (2019)[1] highlights that the benefits of more sophisticated optimizers diminish as batch size decreases, their theoretical results assume a convex quadratic objective that is quite different from language model training. Their empirical experiments on the other hand only cover a single language modeling experiment that is limited to a two-layer transformer. We believe that the training dynamics and optimization differ significantly for large language models and our purpose is to provide practical recommendations for LLM training in particular. Martens (2020)[2] establish theoretical benefits of second-order preconditioners. However, they do not study the effect of batch size on the performance of different optimizers, nor do they focus on LLM optimization. \n\nKidambi et al. (2018)[3] show that standard momentum methods do not provide performance gain over SGD in several simple problem instances, and that their empirical benefits noticed in practice often arise only in large-batch regimes. This result is aligned with our findings that sophisticated optimizers are increasingly beneficial with larger batch sizes, and that the performance gap is significantly smaller for small batch sizes. While Kidambi et al. (2018) focus on vision experiments, we focus on the LLM optimization setting that is not characterized by multiple epoch training like\nvision and is often constrained to a limited token budget, stopping far before convergence and potentially leading to different training dynamics and scaling laws. \n\nSimilarly to our work, Porian et al.(2025)[4] find that increasing the value of $\\beta_2$ is critical to achieving a competitive performance with Adam at small batches. However, the way they choose $\\beta_2$ at smaller batch sizes is fully based on a grid-search while we provide a scaling law for $\\beta_2$ with batch size. We validate in our results throughout the paper that this new scaling law can be used in a zero-shot manner and lead to significant improvement in performance at smaller batch sizes.\nIn summary, while these works offer theoretical and complementary perspectives, our contributions lie in challenging and putting to the test theoretical insights and empirical common wisdom specifically in the context of large language models optimization, providing recommendations that generalize across model size, datasets and optimizers, and deriving empirically-validated scaling laws that make our recommendations actionable. We acknowledge however that it is important to discuss these works, and we cite them indeed in the revised manuscript. We also discuss additional related work mentioned by reviewer 4.\n\n**Learning rate schedule:** We used a linear warmup for the first 5% of the training steps (from zero to peak learning rate), followed by cosine decay (from peak learning rate to zero). We use the same scheduler for all batch sizes. We updated our manuscript to reflect these experimental details. \n\n**Hardware utilization:** Indeed, we are referring to HBM bandwidth of the GPU/TPU. We store the model, optimizer, and full dataset on the GPU (the dataset is small), so there's no meaningful data movement between the GPU memory and CPU memory. Because modern GPUs/TPUs typically require an arithmetic intensity of several hundred FLOPs / byte to saturate tensor cores, they need to perform at least several hundred GEMM FLOPs per one byte of data loaded from HBM. In practice this means that every model parameter loaded from HBM needs to be used for at least several hundred (or even a few thousand) tokens in a batch. We would expect that increasing the batch size beyond this point provides no additional benefits in terms of hardware utilization, even if the accelerator has enough HBM to work with a larger batch size. Hence we recommend using the smallest batch size that achieves optimal or near-optimal hardware utilization. \n\nTo give a clear idea about the wall-clock time in our setting, we report the training time for GPT-2 (124M) using our codebase with a TPU v6e-1 accelerator, SGD optimizer, and sequence length 1024 in the table below. We obtained the highest throughput when using batch size 4 (i.e. 4096 tokens per batch). Batch size 1 resulted in low arithmetic intensity and thus poor throughput. Conversely, batch size 32 required using gradient checkpointing, which increased the FLOPs per step and slowed down training. When training on a single accelerator, small batch sizes can result in good hardware utilization, but the minimum batch size required to achieve good hardware utilization will depend on both hardware and implementation details. Therefore, our general recommendation is to use the smallest batch size that results in good hardware utilization (e.g. in this specific example we would recommend using batch size 4).\n\n|                    |    BS=1        |  BS=4        |  BS=16      | BS=32       | \n| ---             | ---             | ---            |  ---           |     ---        |\n| Time         |  3h:32min  |    ⭐️2h:43min |   3h:14min  |   4h:16min | \n\nWe updated our manuscript to clarify this point. \n\nThank you again for your thoughtful feedback. We believe the extended related work discussion, clarifications and additional experiments meaningfully improve our paper, and we hope you will consider raising your score. Please let us know if you have additional questions we can address."}, {"review_id": "2Ttt1GZxiE", "reviewer": "Reviewer_DoJ4", "summary": "The paper is interested in training with small batch sizes.  By ML theory (McCandlish et al), you'd expect that smaller batches are asymptotically more and more compute efficient (able to reach the same loss with less data), but recent practice (e.g., Porian et al.) has found there is an \"optimal batch size\" below which performance degrades.  Others have speculated that the problem with small batches might be that we should re-tune Adam's Beta parameters (Porian et al., Zhang et al.).  In this paper, they show that if, instead of holding β2 fixed, you hold its \"half-life\" fixed, in tokens (key point), then you can train stably all the way down to B=1.  This scaling rule for β2 essentially means that if you increase the batch size by a factor of K, we should scale optimal β2 by β2^K.  Interestingly, even at B=1, vanilla SGD without momentum performs somewhat reasonably.  This echos other work suggesting that a lot of the advances in optimizers (e.g., Muon, AdamW weight decay, LR decay schedules) are just methods to deal with gradient variance, but another way to deal with it is to just take a lot of very small steps.  Unfortunately, the paper doesn't test using AdamW, only vanilla Adam, and they don't reveal the LR schedule, so we don't really know how their recommendation plays with SOTA LLM training setups.", "questions": "- What is the LR schedule?\n\n- Intuitively, why would it be the case that we should maintain the velocity \\beta_2 timescale but not the momentum \\beta_1?\n\n- The paper recommends making B so small that it maximizes model throughput.  Do you mean the global optimum?  Because then we could just as easily say, \"set the B to the LARGEST value that maximizes model throughput\" or \"the single value that maximizes model throughput.\"  Perhaps you are really looking at some kind of tradeoff between loss and throughput?  That is, if B is too large, then loss will be worse, and it may take longer to train.  So we could frame this as finding the point where we can get the lowest loss in the least amount of wall-clock time, given a fixed amount of compute resources?  This is actually a good contrast case from the work in Bergsma et al (2505.13738), which assumes you want to minimize wall-clock time but have infinite parallel resources.  Furthermore, some kind of experimental results in this section would have been nice.  E.g., you could actually show that following your recommendations does improve a real training run.", "limitations": "Really only a very cursory discussion of limitations in the conclusion, and no real discussion of societal impact (good or bad!).  Having a limitations section, even in the appendices, would have been very nice, because it would have also helped the reviewer/reader understand what exactly was done here: did you use weight decay?  Did you use a LR schedule?  Did you try overtrained models in addition to Chinchilla-optimal?  Did you try very large models?  Etc.", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 2, "strengths_and_weaknesses": "# Strengths\n\n## Quality and clarity:\n\n- Good writing, great organization, good experiments, great figures: easy to follow, easy to evaluate, and easy to read.\n- The idea of keeping the timescale of \\beta_2 fixed is intuitive and interesting and cool.\n- The claims are well-supported.\n\n## Significance and originality:\n\n- The scaling law for \\beta_2 in particular is simple to implement, so could be beneficial to small-batch training runs, although the motivation for such small-batch (pre-training) runs is less clear.  I have not seen this particular scaling law for \\beta_2.\n\n# Weaknesses:\n\n## Discussion of prior work\n\n- Many of the key findings have been observed previously in other contexts, in papers that are not cited or cited for different reasons.\n- Plus, there is work that is very related.\n- I worry one of the other reviewers will be an author of one of these papers and they're going to UNLOAD on you :)\n\nHere are some pointers:\n\nShallue et al. 1811.03600v3\n- SGD without momentum can be competitive at small batch sizes: \"However, at batch sizes small enough that all optimizers are within their perfect scaling region, momentum optimizers perform identically to SGD without momentum. Though initially surprising, this identical performance at small batch sizes is consistent with observations made in Kidambi et al. (2018).\"\n\nZhang - How does critical batch size scale in pre-training - 2410.21676v1:\n- SGD without momentum can be competitive at small batch sizes: \"We observe that in small batch size regimes like 2^6, the performance gap between optimizing with and without momentum β1 is small while the gap increases as we double the batch size (Shallue et al., 2019).\"\n- Tuning \\beta_2 is crucial for AdamW at small batch sizes: \"found that β2 in Adam, the exponential decay rate of the second momentum estimate of gradients to smooth out the model update, also has significant effects on training for small batch sizes. This might be because gradients in small-batch training are sparser. Specifically, we ablate β2 ∈ [0.95, 0.99, 0.999] for all model sizes and batch sizes in [64, 128, 256, 512]. We find that the default value 0.95 in previous works that are set for millions of tokens batch size training might be sub-optimal\"\n\nPorian et al., 2406.19146v3\n- Tuning \\beta_2 is crucial for AdamW at small batch sizes: \"we observe that setting the AdamW β2 parameter to be 0.95 is suboptimal at smaller batch sizes (128 and below), and increasing it allows establishing clear trends from our small-scale hyperparameter sweep\"\n- \"Initially, we kept β2 at its previous value of 0.95. However, this led to poor results at smaller batch sizes: as the batch size gets smaller, the squared gradients become noisier, and AdamW requires more smoothing to obtain a correct denominator. Therefore, we added 0.99 and 0.999 to the sweep, obtaining improved performance on small batch sizes. This empirical observation matches the theoretical work by Zhang et al. [63], that showed that when the batch size is small, higher values of β2 are crucial for the convergence of Adam.\"  FYI, [63] is \"Y. Zhang, C. Chen, N. Shi, R. Sun, and Z.-Q. Luo. Adam can converge without any modification on update rules. Advances in neural information processing systems, 35:28386–28399, 2022.\"\n\nBusbridge - How to scale your EMA - 2307.13813v3:\n- The scaling rule for β2 here is exactly the same as the scaling rule for the momentum parameter --- in a model _weight_ EMA --- in Busbridge.\n\n- Also worth considering the EMA view of AdamW proposed by Wang & Aitchison (2405.13698) – looks at the timescale of AdamW, as a function of η and λ rather than the βs.\n\n- Also, the AdEMAMix optimizer 2409.03137 seems related as both methods essentially try to capture \"older\" information.  Maybe that paper should have experimented more with different \\beta_2 values?\n\n- Related to these prior paper studying AdamW rather than Adam: see scaling weight decay with batch size, in Bergsma et al (2505.13738) as motivated by the EMA view of Wang & Aitchison (2405.13698)\n\nChowdhery et al - PaLM - Scaling Language Modeling with Pathways - 2204.02311\n- \"The second-order moment interpolation value is computed as 1-k^(-0.8), where k is the step number. We have found this to be more stable than the standard B2 when training large language models, because rare embedding tokens can have poorly estimated second moments over\nshorter windows\"\n- Could something to do with embeddings help explain your results?\n\n## Missing theory\n\nFinally, this paper seems to miss that ML theory, e.g., McCandlish et al., specifically predict that smaller batches are (asymptotically) better, i.e., can reach a given target loss in fewer steps.\n- this paper supports this hypothesis over the U-shaped optimal batch size, so it would be good to approach this directly in the discussion.\n\n## Generality of findings\n\nMainly, my other concerns revolve around whether this is a competitive LLM training setup:\n- we use Adam here rather than AdamW, but weight decay is used in all SOTA LLM training, AFAIK\n- nothing is said about the LR schedule.  Decaying the LR can also reduce variance (Bergsma et al., 2502.15938), so what about that?\n- models are undertrained (too large for their token budget).  I was especially confused by the statement: \"We train our 49M model on 600 million tokens from the FineWeb-Edu dataset [48], following the Chinchilla scaling laws [25].\"  The appendix likewise says, \"We trained most models to follow Chinchilla scaling laws, i.e. using 20 tokens per active parameter.\" But 49M on 600M is more like 12 TPP, not 20 TPP.  The GPT-3 model of size 1.3B trained on 9B tokens, that's 6.9TPP -- that is highly under-trained.  No one would train such a model in practice because we could get a better loss with a smaller model (easier to serve) using fewer compute FLOPs.\n- you do go as high as 20TPP in one set of experiments, but overtrained models are also very popular these days so it would be nice to understand how your results generalize to high TPP, where gradient variance is even higher.\n- also, I'm struck by the lack of motivation for small-batch training.  Small batches are unlikely to be used in training giant models.  Yes, we could use smaller batches at smaller scales, but then would we worry that those findings don't transfer as well to larger scales?  Yes, we could use small batches in FT, but this paper mostly looked at PT, right?\n\n# Minor points:\n\n- \"As a consequence, optimizers that perform well in large-batch training are increasingly standard practice\" -- Was this really the motivation behind (cited) Muon, Sophia, Shampoo, SOAP, LAMB?  I think LAMB was motivated to let us train FASTER (greater parallelism), not to address stability issues.  I never heard e.g., SOAP or Muon motivated as a consequence of the belief that large batches are more stable?\n\n- \"4.3 Adam Hyperparameters: How to Scale Them with Batch Size and Why That is Necessary\" - note this section only very loosely addresses the question of *why* tuning β2 is necessary.  The only explanation in this section seems to be: because it works better.\n\n- \"we hypothesize that large batch sizes require careful hyperparameter tuning since they require large step sizes and therefore have to make predictions about the loss surface further away from the current parameter vector.\"  Technically speaking, a \"hypothesis\" is an idea that you can test.  I am not sure that's the right word to the conceptual intuition given here.\n\n- Line 105, just for completeness, should specify \"where d1 and d2 are ...\"", "was_revised": true, "final_justification": "This is interesting work and the authors did a very reasonable job.  After reading the rebuttal, I felt like my main addressable concerns were addressed.  I read the other reviews and feel like Reviewer uKGs and myself are most sympatico, and believe they may also raise their score after the rebuttal.\n\nThe \"Accept\" rating means: \"Technically solid paper, with high impact on at least one sub-area of AI or moderate-to-high impact on more than one area of AI, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.\"\n- I am not 100% sure about \"high impact\" on the field of optimization, but overall this description of the paper most resonates with me.", "rebuttal": "Thank you for your thoutghful feedback. We address the points you raised below. \n\n**Extended related work discussion:** Thank you for providing pointers to related work. We have now updated our manuscript to include the following discussion of how our work extends and contrasts with these previous works.\n\n[1] and [2] report that momentum provides limited benefits when batch sizes are small. These results are aligned with our findings and validate our empirical observations. However, [1] does not compare the performance of SGD vs Adam and [2] claims that SGD performs significantly worse than Adam. We hypothesize this observation could be a result of [2] using 64 as the smallest batch size. We test batch sizes all the way down to 1.\n\n[3] similarly show that increasing the value of $\\beta_2$ is important for small-batch training with AdamW. However, their approach is based on discrete hyperparameter sweeps across a few values of $\\beta_2$. In contrast, we introduce and validate a principled scaling rule that holds the second-moment half-life constant in terms of tokens. This leads to a predictable and interpretable formula for scaling $\\beta_2$ with batch size. We validate this rule across different optimizers, model sizes, and datasets, and show that it leads to better performance in a zero-shot manner without sweeping.\n\n[4] propose a similar token-timescale interpretation of exponential moving averages in the context of model weight averaging. While their work supports the intuition behind our $\\beta_2$-scaling rule, our focus is on second-moment estimation and we validate the scaling law across a broad range of pretraining and fine-tuning settings for LLMs.\n\nThe AdEMAMix optimizer [5] also addresses long-range information retention in gradient estimation by mixing momentums with different timescales. However, our contribution shows that such complexity may not be necessary: by simply tuning $\\beta_2$ using our scaling rule, even Adam can perform competitively at small batch sizes. Unlike AdEMAMix, our recommendations apply to widely used, off-the-shelf optimizers and are validated empirically in memory-constrained fine-tuning settings.\n\n[6] and [7] provide a more theoretical perspective on how optimizer behavior and weight decay can be interpreted in terms of timescale or EMA dynamics. Our contribution complements theirs by operationalizing this perspective into a concrete $\\beta_2$ scaling rule and empirically verifying it across different LLM training configurations. We also show that this scaling rule is critical for enabling training with batch size 1, something not addressed in their work.\n\n[8] note that dynamically adjusting $\\beta_2$ leads to better performance for rare token embeddings. While this is an important insight, our scaling rule provides a simpler, batch-size-aware heuristic that applies even when using constant $\\beta_2$ throughout training.\n\nReferences:\n- [1] Shallue et al. - 1811.03600v3\n- [2] Zhang et al. - 2410.21676v4\n- [3] Porian et al. - 2406.19146v3\n- [4] Busbridge et al. - 2307.13813v3\n- [5] AdEMAMix optimizer - 2409.03137v2\n- [6] Bergsma et al. - 2505.13738v1\n- [7] Wang & Aitchison - 2405.13698v3\n- [8] PaLM - 2204.02311v5\n\n**Generality of our findings and new results:** In the submitted version of our paper, we pretrained models ranging from 30M active parameters to 1.3B. We pretrained all models (except the largest 1.3B model) with 20 TPP. We trained the 1.3B model for 10B tokens following the GPT-2 paper as well as the modded-nanoGPT repository, purely because we were compute-constrained. Since then, we have expanded our results in two ways. First, we tested more optimizers at the 124M and 1.3B scales.\n\n|               |  AdamW, BS=512 (baseline) |  Adam, BS=1, ${\\beta_2}=0.95$ |  Adam, BS=1, ${\\beta_2}=0.9999$ |    Adafactor, BS=1 |  SGD, BS=1 | \n| ---                    | ---                  | ---                       |  ---                | ---                    | ---                  | \n| GPT-2 (124M)    | 3.43           | 4.40                |  3.38               | 3.36⭐️                   | 3.52            | \n| GPT-3 (1.3B)      | 2.95               | 3.45              |  2.79            | 2.78⭐️              | 2.97                  | \n\nSecond, we added fine-tuning. We fine-tuned a non-instruction-tuned Gemma3 4B on the MATH dataset for 5 epochs and measured its accuracy on the validation question set. We found Adafactor (in bfloat16 precision with stochastic rounding) and batch size 1 to perform just as well as Adam in float32 precision with batch size 1 or 16. We believe this experiment is important both to prove the generality of our findings but also to provide a memory-efficient full fine-tuning recipe as an alternative to LoRA.\n\n| Adam, BS=16 | Adam, BS=1 | Adafactor, BS=1 | LoRA, BS=1 | \n| ---                    | ---                  | ---                       |  ---                |\n| 18.2 ± 0.4%    | 18.7 ± 0.4%   | 18.7 ± 0.2%        | 16.7 ± 0.4% |  \n\n**Practical motivation:** We believe that this work provides at least two practical contributions. **First,** we show that smaller batch sizes are much more **robust to hyperparameter** values and optimizer design, meaning that without any hyper-parameter tuning or with limited hyper-parameter tuning, smaller batch sizes should work better in expectation compared to larger batch sizes. It is a common practice to use gradient accumulation to simulate larger batch sizes. We posit that accumulating gradients not only increases memory consumption, it also leads to worse performance. **Second**, we extend our findings to fine-tuning in a memory-constrained setting. Because it is commonplace to fine-tune large models on small datasets, fine-tuning is very often memory-constrained rather than compute-constrained. We show that in a memory-constrained setting, it is possible to use a small batch size combined with a simple optimizer like Adafactor to get the performance benefits of full fine-tuning while having little memory overhead compared to just storing the model in GPU memory, i.e. requiring only around 2 bytes / parameter.\n\n**Chinchilla scaling:** We used Chinchilla scaling as in \"20 tokens per (active) parameter\". We only used dense models for all experiments but we did not tie the input and output embeddings, which resulted in a different number of total vs. active parameters. For example, the \"49M\" model only has 30M active parameters, which is why we trained it on 30M x 20 = 600M tokens. We specify the exact dimensions of each model in Table 1 in the Appendix.\n\n**Learning rate schedule:** We used a linear warmup for the first 5% of the training steps (from zero to peak learning rate), followed by cosine decay (from peak learning rate to zero).\n\n**AdamW:** AdamW is indeed more widely used in practice than plain Adam. Unfortunately, weight decay introduces an additional hyperparameter whose optimal value depends on the batch size. When comparing different optimizers in Figure 1a, we wanted to perform a dense grid search over all hyperparameters for all batch sizes to ensure the comparison is as rigorous as possible. We did not want to use any scaling heuristics for weight decay, as this could bias the results. Since we could not afford to increase the dimensionality of the dense grid search from 4 to 5 dimensions, we did not use weight decay in this experiment at all. Naturally this opens up the question: would the results look any different for AdamW, compared to Adam? While this is not something we measured directly using a grid search, it is something we tested indirectly. Namely in Figure 5, the baselines for GPT-2 and GPT-3 actually do use weight decay, so their labels should be \"AdamW\", not \"Adam\". Hence this result provides a partial answer: Adam with batch size 1 works at least as well as AdamW with batch size 512. AdamW with batch size 1 could only possibly work better, if we tuned the weight decay coefficient instead of fixing it to zero.\n\n**Why fixing $\\beta_2$ half-life works:** We illustrate in the Appendix using a toy example why the first moment is not necessary for small batch sizes, but it also does not hurt performance. As the batch size is increased, momentum starts to become necessary for good performance, however, $\\beta_1$ cannot be set too high, because then the optimizer step direction becomes insensitive to local gradient information. In contrast, we hypothesize that the role of the second moment is to provide preconditioning at a longer token time-scale. We do not have a theoretical foundation to motivate these claims but we believe that our empirical results are well-supported and could be of practical value to the community. We like your intuition that the second moment needs a longer token time-scale to capture rare tokens.\n\n**Throughput:** Our general recommendation is to use the smallest batch size that results in good hardware utilization. Modern ML accelerators typically require an arithmetic intensity of at least several hundred FLOPs / byte (i.e. several hundred tokens in a batch per active model parameter) to saturate their matrix multiplication units. We expect that increasing the batch size beyond this point provides little additional benefits in terms of hardware utilization, even if the accelerator has enough memory to work with a larger batch size. Hence we recommend using the smallest batch size that achieves optimal or near-optimal hardware utilization. We share a table in our response to rewviewer 3pDd that shows the model throughput we were able to achieve for different batch sizes, although this is very much hardware / implementation specific.\n\nThank you again for your thoughtful feedback. We believe the extended related work discussion, clarifications and additional experiments meaningfully improve our paper, and we hope you will consider raising your score. Please let us know if you have additional questions we can address."}]

Author rebuttal:

Meta-review: The paper presents a thorough empirical study challenging the prevailing belief that large batch sizes are essential for stable and efficient language model training. The authors demonstrate that, with appropriate hyperparameter scaling—particularly adjusting the second-moment decay parameter β₂ based on token half-life—small batch sizes, even down to one, enable stable training, increased robustness to hyperparameters, and competitive or superior per-FLOP performance compared to larger batches. Notably, they show that vanilla SGD without momentum performs on par with more sophisticated optimizers, leading to significant memory savings and practical benefits in memory-constrained fine-tuning scenarios. 

The strengths of this work lie in its extensive experimental design, inclusion of diverse optimizers (Adam, Adafactor, Muon), evaluation on multiple datasets (FineWeb, mC4, MATH), and scaling experiments up to 1.3B-parameter models. 

However, the key concern is that the findings are primarily empirical and lack a rigorous theoretical framework explaining why small-batch regimes generalize effectively, limiting broader understanding and transferability.

During the rebuttal, the authors addressed many reviewer concerns by expanding experiments to larger models, multilingual data, and fine-tuning setups, demonstrating consistent trends favoring small batches and showing that Adafactor can achieve full fine-tuning performance with minimal memory overhead. 

Despite the lack of deeper theoretical insights, the comprehensive experimentation, novel scaling heuristics, and practical relevance make this a valuable and timely contribution. I recommend acceptance given its strong empirical findings, actionable recommendations for practitioners, and meaningful impact on optimization strategies for language model training.
Decision: Accept (poster)