Scaling Laws and Optimization

You have $10^{24}$ FLOPs and one attempt. Should you train a 30B model on 5T tokens, or a 100B model on 1.5T? Both cost the same. One of them will be noticeably better, and you will not find out which until the money is spent. Then, having chosen, what learning rate do you set? You cannot sweep it; a single run takes a month.

Ten years ago these questions were answered by folklore. Today they are answered by two bodies of work. Scaling laws say that loss is a smooth, predictable function of model size, data and compute, so you can extrapolate from cheap runs to the expensive one. Optimization practice is the accumulated knowledge of which optimizer settings survive at scale, and why runs blow up when they do. This chapter covers both, because in practice they are used together: scaling laws pick $N$ and $D$; optimization makes the run that uses them succeed.

Part A: Scaling laws

The observation: loss is a power law

The problem in 2019 was that nobody knew whether bigger models would keep getting better, or whether the returns would flatten. Kaplan et al. (2020) at OpenAI trained hundreds of transformers spanning seven orders of magnitude of compute and plotted the test loss against three quantities: non-embedding parameters $N$, dataset tokens $D$, and compute $C$. Every plot was a straight line on log-log axes. That is the signature of a power law.

Read $L(N)$ as "the loss you get from a model of size $N$ when data is not the bottleneck". Kaplan's fits gave $\alpha_N \approx 0.076$ and $\alpha_D \approx 0.095$. The exponents are small, which means the lines are shallow: to halve the loss you need to increase $N$ by a factor of $2^{1/0.076} \approx 9{,}000$. But they are straight over seven orders of magnitude, with no sign of bending. That was the headline: scaling works, and it works predictably.

Kaplan also asked how to split a compute budget between $N$ and $D$. Their answer, from fitting $L$ as a function of both, was that model size should grow much faster than data: $N \propto C^{0.73}$ and $D \propto C^{0.27}$. Between 2020 and 2022 the field followed this advice. GPT-3 (175B on 300B tokens), Gopher (280B on 300B), and Megatron-Turing NLG (530B on 270B) all spent their compute on parameters.

The Chinchilla result

Hoffmann et al. (2022) at DeepMind re-ran the experiment with one change that turned out to matter enormously: they tuned the learning-rate schedule to the length of each run, instead of using a fixed schedule for all of them (which Kaplan had done, leaving short runs under-decayed and so under-performing). With that fixed, the picture changed. Using three separate methods, they found that for compute-optimal training:

Model size and tokens should scale together. Doubling the budget means a model $\sqrt{2}$ times bigger trained on $\sqrt{2}$ times more data. The practical rule that fell out of their isoFLOP curves: about 20 tokens per parameter. By that rule, Gopher's 280B parameters wanted 5.6T tokens, not 300B; Gopher was drastically under-trained. To prove it they trained Chinchilla, a 70B model on 1.4T tokens with exactly Gopher's compute, and it beat Gopher on essentially every benchmark.

The fitted loss form

The most useful output of the Chinchilla paper is a formula. Their third approach fits the final loss as the sum of three pieces: an irreducible floor, a term that shrinks with model size, and a term that shrinks with data.

with fitted constants $E \approx 1.69$, $A \approx 406.4$, $B \approx 410.7$, $\alpha \approx 0.34$, $\beta \approx 0.28$ (loss in nats per token on their held-out data). Read it as: $E$ is the entropy of text itself, which no model can beat; $A/N^\alpha$ is the penalty for having finitely many parameters; $B/D^\beta$ is the penalty for having seen finitely many tokens. Both penalties fall as power laws and both go to zero in the limit.

Deriving the compute-optimal split

Given this formula and the constraint $C = 6ND$, where is the minimum? Set up the trade: spend a little compute moving from $D$ to $N$, and the loss changes by the difference of the two marginal terms. At the optimum they balance. Working through the algebra (a Lagrange multiplier on $\log N + \log D = \text{const}$) gives:

With the published constants, $a = 0.45$, $b = 0.55$, $G = 1.34$. Both exponents are near one half, which is the "scale them together" result. Note that $b$ is slightly larger than $a$: as budgets grow, the fit wants tokens to grow a little faster than parameters.

Why later models over-train past Chinchilla-optimal

Chinchilla answered "what is the best model for a fixed training budget". But that is the wrong question if you are going to serve the model to millions of users. Every generated token costs about $2N$ FLOPs at inference. A model half the size is half the price to run, forever. If inference compute will dwarf training compute, and for a popular model it does, then you should accept a slightly worse loss for a much smaller model, and pay for it with more training tokens.

