Optimizer Lr Schedule Tutorial En

§0 TL;DR Cheat Sheet

1. Why not just plain SGD: real loss landscapes are ill-conditioned (Hessian condition number $\kappa=\lambda_{\max}/\lambda_{\min}$ huge), so plain gradient descent oscillates on steep directions and crawls on shallow ones. Two cures: momentum (accumulate along consistent directions, cancel oscillations) + per-parameter adaptive step (diagonal preconditioning, evening out per-coordinate curvature mismatch). The optimizer lineage is "SGD → +momentum → +adaptive → Adam (both)".
2. SGD / Momentum / Nesterov: plain SGD $\theta_t=\theta_{t-1}-\eta g_t$; heavy-ball momentum $v_t=\mu v_{t-1}+g_t,\ \theta_t=\theta_{t-1}-\eta v_t$ (damps oscillation, accelerates consistent directions); Nesterov evaluates the gradient at the look-ahead point $\theta-\eta\mu v$ (corrects overshoot). Momentum $\approx$ an exponential moving average of gradients with window $\sim 1/(1-\mu)$.
3. AdaGrad → RMSProp → Adam: AdaGrad's per-coordinate step $\propto 1/\sqrt{\sum g^2}$ (accumulates all history → as long as the gradient never stays zero forever ($\Sigma g^2\to\infty$, almost always true on long runs), LR monotonically → 0, "can't learn" on long training); RMSProp replaces the sum with an EMA (fixes the monotonic decay); Adam = RMSProp's second-moment EMA + first-moment momentum + bias correction.
4. Adam bias correction: $m,v$ start at 0, biased early (toward 0); dividing by $1-\beta^t$ de-biases. Since $\beta_2$ (0.999) is closer to 1 than $\beta_1$ (0.9), $v$ is biased more → without correction the early step is too LARGE (with default $\beta$ the first step $\approx\frac{1-\beta_1}{\sqrt{1-\beta_2}}=3.16$× $\eta$), correction returns the first step to $\eta$. Note $\epsilon$ is outside the sqrt (PyTorch).
5. AdamW decoupled weight decay: Adam's L2 regularization $\ne$ weight decay. L2 adds $\lambda\theta$ to the gradient → it gets scaled by $1/\sqrt{\hat v}$ → params with large historical gradients get decayed less (decoupling broken). AdamW (Loshchilov & Hutter) decouples: subtract $\eta\lambda\theta$ directly from the weights, outside the Adam step. This is the modern default; the two coincide for plain SGD (no momentum) but not for Adam (momentum-SGD with coupled L2 accumulates the decay term across the momentum buffer across steps, so it is no longer exactly equivalent to decoupled SGDW).
6. Frontier optimizers: Muon (orthogonalize the momentum of 2D weights via Newton-Schulz iterations, hidden 2D layers only, used by Kimi-K2), Lion (sign momentum, single state → half Adam's memory), Shampoo (Kronecker-factored full-matrix preconditioner), SOAP (run Adam in Shampoo's eigenbasis), Adafactor (factored second moment → sublinear memory), LAMB (layer-wise adaptive, huge batch), Sophia (light second-order, diagonal Hessian).
7. LR schedule: warmup (early $\hat v$ has high variance + large-batch/Post-LN instability → ramp up linearly); cosine (smooth decay to ~0, + warm restarts, needs the total step count up front); inverse-sqrt / Noam (the original Transformer schedule); WSD (warmup-stable-decay, a constant phase + a short final decay → no need to pre-commit the total step count, supports continual training and mid-run checkpoints); one-cycle (super-convergence).
8. Weight decay + LR-batch scaling + no-decay groups: wd does regularization (recently also a "wd $\approx$ effective-LR controller" view); the linear scaling rule (batch ×$k$ → LR ×$k$, SGD; Adam commonly $\sqrt k$); LayerNorm/RMSNorm gains, biases, embeddings get NO wd (split into decay/no-decay groups).
9. Gradient clipping + LLM hyperparameters: clip by global norm (preserves direction) vs by value (changes direction), to prevent gradient explosion / loss spikes; LLM pretraining defaults to $\beta_2=\mathbf{0.95}$ (not 0.999 — long memory is too sluggish for spikes), $\beta_1=0.9$, $\epsilon=\text{1e-8}$, wd $=0.1$, grad-clip $=1.0$.

§1 Why we need optimizers beyond SGD

Plain stochastic gradient descent (SGD)'s fundamental difficulty is that real loss landscapes are almost always "ill-conditioned." Expand the loss to second order near a minimum; the curvature is described by the Hessian $H$. Define the condition number $\kappa=\lambda_{\max}/\lambda_{\min}$ (ratio of $H$'s largest to smallest eigenvalue). Large $\kappa$ means hugely different curvature in different directions — some steep as a canyon wall, some flat as the valley floor. For a quadratic, gradient descent's convergence rate is about $\big(\tfrac{\kappa-1}{\kappa+1}\big)$, slower the larger $\kappa$ is.

Intuitively, plain SGD oscillates in an ill-conditioned valley: with a single global $\eta$, to avoid diverging on the steep direction (large $\lambda$) $\eta$ must be tiny ($\eta\lt 2/\lambda_{\max}$); but that tiny $\eta$ then crawls on the shallow direction (small $\lambda$). The result is zig-zagging back and forth between the canyon walls while the valley-floor direction you actually want to progress along barely moves. Deep-net loss landscapes routinely have $\kappa\sim 10^4$ or more, so plain SGD is both slow and unstable.

The two complementary cures correspond to the two axes of the optimizer lineage:

