perspectiva aiqbenchCercetare AI

This study demonstrates that stochastic noise introduced by mini-batch sampling plays a critical role in generalization, as full-batch gradient descent fails to match the validation performance of standard training despite fitting the data better. By replicating GPT-2 training with exact full-batch updates on a substantial dataset, the author provides empirical evidence that optimization dynamics, rather than just data distribution, determine model outcomes. This work offers valuable insights for researchers investigating learning dynamics and compute-efficient training strategies.

Sursă care acoperăLessWrong

Ce este aceastaLectura unui model asupra articolului și a acoperirii conexe, scrisă pentru aiqbench. Nu sunt cuvintele publicației.

I trained GPT-2 with full-batch gradient descent so you don't have to

LessWrong14 octombrie 202613 min de cititArticolul și perspectiva sunt în engleză.

Note: The results of this are not as clean as they would naively appear, since a lot of training parameters are tuned for non-full-batch gradient descent. See discussion in the comments. I've left this up in case it's still useful, but be careful about drawing conclusions naively from the result, and if you're trying to do anything mathematical and precise look at the reproduction code for confounders. The impetus for doing this was provided in large part by Charles Renshaw-Whitman. Learning dynamics people suggested to me that noise was essential in the training of neural networks. I thought this was kind of counterintuitive--- I usually naively operate from a model in which the pretraining data basically specifies the outcome by virtue of the samples all "pulling" on the loss landscape in some relatively uniform way which incentivizes converging to a point that generalizes well, but instead it seems to be the case that the actual dynamics of the optimization process itself have predictable and measurable effects on the qualitative outcomes of training. The idea seems reasonable based on vague intuitions I have about simulated annealing. I wanted to see directly that this is true at some nontrivial scale, I happened to have access to a pool of free compute that no one else seemed to be using, and I figured this would be an interesting object of study for compute-constrained researchers with an interest in studying learning dynamics closely, so I did some very rough calculations and realized this was feasible to do with GPT-2. The main result is that on a fixed 800M-token subset of FineWeb, ordinary training of GPT-2 (based on the NanoGPT speedrun) with a batch size of 393K, gets to 3.30 validation loss in one epoch (training loss is the same as validation loss given no repeats on the data). Exact full-batch gradient descent on exactly the same tokens fits the training set better (3.09 training loss in 350 steps), but its validation loss converges at about 3.50 in somewhere between 260 and 310 steps (presumably the training loss hasn't converged by this point). This suggests the batch selection noise is in fact doing something essential. (Note that this removes the effect of noise due to batch selection, but not noise due to finite step size.) The code is on GitHub. Everything the runs produced is on Hugging Face: metrics and dynamics statistics for every run, the full gradient and update at every step of the full-batch run, per-token losses on all 800M training tokens every 10 steps, checkpoints, and the exact training set. Appendix: some methodological choices These particular parameters are somewhat arbitrary, and were tuned around the size of dataset which, for some batch size, we could use to train a model to achieve approximately 3.28 validation loss (as a heuristic to save on time spent waiting). 3.28 is completely arbitrary and taken from the NanoGPT objective. For example, I determined that it is also possible to achieve a validation loss of 3.23 in the same 350 steps with a batch size of 50M tokens (consequently requiring 18B tokens). It is also possible to reach this loss with only 329M tokens and 1200 steps using some batch size/learning rate scheduling (as the NanoGPT speedrun record does, but I didn't want to be confounded by optimizations like this). 3.28 is also not the point at which any relevant dataset or batch size configuration actually converges, e.g, if you keep training in the speedrun record configuration for longer you'll get closer to 3.0. Consequently, it's possible that these checkpoints still lack some information about what happens when we get close to convergence points, etc., and generally we haven't "finished training GPT-2" in a sense which can be said of any realistic training configuration. Nevertheless, they likely contain information relevant to early developmental phases of training. An info dump from Claude about learning rate scheduling, as well as datapoints about other possibly relevant training runs follows: Setup The model is the current speedrun record (#360, Sept 2026), with the parts that don't make sense for full-batch training removed: The hashed n-gram embedding table. It's a 65B-parameter lookup table keyed on the previous one or two tokens, and it memorizes any data that repeats: validation loss explodes the moment a token is seen twice. Full batch on a fixed set repeats every token every step, so it has to go. The sampled softmax, which makes the loss biased and stochastic. Every schedule that changes the loss function during training (sequence length, attention window, multi-token-prediction weights, batch size). The whole run uses the record's final-stage settings. What remains is an 11-layer, width-768 model with the record's FP8 and FlashAttention-3 kernels and its optimizer: a Muon-family optimizer (ANVIL) for the matrices and Adam for everything else. These changes cost about 40% more steps to reach 3.28 than the record (about 1,625

The impetus for doing this was provided in large part by Charles Renshaw-Whitman.

Learning dynamics people suggested to me that noise was essential in the training of neural networks. I thought this was kind of counterintuitive--- I usually naively operate from a model in which the pretraining data basically specifies the outcome by virtue of the samples all "pulling" on the loss landscape in some relatively uniform way which incentivizes converging to a point that generalizes well, but instead it seems to be the case that the actual dynamics of the optimization process itself have predictable and measurable effects on the qualitative outcomes of training. The idea seems reasonable based on vague intuitions I have about simulated annealing.

I wanted to see directly that this is true at some nontrivial scale, I happened to have access to a pool of free compute that no one else seemed to be using, and I figured this would be an interesting object of study for compute-constrained researchers with an interest in studying learning dynamics closely, so I did some very rough calculations and realized this was feasible to do with GPT-2.

The essential result is that on a fixed 800M-token subset of FineWeb, ordinary training of GPT-2 (based on the NanoGPT speedrun) with a batch size of 393K, gets to 3.30 validation loss in one epoch (training loss is the same as validation loss given no repeats on the data). Exact full-batch gradient descent on exactly the same tokens fits the training set better (3.09 training loss), but its validation loss converges at about 3.50 in somewhere between 260 and 310 steps (presumably the training loss hasn't converged by this point).

This suggests the batch selection noise is in fact doing something essential. (Note that this removes the effect of noise due to batch selection, but not noise due to finite step size.)

The code is on GitHub. Everything the runs produced is on Hugging Face: metrics and dynamics statistics for every run, the full gradient and update at every step of the full-batch run, per-token losses on all 800M training tokens every 10 steps, checkpoints, and the exact training set.

These particular parameters are somewhat arbitrary, and were tuned around the size of dataset which, for some batch size, we could use to train a model to achieve approximately 3.28 validation loss (as a heuristic to save on time spent waiting). 3.28 is completely arbitrary and taken from the NanoGPT objective.

For example, I determined that it is also possible to achieve a validation loss of 3.23 in the same 350 steps with a batch size of 50M tokens (consequently requiring 18B tokens). It is also possible to reach this loss with only 329M tokens and 1200 steps using some batch size/learning rate scheduling (as the NanoGPT speedrun record does, but I didn't want to be confounded by optimizations like this). 3.28 is also not the point at which any relevant dataset or batch size configuration actually converges, e.g, if you keep training in the speedrun record configuration for longer you'll get closer to 3.0.

Consequently, it's possible that these checkpoints still lack some information about what happens when we get close to convergence points, etc., and generally we haven't "finished training GPT-2" in a sense which can be said of any realistic training configuration. Nevertheless, they likely contain information relevant to early developmental phases of training.

The hashed n-gram embedding table. It's a 65B-parameter lookup table keyed on the previous one or two tokens, and it memorizes any data that repeats: validation loss explodes the moment a token is seen twice. Full batch on a fixed set repeats every token every step, so it has to go.

The sampled softmax, which makes the loss biased and stochastic.

Every schedule that changes the loss function during training (sequence length, attention window, multi-token-prediction weights, batch size). The whole run uses the record's final-stage settings.

What remains is an 11-layer, width-768 model with the record's FP8 and FlashAttention-3 kernels and its optimizer: a Muon-family optimizer (ANVIL) for the matrices and Adam for everything else. These changes cost about 40% more steps to reach 3.28 than the record (about 1,625 instead of 1,190, at 393K tokens per step).

The dataset is the first 800M tokens of FineWeb10B. I picked this size because small-batch training on it barely overfits: 2,372 steps at 393K tokens per step, about 1.2 passes over the data, reaches 3.248 validation loss with a train/val gap of 0.06.

For full batch, each step accumulates the exact gradient of the mean loss over all 800M tokens (16,272 microbatches of 49K tokens, with fp32 accumulators), then takes one optimizer step. I ran 350 steps. The budget comes from a sweep of large batches on fresh data, which extrapolates to roughly 300 to 330 steps to reach 3.28 at infinite batch; a fresh-data run at 50M tokens per step reaches 3.23 in 350 steps. The full-batch run took about 100 GPU-hours on 4 H200s on preemptible compute, surviving 106 preemptions by checkpointing both between and partway through steps.

3.521 (best 3.498)

"Training loss" here is loss on a fixed 10M-token sample of the training set. At equal training loss, full batch's validation loss is about 0.2 nats worse than small batch's. Full batch's validation loss is flat from about step 260 to the end, while its training loss keeps falling. Intermediate batch sizes sampled from the same 800M set (4.7M to 200M tokens per step) fall progressively further behind fresh data at the same batch size as each token gets seen more times.

Consecutive full gradients point in nearly opposite directions: the median cosine between them is −0.945. On 98% of steps, every single microbatch gradient is anti-aligned with the previous step's full gradient.

Measuring the loss along each step's actual update, the step overshoots the minimum in its own direction by about 2× and lands at about the same loss on the far side.

The exact loss over all 800M training tokens goes up on about a third of steps.

The top Hessian eigenvalue keeps growing: 48, 82 and 113 at steps 100, 200 and 300.

This is the picture from Cohen et al. on edge-of-stability, here with a Muon-family optimizer on a 2026 speedrun architecture.

One run per configuration. Seed-to-seed noise in validation loss for this model is about 0.003, so the gap isn't noise, but its exact size isn't pinned down.

Full batch on a fixed set means 350 passes over the same tokens, so this doesn't separate "no gradient noise" from "too many repetitions". Full batch plus artificially injected minibatch-shaped noise would separate the two; I haven't run that.

The learning rate was tuned on fresh data, not on the fixed set.

I also tried a schedule ten times as long on fresh data as a convergence reference. In that run the model's FP8 training forward pass and its bf16 evaluation forward pass drifted apart (3.15 against 4.62 on the same unseen data), so I don't trust its number. The full-batch run doesn't have this problem: its two forward passes agree within 0.015 at the end.

Every run that contributed data to this post, grouped by purpose. "Tokens processed" counts every token in every step, including repeats. "Train" is loss on a fixed 10M-token sample of the 800M training set for runs on that set, and on the training subset for the ablation subset runs; "—" means it wasn't measured. "Val" is the standard speedrun validation set (10.5M tokens) at the end of the run, with the best value in parentheses where it differs. Unless noted, runs use the fixed-objective model described above, on 4 H200s or 8 H100s; metrics for all of them are in the Hugging Face dataset.

the main run; lr multiplier 1.2

about 1.16 epochs; lr multiplier 1.73; at the end of the first epoch (step ~2,034) train ~3.26, val ~3.30

Batch-size ladder on the same 800M set (minibatches drawn without replacement within each pass), with fresh-data runs at the same batch size, learning rate and schedule. Runs marked "halted" were stopped by a rule comparing them against fresh data at matched steps:

Fixed set, 4.7M, lr 1.73

Fresh data, 4.7M

halted; 6.5 epochs

Large-batch sweep on fresh data, used to estimate how many steps full batch would need (about 300 to 330 at infinite batch). Steps to 3.28 is where the validation curve crosses 3.28, for runs that got there:

3.341 / 3.274 / 3.235

0.8 / 1.2 / 1.73 / 2.6

3.322 / 3.301 / 3.303 / 3.364

0.8 / 1.2 / 1.73 / 2.6

3.320 / 3.300 / 3.314 / 3.349

0.8 / 1.2 / 1.73 / 2.6

3.320 / 3.297 / 3.322 / 3.371

memorized the set: training loss 0.04 by step 224

Convergence references on fresh data, at 1×, 3× and 10× the schedule length that reaches 3.28:

FP8 training and bf16 eval passes drifted apart (3.15 against 4.62 on the same unseen data)

3.349 / 3.304 / 3.274, 3.278

3.333 / 3.279, 3.287 / 3.261 / 3.217

3.21 / 3.03 / 2.93

3.292 / 3.326 / 3.468

3.12 / 3.08 / 2.76

3.221 / 3.259 / 3.542

Unveiling the Basin-Like Loss Landscape in Large Language Models (Chen et al., 2025)

The Break-Even Point on Optimization Trajectories of Deep Neural Networks (Jastrzębski et al., 2020) (contains a study of roughly the same thing, but on CIFAR-10)

Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective (Wen et al., 2024)

Understanding Optimization in Deep Learning with Central Flows (Cohen et al.)

Text integral sub licența deschisă a publicației · 1734 cuvinte

Citește pe LessWrong ↗Întreabă despre acest articol