Optimization

The Loss Landscape

Training a neural network means choosing parameters θ that make a loss L(θ) small. Over millions of parameters, L defines a high-dimensional surface, and an optimizer decides how to move across it, one step at a time.

The surface has saddle points, long narrow valleys and minima of different widths. The first sections compare gradient descent, momentum and Adam on one narrow valley; the later ones turn to effects observed when training real networks. Chapter 8 of Goodfellow, Bengio & Courville (2016) covers the classical methods in more depth.

Gradient Descent

The simplest rule takes a small step against the gradient, the direction in which the loss rises fastest. The step size η is the learning rate:

The figure runs this rule on the Rosenbrock function, whose minimum at (1, 1) lies at the bottom of a long, curved, narrow valley. (The button is labelled SGD, but here the gradient is computed from the function itself, not estimated from a random sample, so there is no sampling noise.) At the default η = 0.001 each step overshoots across the steep valley walls, so the path zig-zags from wall to wall while creeping along the floor. From about η = 0.0027 this run diverges.

Momentum

Polyak's heavy-ball method adds a velocity: each step keeps a fraction β of the previous step and adds the new gradient step, like a heavy ball rolling through the landscape with friction (Polyak, 1964):

Select "Momentum" above. Along the valley, successive gradients point the same way and add up in v; across it, they flip sign from step to step and partly cancel. At the default η the path still crosses the valley floor, but less often, and it gets much closer to (1, 1) in the same 200 updates. At β = 0 no velocity carries over, and the update is exactly SGD.

Adam: Adaptive Moments

Adam combines momentum with a separate step size for each parameter. It keeps running averages of the gradient (the first moment) and of its square (the second moment):

Here t = 1, 2, … counts updates, so θ_t is the point after t updates and m₀ = v₀ = 0. Dividing by 1 − β^t removes the pull toward that zero start: on the first update m̂₁ = g₁ and v̂₁ = g₁², so each coordinate moves by η·g₁/(|g₁| + ε), almost exactly η. The small constant ε = 10⁻⁸ is added after the square root, so it only matters when √v̂_t is itself near 10⁻⁸. Squares, roots and division act per coordinate. This run uses β₁ = 0.9, β₂ = 0.999 and ε = 10⁻⁸, the default settings of Kingma & Ba, Algorithm 1, where η is written α.

Dividing by √v̂_t puts every coordinate on the same scale: usually |m̂_t/√v̂_t| is at most about 1, so no coordinate moves much more than η per update, however large its gradient is (Kingma & Ba, 2015, section 2.1). That is why Adam does not bounce off the steep walls here. It is also why, at η = 0.001, it covers little ground in 200 updates. As the run nears the valley floor its gradients fall far below their starting size, but v̂_t still remembers the large early ones, so the steps shrink well below η.

Muon: Orthogonalized Updates

Muon is an optimizer for the two-dimensional weight matrices of hidden layers (Jordan et al., 2024). It first forms an SGD-momentum update G for a matrix, then replaces G by an approximation of UVᵀ, where G = UΣVᵀ is its singular value decomposition. This keeps the directions of the update and sets its singular values to about 1. Only the update is orthogonalized; the weights themselves are not constrained.

UVᵀ is the semi-orthogonal matrix nearest to G. An exact SVD at every step would be slow, so Muon runs five steps of a Newton–Schulz iteration, which uses only matrix products; the figure steps through the classic version of this iteration. Muon's polynomial is tuned to be fast rather than exact, targeting singular values in about [0.7, 1.3]. Five steps can leave very small singular values below that band, and zero stays zero; Jordan et al. report that five steps suffice in their experiments. Their motivation is that momentum updates for transformer weight matrices are nearly low-rank: a few directions dominate, and orthogonalizing gives the rarer directions a comparable step. In a second view, UVᵀ is the direction of steepest descent when the size of a step is measured by the spectral norm (Bernstein & Newhouse, 2024).

Jordan et al. report faster training on the NanoGPT and CIFAR-10 speed benchmarks, and Liu et al. (2025) report about twice the compute efficiency of AdamW for language models once weight decay is added and update sizes are matched. In an independent comparison against carefully tuned AdamW, Wen et al. (2025) found the speedup of matrix-based optimizers such as Muon shrinking with model size, from about 1.4× at 0.1B parameters to 1.1× at 1.2B.