• Momentum: accumulate a velocity of historical gradients. On oscillating directions the gradient flips sign repeatedly and cancels within the velocity; on consistent directions the gradient keeps the same sign and accumulates, accelerating. Equivalent to smoothing the update direction, improving GD's convergence constant from $\kappa$ to $\sqrt\kappa$ scale ($\tfrac{\sqrt\kappa-1}{\sqrt\kappa+1}$). Note: $\kappa\to\sqrt\kappa$ is the classic acceleration result for tuned heavy-ball/Nesterov on a strongly-convex quadratic; real deep nets have stochastic-gradient noise, non-convexity, and time-varying curvature, so it's an intuition, not a guarantee.
• Per-parameter adaptive step (adaptive / per-parameter LR): scale each coordinate's step by its own gradient history — coordinates with persistently large gradients take small steps, persistently small ones take large steps. This acts as a diagonal preconditioner, evening out the per-coordinate effective curvature and directly easing the large-$\kappa$ disease.

The lineage in one line: SGD (first-order gradient only) → + momentum (add a first-moment EMA) → + adaptivity (add a second-moment EMA) → Adam = momentum + adaptivity + bias correction. §2–§5 follow this line; §6 is the frontier (matrix preconditioning / sign / memory-saving); §7–§9 are the accompanying LR schedules, weight decay, and clipping.

§2 SGD / Momentum / Nesterov

2.1　Plain SGD and heavy-ball momentum

Plain SGD looks only at the current mini-batch gradient $g_t$ each step:

$$\theta_t = \theta_{t-1} - \eta\, g_t.$$

Heavy-ball momentum (Polyak, 1964) introduces a velocity buffer $v$, giving the update "inertia":

$$\boxed{\;v_t = \mu\, v_{t-1} + g_t, \qquad \theta_t = \theta_{t-1} - \eta\, v_t\;}\qquad (\mu\ \text{typically }0.9).$$