This is the Llama philosophy, stated explicitly in the Llama 1 paper (Touvron et al., 2023): "the objective of the scaling laws from Hoffmann et al. is to determine how to best scale the dataset and model sizes for a particular training compute budget. However, this objective disregards the inference budget." Llama 3 8B was trained on about 15T tokens, roughly 1,900 tokens per parameter, nearly a hundred times past the 20:1 rule. Under the fitted law its loss ($1.949$) is only slightly worse than Chinchilla 70B's ($1.937$) at a fraction of the serving cost. Sardana and Frankle (2023) formalize this by adding expected inference tokens to the objective; the more you plan to serve, the smaller and longer-trained the optimal model becomes.

The loss surface being flat near the optimum is what makes this a good trade rather than a heroic sacrifice: the $B/D^\beta$ term keeps shrinking as you add tokens, and it partly compensates for the larger $A/N^\alpha$ term of the smaller model.

The power-law plot

Here is the fitted law drawn as a family of curves. Each line is loss against model size for a fixed token count. The curves flatten on the right where $N$ is large enough that the data term dominates: adding parameters to a data-starved model does nothing. Adjust the exponents to see how sensitive the whole picture is to them.

Emergent abilities and the debate

Loss scales smoothly. Do capabilities? Wei et al. (2022) collected examples of tasks (multi-step arithmetic, word unscrambling, some BIG-Bench tasks) where performance sat at chance for small models and then jumped sharply at some scale, and called these emergent abilities: abilities that are absent in smaller models and present in larger ones, not predictable from the small-model trend. If real, this would mean that scaling laws on loss tell you little about what a bigger model will be able to do.

Schaeffer et al. (2023) offered a deflationary reading. Most of the "emergent" tasks were scored with discontinuous metrics: exact match on a five-digit answer, where being off by one digit scores zero. If the model's per-token accuracy improves smoothly, the probability of getting all five digits right goes as accuracy to the fifth power, which looks like a sudden jump when the per-token accuracy crosses a threshold. Replace exact match with a continuous metric like token edit distance, and the curves become smooth. They showed this on the arithmetic tasks and argued that emergence is often a property of the metric, not the model.

Where the debate stands, as best it can be summarized: loss and continuous proxies scale predictably, and that is what you plan a run on. Some downstream capabilities do look sharp in the metrics people care about, and whether you call that "emergent" depends on whether you insist on the metric that users see (exact answers) or the one that is smooth. Both camps agree that per-token loss improves smoothly, and that smooth improvement is what scaling laws promise.

Part B: Optimization

Scaling laws told you what to train. Now you have to train it, and the constraint is brutal: you get one shot, at a scale where you cannot tune anything. Every setting below is a default that was learned the hard way.

AdamW: the optimizer everyone uses

The problem plain gradient descent has with transformers: different parameters see gradients of wildly different scale. Embedding rows for rare tokens get tiny, infrequent gradients; attention output projections get large, dense ones. One learning rate cannot serve both. Adam (Kingma and Ba, 2015) fixes this by keeping, for every parameter, a running average of the gradient and a running average of its square, and dividing one by the root of the other. The result is that each parameter moves by roughly the learning rate per step, regardless of gradient scale.

What just happened, line by line. $m_t$ is momentum: an exponential moving average of gradients, which smooths noise. $v_t$ tracks the typical squared gradient size. The hats divide out the bias from starting both averages at zero (without them the first steps would be tiny). The update divides momentum by the root of the second moment, so a parameter with consistently large gradients and one with consistently small gradients both move by about $\eta$ per step. The last term, $\lambda\theta$, is decoupled weight decay (Loshchilov and Hutter, 2019): the "W" in AdamW. It pulls every weight toward zero by a fixed fraction each step, applied directly rather than through the adaptive scaling.

The hyperparameters for LLMs

The defaults you inherit from PyTorch ($\beta_2 = 0.999$, no weight decay) are wrong for LLM pre-training. The settings that have converged across GPT-3, Llama, and most open reports:

Learning-rate schedules

The problem: the right learning rate is not one number. Early in training, the weights are random and the gradient direction is unreliable; a big step can push the model into a bad region it never recovers from. In the middle, you want big steps to cover ground. At the end, you want small steps so the model settles into a minimum rather than bouncing around it. A schedule encodes all three.

Linear warmup, cosine decay

The standard recipe from GPT-3 onward: ramp the learning rate linearly from zero to its peak $\eta_{\max}$ over the first $W$ steps (GPT-3: 375M tokens; Llama: 2,000 steps), then decay it along a cosine curve to a floor $\eta_{\min}$ (typically 10% of peak) at the final step $S$.