Edge of Stability

On a quadratic with curvature λ, gradient descent with learning rate η diverges when λ > 2/η. Cohen et al. (2021) found that full-batch gradient descent on neural networks drives the sharpness, the largest eigenvalue of the loss Hessian, up until it reaches about 2/η, and then holds it there:

The quadratic picture predicts divergence in this regime, yet training continues to make progress. Damian et al. (2023) explain part of this as self-stabilization: the growing oscillation along the sharpest direction itself pushes the sharpness back down.

Grokking

Power et al. (2022) trained small transformers on tables of arithmetic such as addition modulo a prime. Long after the training set was fitted perfectly, test accuracy jumped from chance to near perfect. They called this delayed generalization grokking:

Smaller training sets needed more optimization before the jump, and of the training choices Power et al. compared, weight decay helped generalization most. For modular addition, Nanda et al. (2023) found that the generalizing circuit forms gradually during the apparent plateau; the sudden drop in test loss comes when weight decay removes the memorizing components.

DPO vs RLHF

RLHF, as used for InstructGPT (Ouyang et al., 2022), fits a reward model to human preference pairs and then optimizes the language model against it with reinforcement learning, with a KL penalty that keeps it close to a reference model. Direct Preference Optimization (Rafailov et al., 2023) shows that for this KL-regularized objective the policy itself determines the reward, up to a term that cancels when two answers are compared. The policy can then be trained on preference pairs directly, with a classification loss:

Here y_w is the preferred answer to prompt x and y_l the rejected one. DPO needs no separate reward model and no sampling during training, which makes it simpler to run. Whether it gives better-aligned models depends on the setting: Xu et al. (2024) found carefully tuned PPO-based RLHF ahead of DPO on several benchmarks, including code generation.

Sharp and Flat Minima

A loss over millions of parameters can be drawn only as a slice: the loss along two directions through the trained weights. Li et al. (2018) choose random directions and rescale them filter by filter, so that slices of different networks can be compared. Near a minimum, where the gradient is zero, the shape of the slice is set by the Hessian H:

The figure takes one gradient-descent step on L(x, y) = x² + 2y², whose Hessian diag(2, 4) is the same at every point. The step multiplies y by 1 − 4η, so along y, the sharper direction, any η above 2/4 = 0.5 leaves y farther from 0 than before: the 2/λ limit from the edge-of-stability section.

Whether flat minima generalize better is debated. Keskar et al. (2017) linked the sharper minima found with large batches to worse test accuracy; Dinh et al. (2017) showed that reparameterizing a network can make a minimum arbitrarily sharp without changing the function it computes.

Backprop Through Attention

In attention, the gradient reaching the queries and keys passes through a softmax. With scores S = QK^T/√d and weights A = softmax(S) taken row by row:

When one weight in a row approaches 1, the softmax Jacobian approaches zero and little gradient gets through. Vaswani et al. (2017) divide the scores by √d to keep them from growing with dimension and saturating the softmax. Residual connections add a second path: with x_l+1 = x_l + F(x_l), the Jacobian is I + ∂F/∂x_l, so gradient can reach early layers without passing through every softmax (He et al., 2016).

Task Vectors

A model fine-tuned from pretrained weights differs from them by a "task vector" in weight space (Ilharco et al., 2023). Task vectors from the same pretrained model can be added and scaled:

Ilharco et al. found that adding task vectors can give one model several of the fine-tuned abilities, and that subtracting one reduces the corresponding behaviour while leaving others largely intact. The scaling coefficients λ are chosen on held-out data, and no training is needed beyond the original fine-tuning runs.

Neural Scaling Laws

For language models, the loss falls as a power law in the number of parameters N and training tokens D over several orders of magnitude (Kaplan et al., 2020). Hoffmann et al. (2022) fit:

Here E is the irreducible loss and C ≈ 6ND is the training compute in FLOPs. With a ≈ b ≈ 0.5, a tenfold larger compute budget should go to a model and a dataset that are each about √10 ≈ 3.2 times larger. Fits like this are used to predict the loss of larger runs before training them.