(This is PyTorch's SGD(momentum=μ) form: $v$ starts at 0, the first step gives $v_1=g_1$, no $(1-\mu)$ damping.) Why it helps:

• Oscillating directions (steep walls): adjacent-step gradients flip sign repeatedly, partially cancelling in the exponentially-weighted sum $v_t=\mu v_{t-1}+g_t$ → the zig-zag is suppressed.
• Consistent directions (valley floor): adjacent-step gradients share a sign and keep accumulating → velocity grows, accelerating forward.

2.2　Nesterov accelerated gradient (NAG)

Heavy-ball evaluates the gradient at the current point $\theta_{t-1}$. Nesterov accelerated gradient instead evaluates at the look-ahead point "where the momentum will carry you":

$$v_t = \mu\, v_{t-1} + g\big(\theta_{t-1} - \eta\mu\, v_{t-1}\big), \qquad \theta_t = \theta_{t-1} - \eta\, v_t.$$

Intuition: "peek at where momentum is about to push me, then compute the gradient there" — if you're about to overshoot, the look-ahead gradient gives an early correction, so NAG has a better convergence constant on convex problems. Sutskever et al. (ICML 2013) showed empirically: carefully designed (Nesterov-style) momentum is crucial for training deep nets; good momentum + initialization can train deep nets even without an adaptive optimizer.

2.3　Momentum ≈ an exponential moving average of gradients

Unrolling the recurrence (with $v_0=0$):

$$v_t = \sum_{i=1}^{t} \mu^{\,t-i}\, g_i.$$

This is a geometric (exponential) weighted sum of historical gradients — more recent gradients weigh more, older ones decay as $\mu^{t-i}$. The effective averaging window is about $1/(1-\mu)$ ($\mu=0.9$ → ~10 steps, $\mu=0.99$ → ~100 steps). So momentum does something simple: average noisy instantaneous gradients into a smoother, more trustworthy descent direction, amplifying the consistent trend and washing out noise/oscillation.

§3 AdaGrad → RMSProp → Adam (the adaptive lineage)

3.1　AdaGrad: per-coordinate adaptive, but "learns to death"

AdaGrad (Duchi, Hazan & Singer, JMLR 2011) gives each coordinate an independent step size that adapts to its gradient history. Maintain the per-coordinate cumulative sum of squares $G_t$:

$$G_t = G_{t-1} + g_t^2 \ (\text{elementwise}), \qquad \theta_t = \theta_{t-1} - \frac{\eta}{\sqrt{G_t}+\epsilon}\odot g_t.$$

Effect: coordinates with persistently large gradients have large $G$ → small step; sparse/rare features have small $G$ → large step. Very effective on convex optimization and sparse features (early NLP, recommendation). Fatal flaw: $G_t$ is a monotonically increasing cumulative sum, so the denominator $\sqrt{G_t}$ only grows → the effective learning rate is monotonically non-increasing; as long as $\Sigma g_t^2\to\infty$ (the gradient doesn't stay exactly zero forever, almost always true on long runs), it tends to 0. On a deep net's long training run the step shrinks to nearly 0 partway through and it can't learn anymore (premature stagnation).

3.2　RMSProp: replace the cumulative sum with an EMA

RMSProp (Hinton, Coursera Neural Networks for ML, Lecture 6e, 2012, never formally published) gets straight to the point: the problem is "accumulating all history." Replace the cumulative sum with an exponential moving average, so old gradients are forgotten:

$$E[g^2]_t = \rho\, E[g^2]_{t-1} + (1-\rho)\, g_t^2, \qquad \theta_t = \theta_{t-1} - \frac{\eta}{\sqrt{E[g^2]_t}+\epsilon}\odot g_t.$$

The EMA denominator no longer grows unboundedly (it tracks the "recent" gradient-square scale), so the effective LR doesn't decay to 0 — AdaGrad's deadlock is undone. This is the direct precursor of Adam's second-moment term.

3.3　Adam = RMSProp + momentum + bias correction

Adam (Kingma & Ba, 2014) merges the two axes: RMSProp's second-moment EMA (adaptive step) + heavy-ball's first-moment EMA (momentum) + a crucial bias correction (§4). One-line mnemonic:

$$\textbf{Adam} = \underbrace{\text{RMSProp}}_{\text{second-moment EMA}} + \underbrace{\text{momentum}}_{\text{first-moment EMA}} + \underbrace{\text{bias correction}}_{\text{remove early bias}}.$$

§4 Adam in depth

4.1　The full update

$$m_t = \beta_1 m_{t-1} + (1-\beta_1)\, g_t \qquad (\text{first moment / momentum, EMA}),$$

$$v_t = \beta_2 v_{t-1} + (1-\beta_2)\, g_t^2 \qquad (\text{second moment / adaptive, EMA}),$$

$$\hat m_t = \frac{m_t}{1-\beta_1^{\,t}}, \qquad \hat v_t = \frac{v_t}{1-\beta_2^{\,t}} \qquad (\textbf{bias correction}),$$

$$\boxed{\;\theta_t = \theta_{t-1} - \eta\,\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon}\;}\qquad (\epsilon\ \text{outside the sqrt, PyTorch}).$$

Defaults $\beta_1=0.9,\ \beta_2=0.999,\ \epsilon=10^{-8}$. $m_0=v_0=0$.

4.2　Why bias correction is needed

$m,v$ both initialize at 0. Early on they are biased estimates of the true moments — biased toward 0, and the closer $\beta$ is to 1 the longer the bias lasts. Taking the expectation for stationary gradients: $\mathbb{E}[m_t]=(1-\beta_1^t)\,\mathbb{E}[g]$, so $\hat m_t=m_t/(1-\beta_1^t)$ exactly de-biases; $\hat v_t$ likewise.

The key point, also frequently probed deeply: without correction, the early effective step is not too small, but too large. Look at the first step ($t=1$):

$$\frac{m_1}{\sqrt{v_1}} = \frac{(1-\beta_1)\,g_1}{\sqrt{(1-\beta_2)}\,\lvert g_1\rvert} = \frac{1-\beta_1}{\sqrt{1-\beta_2}}\,\mathrm{sign}(g_1).$$

Plugging in defaults: $\frac{1-0.9}{\sqrt{1-0.999}}=\frac{0.1}{0.0316}\approx 3.16$. That is, the uncorrected first step is about 3.16× the target $\eta$ — because $v$ (with $\beta_2$ closer to 1) is pushed toward 0 harder than $m$, the denominator $\sqrt{v}$ is too small, which inflates the step instead. Bias correction restores $\hat m_1=g_1,\ \hat v_1=g_1^2$, returning the first step to $\eta\cdot\dfrac{g_1}{\lvert g_1\rvert+\epsilon}\approx\eta\cdot\mathrm{sign}(g_1)$ (the strict form includes $\epsilon$: $\eta\lvert g_1\rvert/(\lvert g_1\rvert+\epsilon)<\eta$, which only approaches $\eta$ when $\lvert g_1\rvert\gg\epsilon$ — training-time gradients are typically $\gg 10^{-8}$, so this approximation holds well in practice). (§A's [d] verifies this 3.16 numerically.)

§5 AdamW (the most-probed section)

5.1　Adam's L2 regularization $\ne$ weight decay

The classic L2 regularization approach: add $\frac{\lambda}{2}\lVert\theta\rVert^2$ to the loss, equivalent to adding $\lambda\theta$ to the gradient:

$$g_t' = g_t + \lambda\,\theta_{t-1}.$$

For plain SGD this is exactly weight decay: $\theta_t=\theta_{t-1}-\eta(g_t+\lambda\theta_{t-1})=(1-\eta\lambda)\theta_{t-1}-\eta g_t$ — each step shrinks the weight by $(1-\eta\lambda)$. But for Adam, disaster strikes: this $\lambda\theta$ is stuffed into $g'$, so it flows into $m,v$ and ends up scaled by $1/(\sqrt{\hat v}+\epsilon)$:

$$\theta_t = \theta_{t-1} - \eta\,\frac{\widehat{m'_t}}{\sqrt{\hat v_t}+\epsilon},\qquad m'\ \text{contains the }\lambda\theta\ \text{contribution}.$$

Consequence: params with large historical gradients ($\hat v$ large) get LESS effective decay, params with small gradients get more — the decay strength is distorted by $\hat v$, completely divorced from the regularizer's intent of "applying a uniform shrinkage to all weights." This is "L2 and weight decay are no longer equivalent under an adaptive optimizer."

5.2　AdamW: decoupled weight decay

AdamW (Loshchilov & Hutter, Decoupled Weight Decay Regularization, 1711.05101) fixes it minimally — take weight decay out of the gradient, bypass the adaptive denominator, and apply it directly to the weights:

$$\boxed{\;\theta_t = \theta_{t-1} - \eta\Big(\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon} + \lambda\,\theta_{t-1}\Big)\;}$$

i.e. compute the Adam step on the raw gradient $g$ (without $\lambda\theta$), then separately subtract an $\eta\lambda\theta_{t-1}$. Equivalently $\theta_t=(1-\eta\lambda)\theta_{t-1}-\eta\,\hat m_t/(\sqrt{\hat v_t}+\epsilon)$ — a clean, $\hat v$-independent uniform shrinkage (PyTorch's AdamW multiplies the decay by lr).

Why it's usually better (more controllable): after decoupling, all params receive the same relative shrinkage $\lambda$, the regularization strength is no longer distorted by each param's gradient history, more faithful to "controlling the weight norm"; Loshchilov & Hutter showed empirically that AdamW is often stably better than Adam+L2 across tasks, and the optimal LR and wd are easier to tune decoupled (no need to absolutize it to "more generalizing on all tasks," but it is the more controllable, more common default). Why plain SGD makes them coincide: plain SGD has neither the $1/\sqrt{\hat v}$ adaptive scaling layer nor a momentum buffer that accumulates $\lambda\theta$ across steps, so $\lambda\theta$ added to the gradient is used as-is, which is exactly uniform shrinkage; once you introduce an adaptive denominator (Adam) or a momentum buffer (SGD+momentum — the SGDW-vs-coupled-L2+momentum distinction), adding to the gradient (L2) and shrinking directly (decay) part ways.

§6 Frontier optimizers

Remember just one core idea for each — being able to state "which part of Adam it changes" in an interview is enough.

6.1　Muon: orthogonalize the momentum

Muon (Keller Jordan, 2024 tech blog) targets 2D weight matrices: first compute the update matrix $M$ from momentum, then use a few Newton-Schulz iterations to orthogonalize it into $O\approx\mathrm{NewtonSchulz}(M)$ (mathematically approximating $UV^\top$ of $M=U\Sigma V^\top$, i.e. the nearest semi-orthogonal matrix, replacing an expensive SVD with iteration), and update the weight with $O$. Intuition: orthogonalization makes the update push equally on all singular directions, not dominated by a few large singular values → more balanced updates, faster effective learning. Only for hidden-layer 2D matrices; embeddings, the output head, norm gains, biases still go through AdamW (orthogonalization is meaningless for them). Moonshot's Kimi-K2 trains with Muon; for scalability see Muon is Scalable for LLM Training (2502.16982).

6.2　Lion: sign momentum, single state

Lion (Chen et al., Symbolic Discovery of Optimization Algorithms, 2302.06675) was discovered automatically by symbolic program search. The update takes the sign of the momentum: $u_t=\mathrm{sign}(\beta_1 m_{t-1}+(1-\beta_1)g_t)$, $\theta_t=\theta_{t-1}-\eta(u_t+\lambda\theta_{t-1})$, then updates the momentum $m_t=\beta_2 m_{t-1}+(1-\beta_2)g_t$. It stores only 1 state (momentum) → optimizer memory is half Adam's (which stores both $m,v$). The sign makes every coordinate take an equal-magnitude step; usually paired with a smaller LR and a larger wd.

6.3　Shampoo / SOAP: full-matrix preconditioning

• Shampoo (Gupta, Koren & Singer, 1802.09568): goes beyond Adam's diagonal preconditioning to full-matrix / Kronecker-factored preconditioning. For a weight matrix it maintains left and right preconditioners $L=\sum GG^\top$, $R=\sum G^\top G$, and updates with $L^{-1/4} G\, R^{-1/4}$ — approximating full-matrix AdaGrad with a Kronecker product, capturing cross-coordinate coupling. Converges faster but each step is expensive (computing matrix inverse-roots).
• SOAP (Vyas et al., 2409.11321): runs Adam in Shampoo's (slowly-changing) eigenbasis — rotate the gradient into Shampoo's preconditioner's eigenspace, do standard Adam there, rotate back. It connects Shampoo with Adam/Adafactor, more stable with fewer hyperparameters.

6.4　Adafactor / LAMB / Sophia

• Adafactor (Shazeer & Stern, 1804.04235): factorize the second-moment matrix $v\in\mathbb{R}^{m\times n}$ into row and column vectors (rank-1 reconstruction $v\approx rc^\top$) → memory drops from $O(mn)$ to $O(m+n)$ (sublinear). T5 used it to train a large model on limited memory.
• LAMB (You et al., 1904.00962): on top of Adam's update direction, apply a further layer-wise adaptation — scale by each layer's $\lVert\theta\rVert/\lVert\text{update}\rVert$ (trust ratio) (the Adam version of LARS). Stabilizes very large batch training (BERT in 76 minutes).
• Sophia (Liu et al., 2305.14342): light second-order. Uses a diagonal Hessian estimate (Gauss-Newton-Bartlett or Hutchinson) for preconditioning + clipping the update, updating the Hessian only every $k$ steps → converges faster than AdamW with the second-order cost amortized.

§7 Learning-rate schedules

7.1　Warmup: why start slow

Warmup linearly ramps the LR from 0 (or very small) up to the peak over the first few hundred ~ few thousand steps, then hands off to the subsequent decay. Two motivations:

1. Adam's early variance is high: $\hat v$ is estimated from very few samples in the first steps — even though bias correction guarantees $\mathbb{E}[\hat v_t]$ is unbiased (§4.2), any single realization of $\hat v_t$ is still a high-variance estimate, so on individual steps $\hat v_t$ can fall well below the true second moment, making $\hat m/\sqrt{\hat v}$ swing wildly and occasionally spike into an oversized effective step (the core motivation behind RAdam, Liu et al. 2019). This is a different issue from the "uncorrected first step is systematically 3.16× too large" discussed in §4.2 — that's a deterministic bias standard Adam already cancels exactly via bias correction, not the real source of early instability. Use a small LR first to let $m,v$'s statistics "warm up" to reliability, then scale up the step.
2. Large-batch / Post-LN instability: under mean-reduction the batch-averaged gradient is unbiased ($\mathbb{E}[g_B]=\mathbb{E}[g]$); enlarging the batch only lowers the variance ($\mathrm{Var}\propto 1/B$), it doesn't inflate the gradient's scale itself. The real source of instability is that a larger batch is paired with a larger LR via the linear scaling rule (§8.2), and early on the weights are far from converged and change fast, so a large LR diverges more easily; on top of that, Post-LN Transformers have naturally large top-layer gradients (see the normalization tutorial, §5) — the two stack, so hitting the network with a large LR right away easily diverges. Warmup suppresses early instability.

7.2　Cosine annealing (+ warm restarts)

Cosine (Loshchilov & Hutter, SGDR, 1608.03983) decays the LR smoothly with a cosine from the peak to $\sim\eta_{\min}$:

$$\eta_t = \eta_{\min} + \tfrac12(\eta_{\max}-\eta_{\min})\Big(1+\cos\big(\pi\, t/T\big)\Big).$$

Early large LR explores quickly, late small LR refines convergence (less noise, easier to settle into a flatter minimum). SGDR also adds warm restarts: periodically pull the LR back to the peak to help escape local solutions. Cost: you must fix the total step count $T$ in advance ($\cos$ needs to know the endpoint) — exactly the pain point WSD solves.

7.3　Inverse-sqrt / Noam (the original Transformer schedule)

The schedule the original Transformer (Vaswani et al., 1706.03762) used, first a linear warmup then a $1/\sqrt t$ decay:

$$\eta_t = d_{\text{model}}^{-1/2}\cdot\min\!\Big(t^{-1/2},\ t\cdot t_{\text{warmup}}^{-3/2}\Big).$$

For $t\lt t_{\text{warmup}}$ the second term applies (linear rise), after that the first (decaying as $t^{-1/2}$). No need to pre-fix the total step count; it was the standard for early machine translation / language models.

7.4　WSD: Warmup-Stable-Decay

WSD (MiniCPM, Hu et al., 2404.06395) splits the schedule into three phases: warmup → a long constant phase (stable, at the peak LR) → a short final decay. The biggest selling point: no need to pre-commit the total step count — the constant phase can extend indefinitely, and you append a short decay only when you want to stop (loss drops sharply during the decay phase to a good level). Benefits:

• supports continual training / adding data (the constant phase can resume anytime);
• each decay yields a usable checkpoint, so you don't train a separate cosine run for each target step count;
• makes it easy to explore "how long to train" under a fixed compute budget, instead of betting on $T$ once and for all.

7.5　One-cycle / super-convergence

One-cycle (Smith, Super-Convergence, 1708.07120): the whole training is one cycle — the LR rises to a very large peak then falls to well below the initial value, while the momentum cycles inversely in sync (low momentum when LR is high). In some settings it achieves "super-convergence" (training very fast with a very large LR). Like cosine and WSD it's in the "rise then fall" family, distinguished by the single cycle + the very large peak LR.

🎯 Picking a schedule in one line Fixed budget, chasing peak quality → cosine (but pre-fix the step count); uncertain how long to train / want continual training / want mid-run checkpoints → WSD; classic Transformer reproduction → Noam; want to train super-fast with a very large LR → one-cycle. All of these should be preceded by warmup.

§8 Weight decay + LR-batch scaling + no-decay groups

8.1　What weight decay does (incl. a modern view)

Classic view: weight decay is regularization — each step shrinks the weight toward 0, bounding the weight norm, limiting model complexity, resisting overfitting. Modern view (for the part of the network's weights that are scale-invariant): a weight matrix $W$ that feeds directly and exclusively into a normalization layer, with no intervening scale-sensitive path such as a residual add (e.g. the linear-layer weights immediately followed by LayerNorm/RMSNorm inside a Pre-LN block) is scale-invariant — multiply $W$ by a positive constant and the normalized output is unchanged; that weight's "absolute size" doesn't directly affect the function, so wd's real role there is more like tuning the effective learning rate / controlling the weight-norm equilibrium — larger wd, smaller weight norm, relatively larger effective gradient step. Caveat: this argument does not apply to weights on the residual branch, the norm layer's own $\gamma/\beta$, biases, or the output head — their scale directly affects the network's output, so they are not scale-invariant. This is an important observation in recent training-dynamics research (treating wd as an "effective-LR knob" rather than pure regularization, but only for the part of the weights that are actually scale-invariant).

8.2　Scaling LR with batch size

In large-batch training, LR must scale with the batch, else you waste compute or destabilize:

• Linear scaling rule (Goyal et al., 1706.02677): batch ×$k$ → LR ×$k$ (for large-batch SGD), paired with warmup. Intuition: the mini-batch mean gradient's variance $\propto 1/B$, so enlarging the batch $k$× reduces the noise, letting you safely enlarge the step $k$× to keep "the effective update each sample contributes" roughly constant.
• Adam's $\sqrt k$ scaling: Adam's update is approximately normalized by $\sqrt{\hat v}$ (close to sign), so empirically scaling the batch by a square root (LR ×$\sqrt k$) is more stable than linear.
• Relatedly, LARS (You et al., 1708.03888) / LAMB (§6.4) use a per-layer trust ratio to achieve adaptive scaling under very large batches.

8.3　No-decay parameter groups

Not all parameters should get weight decay. Standard practice splits params into two groups:

• with wd: each linear / conv layer's weight matrix.
• without wd (no-decay): LayerNorm/RMSNorm gains $\gamma$, all biases, and (usually) embeddings.

Reason: these are scale / shift parameters, and shrinking them toward 0 directly harms representation — pushing an LN gain toward 0 wipes out that layer's activation scale; a bias shouldn't be regularized at all (it only handles shifts). Codebases like GPT/LLaMA explicitly construct decay / no-decay param groups.

§9 Gradient clipping + LLM hyperparameter details

9.1　Gradient clipping: global norm vs by value

Occasional gradient explosions / loss spikes during training (especially RNNs, long sequences, large LR) can blow the weights away in one step or even produce NaN. Gradient clipping (Pascanu, Mikolov & Bengio, 1211.5063) puts a gate on the gradient:

• Clip by global norm (most common): concatenate all params' gradients into one big vector, compute its $\lVert g\rVert_2$; if it exceeds a threshold $c$, scale the whole thing $g\leftarrow g\cdot c/\lVert g\rVert_2$. Preserves the direction, only caps the magnitude.
• Clip by value: clamp each component to $[-v, v]$. This changes the gradient direction (different components are truncated differently), generally worse than global norm.

Why it stabilizes: occasional huge gradients are clamped within the threshold, the single-step update magnitude is bounded → one bad batch can't crash training. LLM pretraining almost always uses global-norm clip with threshold 1.0.

9.2　Adam hyperparameter defaults for LLM training

The Adam/AdamW recipe for LLM pretraining differs from the "textbook defaults" in a few key places:

§10 Picking an optimizer + the memory bill

10.1　Optimizer-state memory

In mixed-precision training, each trainable parameter costs roughly: bf16 weight 2 B + bf16 gradient 2 B + fp32 master 4 B + optimizer state. The optimizer state is exactly the watershed between methods:

• SGD+momentum: 1 state (the momentum $v$).
• Adam / AdamW: 2 states ($m, v$), $=8$ B/param in fp32 — often the big chunk, the source of the "save memory" motivation.
• Lion: 1 state (momentum) → half of Adam.
• Adafactor: factorized $v$ → sublinear, approximately leaving only "param-level" cost.

The memory-saving routes this spawns:

• 8-bit Adam (Dettmers et al., 2110.02861): block-quantize $m,v$ to 8-bit storage (dequantize on use), dropping optimizer-state memory to about $1/4$, with nearly no quality loss.
• Adafactor: factorized second moment, sublinear memory (T5).
• Lion: one fewer state, directly halved.

10.2　When to use which

🎯 The "which optimizer" decision tree

• Transformer / LLM fine-tuning or pretraining: AdamW (safe default, $\beta_2=0.95$, wd 0.1, grad-clip 1.0, with warmup+cosine/WSD).
• Classic vision (ResNet, etc.): SGD+momentum often generalizes better and saves memory — much CV SOTA still uses it.
• Want extra training speedup / large scale: Muon (2D hidden layers), Shampoo/SOAP (full-matrix preconditioning, if you'll pay the per-step compute).
• Memory-constrained: 8-bit Adam / Adafactor / Lion.
• Very large batch: LAMB / LARS (per-layer trust ratio).

§11 Engineering comparison + common misconceptions

11.1　Optimizer comparison table

11.2　Common footguns

§12 25 high-frequency interview questions

Sorted into three tiers; click for answer points + easy traps. L2/L3 are top-lab deep water (AdamW decoupling, bias-correction direction, $\beta_2$ choice, scaling rules, Muon/Shampoo/Lion/Sophia, etc.).

L1 must-know

L2 advanced

L3 advanced

§A Appendix: sanity check

This tutorial's from-scratch implementations should satisfy the following key invariants (verifiable with a short pure-PyTorch script, a few seconds on CPU):

1. [a] from-scratch SGD+momentum == torch.optim.SGD: with $v_t=\mu v_{t-1}+g_t,\ \theta\mathrel{-}=\eta v_t$ (PyTorch convention: $v$ init 0, no damping), after a few steps the two are elementwise equal within float error (atol≈1e-5).
2. [b] from-scratch Adam == torch.optim.Adam: $\hat m=m/(1-\beta_1^t),\ \hat v=v/(1-\beta_2^t),\ \theta\mathrel{-}=\eta\hat m/(\sqrt{\hat v}+\epsilon)$ ($\epsilon$ outside the sqrt) should match PyTorch elementwise.
3. [c] AdamW $\ne$ Adam+L2, and from-scratch decoupled decay == torch.optim.AdamW: at the same weight_decay=λ, torch.optim.AdamW (decoupled) and torch.optim.Adam(weight_decay=λ) (L2) produce different updates ($\lVert\Delta\rVert\gt 0$); the from-scratch "Adam step first, then subtract $\eta\lambda\theta$" decoupled implementation should match AdamW. This is the core of §5.
4. [d] bias correction returns the first step approximately to $\eta$ (when $|g_1|\gg\epsilon$; the strict form is $\eta|g_1|/(|g_1|+\epsilon)<\eta$): at $t=1$ the uncorrected effective step $\approx 3.16\eta$ ($=\frac{1-\beta_1}{\sqrt{1-\beta_2}}$, because $v$ is pushed toward 0 harder by $\beta_2$ than $m$ → denominator too small → step too large), corrected it returns approximately to $\eta$ (again requiring $|g_1|\gg\epsilon$). Verify the corrected/uncorrected ratio $=\frac{\sqrt{1-\beta_2}}{1-\beta_1}\approx 0.316$.
5. [e] cosine-with-warmup schedule shape: $\eta(0)\approx 0$, the peak $\eta_{\max}$ is hit at the warmup boundary, decays to $\sim\eta_{\min}$ at the end, and is monotonically non-increasing after warmup.
6. [f] momentum accelerates an ill-conditioned quadratic: on a quadratic with large condition number $\kappa$, using a small $\eta$ inside the stability bound for $N$ steps, GD+momentum's final loss is markedly lower than plain GD's (momentum accumulates and accelerates on the shallow direction).

Below are a few illustrative snippets (CPU-runnable, English comments, with shape/numeric notes; they demonstrate several of the invariants — the full 6 of [a]–[f] are in the standalone script below).

[a]/[b] SGD+momentum and Adam from scratch, aligned with torch.optim:

import torch

def sgd_momentum_from_scratch(params, grads, lr=0.1, mu=0.9, steps=5):
    """heavy-ball: v = mu*v + g; theta -= lr*v. v init 0 (first step v=g), aligned with PyTorch SGD."""
    theta = params.clone()                      # [d] params
    v = torch.zeros_like(theta)                 # [d] momentum buffer
    for t in range(steps):
        g = grads[t]                            # [d] step-t gradient (a fixed sequence here, for alignment)
        v = mu * v + g                          # velocity = momentum*old velocity + current gradient
        theta = theta - lr * v                  # update along the velocity
    return theta

def adam_from_scratch(grads, theta0, lr=0.1, b1=0.9, b2=0.999, eps=1e-8, steps=5):
    """EMA of m,v + bias correction + eps outside the sqrt, aligned elementwise with PyTorch Adam."""
    theta = theta0.clone()
    m = torch.zeros_like(theta)                 # first moment [d]
    v = torch.zeros_like(theta)                 # second moment [d]
    for t in range(1, steps + 1):
        g = grads[t - 1]                        # [d]
        m = b1 * m + (1 - b1) * g               # momentum EMA
        v = b2 * v + (1 - b2) * g * g           # squared-gradient EMA
        m_hat = m / (1 - b1 ** t)               # bias correction (remove early bias toward 0)
        v_hat = v / (1 - b2 ** t)
        theta = theta - lr * m_hat / (v_hat.sqrt() + eps)   # eps outside the sqrt
    return theta

# Align with torch.optim (feed the same fixed gradient sequence to both):
torch.manual_seed(0)
d, steps = 8, 5
theta0 = torch.randn(d)
grads = [torch.randn(d) for _ in range(steps)]   # fixed gradient sequence, ensures comparability

# --- SGD ---
p = theta0.clone().requires_grad_(True)
opt = torch.optim.SGD([p], lr=0.1, momentum=0.9)
for t in range(steps):
    opt.zero_grad(); p.grad = grads[t].clone(); opt.step()
ref_sgd = p.detach()
mine_sgd = sgd_momentum_from_scratch(theta0, grads, lr=0.1, mu=0.9, steps=steps)
assert torch.allclose(mine_sgd, ref_sgd, atol=1e-6)    # from-scratch == torch.optim.SGD

[d] bias correction: uncorrected first step too large (at default $\beta$ $\approx 3.16\eta$), corrected $=\eta$:

import torch

b1, b2, eps = 0.9, 0.999, 1e-8
g1 = torch.tensor([2.0])                          # any nonzero first-step gradient
m1 = (1 - b1) * g1                                # uncorrected first moment
v1 = (1 - b2) * g1 * g1                           # uncorrected second moment
uncorrected = (m1 / (v1.sqrt() + eps)).item()     # uncorrected effective step / lr
corrected = ((m1 / (1 - b1)) / ((v1 / (1 - b2)).sqrt() + eps)).item()  # after correction
ratio = (1 - b1) / (1 - b2) ** 0.5                # = 3.162..., the theoretical first-step inflation
# invariant: uncorrected ≈ ratio ≈ 3.16 (too large), corrected ≈ 1.0 (= lr scale)
assert abs(uncorrected - ratio) < 1e-3
assert abs(corrected - 1.0) < 1e-3

[c] AdamW (decoupled decay) $\ne$ Adam+L2 (coupled), from-scratch decoupled == torch.optim.AdamW:

import torch

def adamw_decoupled_from_scratch(grads, theta0, lr=0.1, b1=0.9, b2=0.999,
                                 eps=1e-8, wd=0.05, steps=5):
    """Decoupled: do the Adam step on the raw gradient g, then subtract lr*wd*theta directly (bypassing the adaptive denom)."""
    theta = theta0.clone()
    m = torch.zeros_like(theta); v = torch.zeros_like(theta)
    for t in range(1, steps + 1):
        g = grads[t - 1]                          # raw gradient without wd
        m = b1 * m + (1 - b1) * g
        v = b2 * v + (1 - b2) * g * g
        m_hat = m / (1 - b1 ** t); v_hat = v / (1 - b2 ** t)
        theta = theta - lr * (m_hat / (v_hat.sqrt() + eps) + wd * theta)  # decoupled decay
    return theta

torch.manual_seed(0)
d, steps, wd = 8, 5, 0.05
theta0 = torch.randn(d)
grads = [torch.randn(d) for _ in range(steps)]

# torch.optim.AdamW (decoupled) vs torch.optim.Adam(weight_decay=) (L2/coupled)
def run(opt_cls, **kw):
    p = theta0.clone().requires_grad_(True)
    opt = opt_cls([p], lr=0.1, betas=(0.9, 0.999), eps=1e-8, weight_decay=wd, **kw)
    for t in range(steps):
        opt.zero_grad(); p.grad = grads[t].clone(); opt.step()
    return p.detach()

ref_adamw = run(torch.optim.AdamW)
ref_adam_l2 = run(torch.optim.Adam)
assert (ref_adamw - ref_adam_l2).norm() > 1e-4         # AdamW != Adam+L2 (differ at the same wd)
mine = adamw_decoupled_from_scratch(grads, theta0, wd=wd, steps=steps)
assert torch.allclose(mine, ref_adamw, atol=1e-6)      # from-scratch decoupled == torch.optim.AdamW

[e] cosine-with-warmup schedule (check first/peak/end + monotonicity):

import math

def cosine_warmup_lr(step, warmup, total, lr_max, lr_min=0.0):
    """Linear warmup to lr_max, then cosine decay to lr_min. Needs the total step count `total` up front."""
    if step < warmup:
        return lr_max * step / warmup                       # linear rise (step=0 -> ~0)
    progress = (step - warmup) / (total - warmup)           # 0 -> 1
    return lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * progress))

warmup, total, lr_max = 100, 1000, 1.0
lrs = [cosine_warmup_lr(t, warmup, total, lr_max) for t in range(total + 1)]
assert lrs[0] < 1e-9                                        # start ~0
assert abs(lrs[warmup] - lr_max) < 1e-9                     # peak hit at the warmup boundary
assert lrs[-1] < 1e-6                                       # decays to ~lr_min(=0) at the end
assert all(lrs[i] >= lrs[i + 1] - 1e-12 for i in range(warmup, total))  # monotonically non-increasing after warmup

The full runnable script is at code/optimizer_lr_schedule.py (pure PyTorch, a few seconds on CPU, 6 asserts covering [a]–[f]). The real output of running python code/optimizer_lr_schedule.py:

[a] SGD+momentum from scratch vs torch.optim.SGD: max|Δ| = 1.19e-07  OK
[b] Adam from scratch vs torch.optim.Adam: max|Δ| = 1.19e-07  OK
[c] AdamW vs Adam+L2 (same wd=0.1): max|Δ| = 2.13e-01 (DIFFER, decoupled≠coupled); from-scratch decoupled vs AdamW = 2.16e-07 (MATCH)  OK
[d] bias correction at step 1: corrected/uncorrected = 0.316 (= √(1-β2)/(1-β1) ≈ 0.316); without it the under-estimated v̂ makes the first step ~3.2× too large  OK
[e] cosine+warmup lr: t=0 -> 0.000, t=warmup -> 1.000 (peak), t=end -> 0.000 (min); linear warmup midpoint -> 0.500  OK
[f] ill-conditioned quadratic (κ=20) after 100 steps: GD loss = 4.04e-02  vs  GD+momentum loss = 1.46e-05  (momentum lower)  OK

all optimizer / LR-schedule sanity checks passed ✓

📚 References

• Adam — Kingma & Ba, Adam: A Method for Stochastic Optimization, arXiv 1412.6980 (2014), ICLR 2015.
• AdamW (Decoupled Weight Decay) — Loshchilov & Hutter, Decoupled Weight Decay Regularization, arXiv 1711.05101 (2017), ICLR 2019.
• AdaGrad — Duchi, Hazan & Singer, Adaptive Subgradient Methods for Online Learning and Stochastic Optimization, JMLR 12 (2011) (no arXiv id, cited by JMLR venue).
• RMSProp — Hinton, Neural Networks for Machine Learning, Coursera Lecture 6e (2012) (course lecture, no formal paper / arXiv id).
• Nesterov accelerated gradient — Nesterov, A method of solving a convex programming problem with convergence rate $O(1/k^2)$, Soviet Math. Doklady 27 (1983) (no arXiv id).
• Momentum for deep nets — Sutskever, Martens, Dahl & Hinton, On the importance of initialization and momentum in deep learning, ICML 2013 (PMLR proceedings; no standalone arXiv id, cited by venue).
• AMSGrad (Adam convergence) — Reddi, Kale & Kumar, On the Convergence of Adam and Beyond, arXiv 1904.09237 (2019), ICLR 2018.
• SGDR / cosine annealing — Loshchilov & Hutter, SGDR: Stochastic Gradient Descent with Warm Restarts, arXiv 1608.03983 (2016), ICLR 2017.
• Transformer / Noam schedule — Vaswani et al., Attention Is All You Need, arXiv 1706.03762 (2017), NeurIPS 2017.
• Linear scaling rule — Goyal et al., Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour, arXiv 1706.02677 (2017).
• LARS — You, Gitman & Ginsburg, Large Batch Training of Convolutional Networks, arXiv 1708.03888 (2017).
• LAMB — You et al., Large Batch Optimization for Deep Learning: Training BERT in 76 minutes, arXiv 1904.00962 (2019), ICLR 2020.
• Lion — Chen et al., Symbolic Discovery of Optimization Algorithms, arXiv 2302.06675 (2023), NeurIPS 2023.
• Shampoo — Gupta, Koren & Singer, Shampoo: Preconditioned Stochastic Tensor Optimization, arXiv 1802.09568 (2018), ICML 2018.
• SOAP — Vyas et al., SOAP: Improving and Stabilizing Shampoo using Adam, arXiv 2409.11321 (2024).
• Adafactor — Shazeer & Stern, Adafactor: Adaptive Learning Rates with Sublinear Memory Cost, arXiv 1804.04235 (2018), ICML 2018.
• Sophia — Liu et al., Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training, arXiv 2305.14342 (2023), ICLR 2024.
• One-cycle / Super-Convergence — Smith & Topin, Super-Convergence: Very Fast Training of Neural Networks Using Large Learning Rates, arXiv 1708.07120 (2017).
• WSD / MiniCPM — Hu et al., MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies, arXiv 2404.06395 (2024), COLM 2024.
• 8-bit Adam — Dettmers et al., 8-bit Optimizers via Block-wise Quantization, arXiv 2110.02861 (2021), ICLR 2022.
• Gradient clipping — Pascanu, Mikolov & Bengio, On the difficulty of training Recurrent Neural Networks, arXiv 1211.5063 (2012), ICML 2013.
• Muon (scaling) — Liu et al. (Moonshot AI), Muon is Scalable for LLM Training, arXiv 2502.16982 (2025).
• Muon (original writeup) — Keller Jordan et al., Muon: An optimizer for the hidden layers of neural networks, 2024 tech blog (no standalone arXiv paper, cited as a blog writeup).