The cosine is nothing magic; it is a smooth curve that spends a long time near the peak and a long time near the floor, with a fast transition between. Its weakness is that it needs to know $S$ in advance. Decide to train longer, and the schedule you already ran was wrong (it decayed too early). This is also exactly what tripped up the Kaplan scaling-law fits.

Warmup-stable-decay (WSD)

The alternative that has spread since 2024 (Hu et al., 2024 in MiniCPM; analyzed by Hägele et al., 2024): warm up, then hold the learning rate constant for most of training, then decay it quickly (linearly, or along a square-root-shaped curve) over the last 10 to 20 percent of steps. The constant phase is agnostic to total length: you can branch off a decay at any point and get a good model, or keep going. It also makes scaling-law experiments cheap, because one long stable run yields models at many token counts, each finished by a short decay. Loss during the stable phase sits above the cosine run's, and then drops sharply during decay to match or beat it.

Why warmup matters

Three reasons, all early-training. First, Adam's $v_t$ estimate is built from a handful of samples in the first steps, and dividing by an unreliable $\sqrt{\hat v}$ produces erratic step sizes; warmup keeps those steps small until the estimate settles. Second, at random initialization the loss surface is sharp in some directions (large gradient norms), and a full-size step overshoots; warmup lets the model move into a flatter region first. Third, attention logits and layer norms are unconstrained at initialization and can be driven into saturation by a few large updates. Skipping warmup on a large model reliably produces a spike in the first few hundred steps.

Batch size in tokens and the critical batch size

The problem: with thousands of GPUs you want a huge batch, because each GPU needs enough work to stay busy and because fewer, bigger steps means less communication per token. But does a bigger batch actually learn more per token? Not without limit. LLM batch sizes are quoted in tokens: GPT-3 used 3.2M tokens per batch; Llama 3 405B ramped from 4M to 16M tokens over the run.

McCandlish et al. (2018) gave the tool for thinking about this: the gradient noise scale. A mini-batch gradient is the true gradient plus noise whose variance shrinks as $1/B$. Averaging over more examples helps until the noise is small compared to the signal, and then it stops helping. The noise scale, approximately

where $G$ is the true gradient and $\Sigma$ the per-example covariance, estimates the batch size where that transition happens: the critical batch size. Below it, doubling the batch nearly halves the number of steps needed (perfect scaling). Above it, doubling the batch buys almost nothing; you are just spending more tokens per step. Their empirical rule ties the two limits together: if $S_{\min}$ is the minimum steps (with infinite batch) and $E_{\min}$ the minimum examples (with tiny batch), any run satisfies roughly $(S/S_{\min} - 1)(E/E_{\min} - 1) = 1$.

The noise scale grows as the loss falls: as the model gets better, the "signal" gradient gets smaller while the per-example noise stays, so larger batches become useful later in training. That is why Llama 3 and others ramp the batch size up in stages rather than fixing it.

Gradient clipping

The problem: occasionally a batch produces a gradient far larger than usual (a rare token pattern, a numerical hiccup). One such step with Adam is partly buffered by the $\sqrt{\hat v}$ division, but with $\beta_2 = 0.95$ the second moment adapts within a few steps and a burst of large gradients gets through. Clipping rescales the entire gradient vector if its global $L_2$ norm exceeds a threshold, almost always 1.0:

The direction is kept; only the magnitude is capped. Logging how often the clip fires is one of the most useful health signals in a run. If it fires on most steps, the learning rate is too high or something is wrong; if it never fires, it is not doing anything.

Loss spikes and their remedies

Every team that has trained a large model has stared at a loss curve that suddenly leapt upward, sometimes recovering, sometimes not. Spikes are the dominant failure mode of long runs, and by now the causes and fixes are reasonably well catalogued (Wortsman et al., 2023 is a good survey with small-scale reproductions).

μP and hyperparameter transfer

The problem: the optimal learning rate depends on model width, so a value tuned on a 100M-parameter proxy is wrong for the 100B target, and you cannot tune the target. Yang et al. (2022) showed that this dependence is an artifact of how we initialize and scale layers. In the standard parametrization, wider layers effectively get larger updates relative to their activations, so the safe learning rate shrinks with width. Their maximal update parametrization (μP) rescales initialization variances and per-layer learning rates with width in a specific way (for example, hidden-layer Adam learning rates scale as $1/\text{width}$) such that the optimal learning rate becomes independent of width. Tune once on a small model, transfer the values to the large one. It has been used for GPT-4-scale runs (reported), and MiniCPM and others document it in the open. The practical catch is that it must be implemented carefully in every layer, and depth transfer is less clean than width transfer.

