Thank you for your thoutghful feedback. We address the points you raised below.

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.

[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.

[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.

[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.

The 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.

[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.

[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.

References:

• [1] Shallue et al. - 1811.03600v3
• [2] Zhang et al. - 2410.21676v4
• [3] Porian et al. - 2406.19146v3
• [4] Busbridge et al. - 2307.13813v3
• [5] AdEMAMix optimizer - 2409.03137v2
• [6] Bergsma et al. - 2505.13738v1
• [7] Wang & Aitchison - 2405.13698v3
• [8] PaLM - 2204.02311v5

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.

Second, 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.

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.

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.

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).

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.

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.

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.

Thank 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.

I have read the rebuttal to my original review and I find the following points compelling:

• willingness to expand their discussion of related work, demonstrating good grasp of it ✔️
• interesting new experiments with other optimizers and FT ✔️
• convincing discussion of practical motivation ✔️
• clarifications on other points that I raised (which I gather will be used to "meaningfully improve [the] paper") ✔️

I will revise my score.

Thank you for your detailed and constructive feedback. We address your feedback below.

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.

While 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.

Kidambi 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 vision and is often constrained to a limited token budget, stopping far before convergence and potentially leading to different training dynamics and scaling laws.

Similarly 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. In 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.

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.

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.

To 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).

We updated our manuscript to clarify this point.

Thank 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.

Thank you for your feedback. We address your questions and concerns below.

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.

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.

However, 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.

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.

         Batch size 1          Batch size 512
      +----------------+     +----------------+
0.999 |████████████████|     |                |
0.99  |████████████████|     |                |
0.97  |████████████████|     |                |
0.9   |████████████████|     |     ███████    |
0.7   |████████████████|     |                |
0     |████████████████|     |                |
      +----------------+     +----------------+
      2⁻¹⁴ 2⁻¹² 2⁻¹⁰ 2⁻⁸     2⁻¹¹ 2⁻⁹  2⁻⁷  2⁻⁵
        Learning rate          Learning rate

We 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.

We 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.

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.

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.

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.

Thank 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.

Thank you for your detailed review. We address your feedback below.

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.

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.

         Batch size 1          Batch size 512
      +----------------+     +----------------+
0.999 |████████████████|     |                |
0.99  |████████████████|     |                |
0.97  |████████████████|     |                |
0.9   |████████████████|     |     ███████    |
0.7   |████████████████|     |                |
0     |████████████████|     |                |
      +----------------+     +----------------+
      2⁻¹⁴ 2⁻¹² 2⁻¹⁰ 2⁻⁸     2⁻¹¹ 2⁻⁹  2⁻⁷  2⁻⁵
        Learning rate          Learning rate

We 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.

We 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.

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.

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.

However, 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.

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.

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.

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.

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.

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.

Thank 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.

References:

• [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.
• [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.
• [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.