Learning rate vs model size: rules of thumb

Without μP, you fall back on empirical fits. Kaplan et al. (2020) give one for standard-parametrization transformers: $\eta_{\text{opt}}(N) \approx 0.003239 - 0.0001395 \ln N$, which yields about $3.5\times 10^{-4}$ at 1B parameters; it was fit on models up to about a billion parameters and falls off too fast to extrapolate much beyond that. The GPT-3 table (Brown et al., 2020) is the reference most people actually copy: $6\times 10^{-4}$ for 125M, $2\times 10^{-4}$ for 1.3B, $1\times 10^{-4}$ for 13B, $6\times 10^{-5}$ for 175B. DeepSeek LLM (2024) fit both learning rate and batch size directly against compute, finding $\eta_{\text{opt}} \propto C^{-0.125}$ and $B_{\text{opt}} \propto C^{0.33}$. The direction is consistent everywhere: bigger model, smaller learning rate, bigger batch.

Companion code

code/lumen/optimizers.py implements SGD, momentum, Adam and AdamW from scratch as plain functions over parameter tensors, so you can read every line of the update. code/lumen/train.py has the warmup-cosine and WSD schedules, gradient clipping, and the training loop that logs the clip rate. The core of both fits on one screen.

import math, torch

def adamw_step(params, grads, state, lr, betas=(0.9, 0.95), eps=1e-8, wd=0.1):
    b1, b2 = betas
    state["t"] = state.get("t", 0) + 1
    t = state["t"]
    for i, (p, g) in enumerate(zip(params, grads)):
        m = state.setdefault(("m", i), torch.zeros_like(p))
        v = state.setdefault(("v", i), torch.zeros_like(p))
        m.mul_(b1).add_(g, alpha=1 - b1)              # m = b1*m + (1-b1)*g
        v.mul_(b2).addcmul_(g, g, value=1 - b2)       # v = b2*v + (1-b2)*g^2
        m_hat = m / (1 - b1 ** t)
        v_hat = v / (1 - b2 ** t)
        p.mul_(1 - lr * wd)                           # decoupled weight decay
        p.addcdiv_(m_hat, v_hat.sqrt().add_(eps), value=-lr)

def cosine_lr(step, warmup, total, peak, floor):
    if step < warmup:
        return peak * step / max(1, warmup)
    frac = (step - warmup) / max(1, total - warmup)
    return floor + 0.5 * (peak - floor) * (1 + math.cos(math.pi * frac))

def clip_grad_norm(grads, max_norm=1.0):
    total = math.sqrt(sum(float((g * g).sum()) for g in grads))
    scale = min(1.0, max_norm / (total + 1e-6))
    for g in grads:
        g.mul_(scale)
    return total  # log this: how often scale < 1 is a health signal

Practice

Further reading

• Kaplan et al. (2020). Scaling Laws for Neural Language Models. The original power laws, and the learning-rate fit in Appendix D.
• Hoffmann et al. (2022). Training Compute-Optimal Large Language Models. Chinchilla; the fitted constants come from Section 3.3 and Appendix D.
• Besiroglu et al. (2024). Chinchilla Scaling: A Replication Attempt. Why the published approach-3 constants disagree with the 20:1 rule, and refitted values.
• Sardana and Frankle (2023). Beyond Chinchilla-Optimal: Accounting for Inference in Language Model Scaling Laws. The over-training argument made quantitative.
• Wei et al. (2022). Emergent Abilities of Large Language Models. The case for emergence.
• Schaeffer et al. (2023). Are Emergent Abilities of Large Language Models a Mirage? The case that it is the metric.
• Kingma and Ba (2015). Adam: A Method for Stochastic Optimization. The optimizer.
• Loshchilov and Hutter (2019). Decoupled Weight Decay Regularization. AdamW.
• McCandlish et al. (2018). An Empirical Model of Large-Batch Training. The gradient noise scale and critical batch size.
• Hägele et al. (2024). Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations. WSD schedules and why they make scaling studies cheap.
• Hu et al. (2024). MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies. WSD and μP in a documented open run.
• Wortsman et al. (2023). Small-scale proxies for large-scale Transformer training instabilities. Reproduces spikes at small scale and tests the fixes.
• Yang et al. (2022). Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer. μP.
• DeepSeek-AI (2024). DeepSeek LLM: Scaling Open-Source Language Models with Longtermism. Fits of learning rate and batch size against compute.