Basics of Optimizers

Training a model means choosing its parameters so that its predictions incur a small loss. If all trainable parameters are collected in a vector \(\theta\in\mathbb R^d\), the problem is

For a neural network, \(d\) may be billions and \(L(\theta)\) is a deeply nested nonlinear function. We cannot list every possible \(\theta\), evaluate the loss at each one, and select the best. Nor can we usually rearrange the equation \(\nabla L(\theta)=0\) and solve directly for \(\theta\). Training must therefore be iterative: start from some \(\theta_0\), make a small informed change, inspect the new local information, and repeat.

An optimizer answers the decision made at every such iteration:

Given what we know at the current parameters, which update \(\Delta\theta_t\) should we apply next?

This wording also separates two algorithms that are often conflated. Backpropagation computes how sensitive the current loss is to each parameter—the gradient. The optimizer consumes that gradient, combines it with its state and hyperparameters, and produces the actual parameter update. Backprop provides local evidence; the optimizer decides how to act on it.

Throughout, \(L(\theta)\) is the objective, \(g_t\) is the gradient available at step \(t\), \(\Delta\theta_t=\theta_{t+1}-\theta_t\) is the chosen update, and \(\eta_t>0\) is the learning rate.

1. Gradient Descent Fundamentals

For supervised learning, the objective is usually an average of per-example losses:

Here \(\xi\) is one example drawn from a data distribution \(\mathcal D\), \(\ell(\theta;\xi)\) measures how badly the current model handles that example, and \(L\) is the average behavior we actually care about. In a finite dataset, the expectation can simply be read as an average over its examples.

Notice the mismatch between the question and the information available. The question is global—which parameters have the smallest loss?—but at step \(t\) we stand at only one point \(\theta_t\). A forward pass tells us \(L(\theta_t)\). Backpropagation tells us the local slope

for a full batch, or an estimate of it for a mini-batch. The symbol \(\nabla_\theta\) collects one partial derivative per parameter,

Thus \(g_{t,i}\) answers a deliberately local question: if only \(\theta_i\) increased by a tiny amount, at what rate would the loss change? Neither the loss value nor this vector tells us the loss everywhere else. The optimizer must use this limited local information to choose \(\Delta\theta_t\).

1.1 Why Does Taylor Expansion Appear?

Suppose we are considering many possible updates. Evaluating the full network and loss after every candidate would be prohibitively expensive. We instead want a cheap local model that answers:

If I move by a small \(\Delta\theta\), approximately how will the loss change?

The derivative is defined precisely to answer how a function changes under a small displacement. Collecting the derivatives for all parameter coordinates gives the first-order Taylor expansion

The first term is the loss we already know. The dot product in the second term is the predicted change:

Every parameter’s slope is multiplied by how far we propose to move that parameter, and the contributions are added. For example, a positive component of \(g_t\) says that increasing that parameter alone would locally increase the loss; a negative component says the reverse.

Why keep only the first-order term? Higher-order terms contain curvature and can make the prediction more accurate, but constructing and solving with them is far more expensive at neural-network scale. When \(\Delta\theta\) is small, their contribution also shrinks faster than the linear term. Taylor expansion is therefore not an arbitrary proof trick: it is the simplest affordable model of the nearby loss surface.

Before optimizing the expression, separate what is known from what is being chosen:

Thus \(\Delta\theta\) is the update direction being chosen. The known vector \(g_t\) tells us how each possible direction would score. To compare directions fairly, first give every candidate the same length \(\lVert\Delta\theta\rVert_2=r\). We then ask which direction makes \(g_t^\top\Delta\theta\) as negative as possible. Cauchy–Schwarz gives

and the lower bound is attained only when \(\Delta\theta\) points exactly opposite \(g_t\). Therefore the best fixed-length candidate is

This answers the direction question: \(g_t\) points toward steepest local increase, while \(-g_t\) points toward steepest local decrease. Plain gradient descent does not usually choose a fixed radius \(r\) explicitly. It introduces a positive learning rate \(\eta_t\) and sets

The negative sign chooses the direction; \(\eta_t\) scales the distance. In fact,

so a fixed radius \(r\) would correspond to \(\eta_t=r/\lVert g_t\rVert_2\). Substitution into the Taylor prediction gives

For example, if \(g_t=(3,4)\) and we allow radius \(r=1\), then \(\lVert g_t\rVert_2=5\) and

Its predicted change is \((3,4)^\top(-3/5,-4/5)=-5\), the smallest possible among all unit-length updates. If instead \(g_t=0\), the first-order model predicts no change in any direction, so it cannot choose a descent direction from gradient information alone.

Thus the negative gradient is not merely a downhill direction: under the ordinary Euclidean notion of a fixed-size move, it gives the largest decrease predicted by the local linear model. Gradient descent applies exactly this update:

This conclusion has two qualifications. It is about the local approximation, not the entire landscape, so it does not promise the global minimum. And it determines a direction but not a safe travel distance. An arbitrarily large step can leave the neighborhood where the Taylor model is accurate and increase the real loss. That remaining decision is the learning rate.

1.2 From a Direction to a Distance: the Learning Rate

The learning rate converts the gradient direction into a displacement. Too small wastes iterations; too large makes the local prediction unreliable and can overshoot. A quadratic is a degree-two function; it is the simplest function with nonzero curvature. A one-dimensional quadratic lets us see this tradeoff without any neural-network detail:

Gradient descent becomes

This one equation shows three regimes:

• if \(0<\eta a<1\), the parameter approaches zero without changing sign;
• if \(1<\eta a<2\), it crosses zero on every step but the oscillation shrinks;
• if \(\eta a>2\), the magnitude grows and training diverges.

Here \(a\) measures curvature: large \(a\) means the loss bends upward sharply. The three regimes show that “the gradient points downhill” is not enough. The step must be small relative to the distance over which the surface bends. In many dimensions, different directions have different curvatures, so the next difficulty is that one global learning rate must somehow serve all of them.

2. Stochastic Gradient Descent

The update derived above is the cleanest possible baseline, but it quietly assumes two conveniences. First, \(g_t\) faithfully describes the objective we care about. Second, measuring every parameter coordinate with the same Euclidean ruler is sensible. Neural-network training violates both assumptions: the gradient is estimated from sampled data, and the loss bends at very different rates in different directions. We will introduce these difficulties separately so that each later modification has a clear job.

2.1 Mini-Batches Turn the Gradient into an Estimate

The population objective averages over every possible data example, so its exact gradient is itself an expectation. The true data distribution is unknown, and even averaging over every example in a large stored dataset before every update would be expensive. A mini-batch makes one update affordable: for a sampled set \(B_t\) of size \(b\), we use

Why is this a reasonable substitute? Under ordinary independent uniform sampling, each example has the correct chance to appear. Averaging over all possible choices of the mini-batch therefore recovers the full dataset gradient:

Here \(\mid\theta_t\) means that \(\theta_t\) is held fixed while the expectation averages only over the random choice of mini-batch.

Here \(L\) denotes the finite-dataset average. This equality does not say that the dataset gradient equals the unknown population gradient; the dataset itself is only a sample from the world. There are therefore two sampling gaps: a mini-batch approximates the dataset, and the dataset approximates the population.

The word unbiased describes the first gap—an average over hypothetical mini-batches. It does not say that the particular batch in front of us equals the full gradient. One batch may contain unusually easy examples, another unusually hard ones, so their gradients can differ even at the same \(\theta_t\). Increasing batch size usually lowers this sampling variation, but costs more computation per update.

The tradeoff is now visible. A full-batch gradient obtains a more stable direction but waits longer before moving. A mini-batch gradient is cheap and frequent, but every estimate is uncertain. Combining this estimate with the gradient-descent update produces SGD, whose complete update appears below in this section. Momentum will then use history to distinguish persistent signal from batch-to-batch fluctuation.

2.2 Why Plain SGD Needs More Context

Even with an exact full-batch gradient, one global learning rate can be inefficient. Section 1 used only the first-order Taylor term to choose a direction. To understand how long that direction remains trustworthy, we must look at the next term:

The Hessian \(H_t\) is the matrix of second derivatives. It records curvature: how the gradient itself changes as we move. Near a smooth local minimum \(\theta^*\), the gradient is approximately zero, leaving the familiar quadratic model

An eigenvector of \(H\) is a direction in which this quadratic does not mix with the other directions; its eigenvalue \(\lambda_i\) is the curvature along that direction. Gradient descent there behaves exactly like the one-dimensional example with \(a=\lambda_i\). Stability is limited by the largest curvature, while progress in a flat direction is controlled by the smallest.

The objective in Figure 1 makes this concrete:

The same displacement in \(y\) changes the loss twelve times as sharply as in \(x\). Stability in the \(y\) direction requires \(\eta<2/12=1/6\). But after respecting that limit, the update along the flatter \(x\) direction is necessarily modest. The path therefore tends to bounce across the narrow valley while advancing slowly along it. More generally, a large condition number

creates exactly this separation of scales.

We have now exposed two different losses of information. A current mini-batch gradient does not tell us which components will persist on the next batch, and its raw values do not tell us whether a large component comes from a useful direction or merely from a differently scaled, highly curved direction.

Here a coordinate means one scalar position in the parameter vector. If

then \(g_{t,i}=\partial L/\partial\theta_i\) is the gradient component associated with parameter \(\theta_i\). “Across coordinates” means comparing these components and their histories. For example, suppose several successive gradients look like

The first coordinate repeatedly has magnitude near \(100\), whereas the second is near \(0.01\). Plain SGD multiplies both by the same learning rate, so its first-coordinate update will be thousands of times larger. But a persistently large gradient does not by itself prove that this parameter should move much farther: the difference can also come from parameter units, network parameterization, or curvature. An optimizer may therefore accumulate two kinds of context:

• across steps: has a gradient component kept the same direction, or has its sign fluctuated over time?
• across coordinates: do different parameters’ gradients persistently live on different magnitude scales, so their update scales should be adjusted separately?

Momentum uses the first kind of history. RMSProp records a typical squared-gradient magnitude separately for every coordinate and uses the second. Adam uses both.

2.3 An Optimizer Is a State-Update Rule

A stateless rule sees only the current \(g_t\). It cannot distinguish a component that has pointed the same way for twenty steps from one that alternates sign on every step; nor can it know whether a large coordinate is consistently large. To make those distinctions, the optimizer must carry information from earlier steps. We can write such an optimizer abstractly as

where \(s_t\) is optimizer state: a compact summary of the gradient history. The function \(F\) updates that memory; \(U\) turns the memory and current gradient into a parameter change. SGD without momentum has no persistent per-parameter state. Momentum stores one moving average. Adam stores two. These buffers survive across training steps and must be checkpointed if training is to resume with identical dynamics. Their memory cost is discussed separately in LLM Optimization Basics: Memory.

2.4 The Plain SGD Algorithm

We now have the two ingredients needed to define the standard baseline. Gradient descent says how to update from a gradient; mini-batching says how to obtain an affordable estimate of that gradient. Stochastic gradient descent (SGD) combines them:

One SGD step has a concrete order:

1. sample a mini-batch \(B_t\);
2. run the model and average its per-example losses on that batch;
3. use backpropagation to compute the mini-batch gradient \(g_t\);
4. multiply it by \(-\eta_t\) and update the parameters.

The word stochastic refers to step 1: because the batch is random, \(g_t\) and the resulting update are random. Once \(B_t\), \(\theta_t\), and \(\eta_t\) are fixed, the plain SGD rule itself is deterministic. When the mini-batch estimator is unbiased and the learning rate is fixed during the conditional expectation,

Thus one update need not lower the full objective, but its expected local direction matches full-batch gradient descent. The computational advantage is that one step processes \(b\) examples rather than the entire dataset. The price is sampling noise.

Plain SGD stores no moving average or per-coordinate scale: apart from the parameters themselves and an external step counter or schedule, the next update uses only the current batch gradient. This makes it a low-memory, transparent baseline. It also leaves both difficulties identified above unresolved: batch fluctuations pass directly into the update, and one global learning rate still serves every coordinate.

An optimizer step and an epoch are different units. One step consumes one effective batch and updates parameters once; one epoch processes roughly one dataset’s worth of examples. With gradient accumulation, parameters stay fixed while \(K\) microbatch gradients are combined, typically as \(g_t=K^{-1}\sum_{k=1}^K g_t^{(k)}\), before one optimizer update. It enlarges the effective batch without storing all examples’ activations at once; it does not create \(K\) optimizer steps.

2.5 SGD with Momentum: Consistency across Steps

Write a mini-batch gradient informally as “persistent direction + sampling fluctuation.” We do not know either term separately, but repeated observations help: a component that keeps the same sign is more credible evidence of a persistent direction, while rapidly alternating components tend to cancel when averaged.

A simple average of every past gradient would eventually react very slowly because ancient gradients—computed at quite different parameters—would retain the same weight as recent ones. Momentum instead uses an exponential moving average (EMA), which gives the newest gradient weight \(1-\beta\) and discounts the existing history by \(\beta\):

The first equation updates the memory; the second uses that smoothed direction in place of the raw current gradient. Expanding the recurrence shows the exact weights:

when \(m_0=0\). Within the current buffer \(m_t\), a more recent gradient has a larger direct coefficient: the coefficient on a gradient \(j\) steps old is \((1-\beta)\beta^j\). The characteristic averaging window is on the order of \(1/(1-\beta)\) steps: about \(10\) for \(\beta=0.9\) and \(100\) for \(\beta=0.99\). This is a scale, not a hard cutoff; older gradients never disappear abruptly.

This statement does not say that the newest data always has the greatest influence on the entire training trajectory. At the first step,

so \(g_1\) is the only observed direction and therefore determines the direction completely; the factor \(1-\beta\) only shrinks the buffer magnitude under this convention. That first update then changes \(\theta_1\), which changes every gradient computed afterward. Early gradients can therefore have a lasting indirect effect through the parameters, even after their direct coefficients inside a later buffer \(m_t\) have decayed. The expansion above describes the latter only.

The name “momentum” comes from the qualitative analogy with velocity: the update carries some previous direction forward instead of responding only to the force-like current gradient. The formula is an optimization rule, not a literal simulation of physical mechanics.

Why does this help in a narrow valley? Across the steep direction, successive gradients often alternate signs, so they cancel in the moving average. Along the shallow direction, their signs remain consistent, so they survive the averaging. Momentum therefore damps side-to-side oscillation while preserving motion along the valley. With the normalized EMA convention used here, any additional acceleration depends on the learning rate and the dynamics—not on the buffer’s raw magnitude growing without bound.

3. Adam

3.1 Coordinate-wise Scaling

Momentum addresses when gradient evidence is consistent, but it still measures every coordinate in raw gradient units. Return to the narrow valley: the steep coordinate may be numerically large even when we are already close to the valley floor, while a smaller gradient along the flat coordinate may be the direction in which substantial progress remains.

Does this contradict the claim that \(-g_t\) is the best direction? No. Section 1 proved a narrower statement: using only the linear Taylor model, among updates with the same Euclidean length, \(-g_t\) gives the largest immediate predicted decrease. The gradient’s magnitude is a slope—loss change per unit parameter change—not a measurement of the distance to the minimum or of how long that slope will remain valid.

A one-dimensional quadratic makes the missing information explicit:

The same observed gradient \(g=1\) could mean \(\theta=1\) when \(a=1\), but only \(\theta=0.01\) when \(a=100\). The gradient has the same sign and magnitude in both cases, yet the distances to the minimum differ by a factor of one hundred. We need the curvature \(a\) to recover the appropriate quadratic step:

In several dimensions, the analogous quadratic step is \(-H^{-1}g\) when the Hessian \(H\) is positive definite; it is generally not parallel to \(-g\). Thus coordinate-wise scaling does not repair a gradient that “forgot” its magnitude. It uses additional scale information to reinterpret how far each gradient component should move. Adaptive optimizers do not know the exact Hessian: their historical squared-gradient statistics provide a practical diagonal scaling heuristic.

Here global multiplier means the one scalar learning rate shared by all coordinates, not an optimizer that can find a global optimum. In plain gradient descent,

so \(\Delta\theta_i/g_i=-\eta\) for every \(i\). Changing \(\eta\) expands or contracts the whole update, but it cannot make the steep coordinate cautious relative to its gradient while making the flat coordinate aggressive. For example, consider

At \((x,y)=(0.1,1)\), the gradient is \((10,1)\) even though the remaining distances to the minimum \((0,0)\) are only \(0.1\) in \(x\) and \(1\) in \(y\). The shared multiplier produces

Every choice of \(\eta\) preserves this \(10:1\) update ratio. The steep \(x\) direction forces \(\eta\) to remain small for stability; that same small number then makes progress along the flat \(y\) direction slow. Momentum may smooth sign changes over time, but its update \(-\eta m_t\) still applies one scalar to all components of the raw-unit vector \(m_t\). Distinguishing coordinate scales requires a vector of effective learning rates—or, equivalently, a separate scale estimate for every coordinate.

Adaptive methods therefore build a separate recent scale for every coordinate and divide by it. To estimate magnitude without positive and negative gradients cancelling, they accumulate \(g_t^2\) rather than \(g_t\). Taking a square root converts the squared quantity back to gradient units, so those units cancel in “gradient divided by typical gradient magnitude.” The ratio is dimensionless—it has no remaining physical unit—before the global learning rate is applied.

AdaGrad: Accumulate Squared Gradients

The most direct version is AdaGrad. Beginning from \(v_0=0\), it adds every squared gradient observed so far:

where \(\odot\) denotes element-wise multiplication, not a dot product or matrix multiplication. If

then

Thus \(v_t\) has the same shape as \(g_t\), and coordinate \(v_{t,i}\) accumulates only the squared gradients previously observed for parameter \(\theta_i\). AdaGrad then updates

The square root, addition, and division in this update are also element-wise. Written for one coordinate,

Thus coordinate \(i\) has effective learning rate \(\eta/(\sqrt{v_{t,i}}+\epsilon)\). A coordinate that repeatedly receives large gradients builds a large denominator and takes smaller future steps. A rare feature—one active in only a small fraction of examples—produces zero gradient on most steps and retains a relatively large effective step, which makes AdaGrad attractive for sparse data. The small positive \(\epsilon\) prevents division by zero before a coordinate has accumulated any scale.

Its weakness is visible in the recurrence: \(v_t\) only increases. Effective learning rates continually shrink and can become too small during long nonstationary neural-network training.

RMSProp: Forget Stale Squared Gradients

AdaGrad treats a gradient from the first training step as permanently relevant. But the model and its loss geometry change throughout training, so an old scale estimate can become stale. RMSProp makes one surgical change: replace the never-forgetting sum with an exponential moving average,

The denominator estimates each coordinate’s recent root-mean-square gradient. Dividing by it reduces persistent scale differences: after enough steps, multiplying one coordinate’s gradients by a positive constant approximately multiplies both numerator and denominator by that constant, leaving the normalized update similar. The equivalence is only approximate because of initialization, \(\epsilon\), momentum, clipping, and nonstationarity.

The quantity \(v_t\) is called an uncentered second moment because a moment is an average of a power: the first raw moment averages \(g\), and the second averages \(g^2\). It is not the statistical variance; variance would be \(\mathbb E[g^2]-\mathbb E[g]^2\).

3.2 Combining Momentum and Coordinate-wise Scaling

Momentum and RMSProp changed two independent parts of plain SGD. Momentum replaced the noisy numerator by a smoothed direction. RMSProp kept the current numerator but divided it by a recent coordinate-wise scale. Adam applies both ideas at once: smooth the direction in the numerator and normalize it by the smoothed magnitude in the denominator.

It therefore maintains two buffers:

The names “first moment” and “second moment” refer here to averages of the observed gradient stream: \(m_t\) averages \(g_t\), whereas \(v_t\) averages \(g_t^2\). They are not moments of the training-data distribution or of the model’s predictions.

Both buffers start at zero. During the first few steps, part of their notional averaging window is therefore filled with zeros rather than observations, which pulls them toward zero. Adam corrects this initialization bias:

The update is

Every part now has a distinct role:

Does \(v_t\) affect the update direction? For one scalar coordinate, its nonnegative denominator cannot change the sign set by \(m_t\); it changes only that coordinate’s magnitude. For the full parameter vector, however, different entries of \(v_t\) rescale coordinates by different amounts and therefore generally rotate the update. Ignoring \(\epsilon\), for example,

which is not parallel to \(-\hat m=(-1,-1)\). The numerator supplies each coordinate’s sign; the denominator changes the coordinates’ relative magnitudes.

The following pseudocode puts one complete Adam step in execution order. Hover a highlighted symbol—or focus it with the keyboard—to inspect what it stores and why it is needed.

θ ← initial_parameters()
m ← zeros_like(θ)
v ← zeros_like(θ)
t ← 0

repeat until the stopping condition:
    t ← t + 1
    B ← sample_minibatch()
    g ← minibatch_gradient(θ, B)
    m ← β₁·m + (1−β₁)·g
    v ← β₂·v + (1−β₂)·(g ⊙ g)
    m̂ ← m / (1−β₁^t)
    v̂ ← v / (1−β₂^t)
    θ ← θ − ηₜ·m̂ / (sqrt(v̂) + ε)

return θ

3.3 Why Is Bias Correction Necessary?

Suppose the gradient is the same vector \(g\) for the first several steps. Starting from zero,

Why do those denominators have exactly this form? The EMA weights accumulated through step \(t\) sum to \(1-\beta^t\), not one. Dividing by that sum renormalizes the observed weights. Equivalently, in the constant-gradient example, division by \(1-\beta_1^t\) and \(1-\beta_2^t\) recovers \(g\) and \(g^2\). The correction approaches one as the history fills, so it matters most near the start of training.

An instructive consequence appears on the first step. Ignoring \(\epsilon\),

coordinate-wise. Adam’s first update depends mainly on the signs of nonzero gradient coordinates, not their raw magnitudes. This is one reason an Adam learning rate is not numerically interchangeable with an SGD learning rate.

Here \(\operatorname{sign}(a)\) is \(+1\) when \(a>0\), \(-1\) when \(a<0\), and \(0\) when \(a=0\). The statement is an approximation for nonzero coordinates because the actual \(\epsilon\) prevents exact cancellation.

3.4 What Problem Does Adam Solve—and Not Solve?

Adam is helpful when gradients are noisy, sparse, or very differently scaled across coordinates. It often reaches a useful training regime with less tuning than plain SGD. But its denominator rescales coordinates independently; it cannot rotate the update to undo off-diagonal interactions between parameters, infer how far the local approximation remains valid, or guarantee a better minimum. A small recent \(v_{t,i}\) can also amplify coordinate \(i\). The floor \(\epsilon\) limits that amplification; learning-rate warmup begins with deliberately small steps; gradient clipping caps unusually large gradient norms; and sufficient numerical precision keeps these small statistics representable.

The usual defaults \(\beta_1=0.9\), \(\beta_2=0.999\), and a small \(\epsilon\) are starting points, not laws. The much larger \(\beta_2\) gives the scale estimate a longer, steadier window than the direction estimate. Changing batch size, loss normalization, model scale, or precision can change the appropriate learning rate and sometimes the moment constants.

4. AdamW

So far every modification has tried to use the data gradient more effectively. Weight decay asks a separate question: besides fitting the training objective, do we want to prefer parameter vectors with smaller norms? Such a preference can act as regularization or norm control, though its effect depends on the architecture. It is easy to conflate with the optimizer because it is applied during the same update.

The traditional way to express this preference is to add an \(L_2\) penalty to the loss. This gives

and gradient

For plain SGD, substituting this gradient into the update and collecting the two terms gives

The factor \(1-\eta\lambda\) literally shrinks every selected parameter by the same proportion on that step. This is why \(L_2\) regularization and weight decay are equivalent for plain SGD.

For Adam, however, inserting \(\lambda\theta\) into the gradient sends the penalty through the first- and second-moment normalization. A coordinate’s penalty is then divided by a history-dependent scale, so different coordinates experience different effective shrinkage. Coupled \(L_2\) regularization is no longer the same operation as uniform weight decay.

AdamW keeps the loss gradient and decay separate:

The moments are computed from the data gradient \(g_t\), not from \(g_t+\lambda\theta_t\). This is the meaning of decoupled weight decay.

Here is a minimal array-level implementation. A production model stores parameters in a parameter tree, a nested collection of named arrays. Distributed training also reduces gradients by combining contributions from multiple devices. The example omits those structures, mixed precision, and fused kernels so the update equation remains visible.

import jax.numpy as jnp

def adamw_step(params, grads, state, lr,
               beta1=0.9, beta2=0.999,
               eps=1e-8, weight_decay=0.01):
    t = state["t"] + 1
    m = beta1 * state["m"] + (1.0 - beta1) * grads
    v = beta2 * state["v"] + (1.0 - beta2) * jnp.square(grads)

    m_hat = m / (1.0 - beta1 ** t)
    v_hat = v / (1.0 - beta2 ** t)

    params = ((1.0 - lr * weight_decay) * params
              - lr * m_hat / (jnp.sqrt(v_hat) + eps))
    return params, {"t": t, "m": m, "v": v}

The code assumes \(t\) counts optimizer updates beginning at zero and that \(m,v\) have the same shape as the parameter array. Production libraries may place \(\epsilon\) differently, use AMSGrad, apply capturable step counters, or fuse decay and updates. Here, placing \(\epsilon\) differently means choosing \(\sqrt{\hat v}+\epsilon\) versus \(\sqrt{\hat v+\epsilon}\); AMSGrad replaces the current second-moment denominator by a running maximum; a capturable counter stores \(t\) where a compiled accelerator graph can access it; and a fused kernel performs several array operations together to reduce memory traffic. These choices preserve the central idea but can matter for exact numerical reproduction.

5. Muon

Adam’s scale estimate is coordinate-wise: it would work the same way if a weight matrix were flattened into a long vector. But a matrix in a linear layer is not merely stored in two dimensions by accident. If \(W\) maps an activation \(h\) to \(Wh\), its rows and columns jointly describe input and output directions. A matrix update can be strong along one pair of directions and weak along another. Treating each entry independently discards this structure.

Muon starts from the same temporal question as momentum and then changes the geometry of the resulting matrix update. For a weight \(W\in\mathbb R^{m\times n}\), its gradient \(G_t\) and momentum \(M_t\) have the same shape. First form the momentum matrix

To see the next operation, recall what a singular value decomposition says. Any matrix can be written

The columns of \(V\) identify orthogonal input directions, the columns of \(U\) identify the corresponding output directions, and the nonnegative diagonal entries of \(\Sigma\)—the singular values—say how strongly the update acts along each paired direction. If one singular value is much larger than the others, that mode dominates the matrix update.

Ideal orthogonalization keeps the two sets of directions but replaces every nonzero singular value by one:

For a rectangular matrix the result is more precisely semi-orthogonal, but the important statement is the same: its nonzero singular modes now have equal strength. This is a matrix-level analogue of normalization. It does not assert that the true loss has identical curvature in every singular direction; it chooses a better-balanced candidate update from the information in \(M_t\).

Computing a full SVD for every eligible weight on every step would be expensive. Muon instead normalizes the starting matrix and applies a short Newton–Schulz-style polynomial iteration. Each iteration uses matrix multiplications and moves the singular values toward one, thereby approximating \(UV^\top\) without explicitly constructing \(U\), \(\Sigma\), and \(V\). The implementation then applies its prescribed dimension-dependent scale and learning rate. Thus a practical Muon configuration includes more than the symbol \(UV^\top\): one must specify the momentum convention, orthogonalization approximation, scaling, decay, and parameter grouping.

This is a different kind of adaptivity from Adam:

• Adam rescales individual scalar coordinates using their historical squared gradients;
• Muon reshapes an entire matrix update using its singular directions.

Why not apply this to every parameter? A bias or normalization gain is a scalar or vector and has no two-dimensional input–output structure to orthogonalize. An embedding table maps discrete token identities to vectors, while an output head maps hidden states to output scores; these matrices play special boundary roles for which the same scaling rule need not be appropriate. Muon is therefore normally used for eligible two-dimensional hidden-layer weights, while embeddings, output heads, normalization gains, biases, and other vector or scalar parameters are handled by AdamW.

Muon is included here because it makes the design space clearer, not because it universally dominates AdamW. Orthogonalization adds matrix multiplications. If a matrix is split across devices, those multiplications may also require communication, so operation count alone does not determine wall-clock cost. A fair comparison should match compute or training time, number of processed tokens, learning-rate schedules, and the amount of hyperparameter tuning given to each method.

训练模型，就是选择一组参数，使模型预测产生的 loss 尽可能小。把全部可训练参数收集成 vector \(\theta\in\mathbb R^d\)，我们总体上要解决的是

对 neural network 而言，\(d\) 可能达到数十亿，\(L(\theta)\) 又是一个层层嵌套的 nonlinear function。我们不可能列出每一种 \(\theta\)，逐个计算 loss，再挑出最好的一个；通常也不能把方程 \(\nabla L(\theta)=0\) 重新排列，直接解出 \(\theta\)。因此，训练只能采用 iterative 的方式：从 \(\theta_0\) 出发，先作一次有依据的小改动，到达新位置后再取得新的局部信息，如此重复。

Optimizer 回答的正是每次 iteration 中都要作出的决定：

根据当前位置掌握的信息，下一次应该施加怎样的更新 \(\Delta\theta_t\)？

这个问法也把两个经常混在一起的算法分开。Backpropagation 计算当前 loss 对各参数有多敏感，也就是 gradient；optimizer 接收 gradient，再结合自身状态与超参数，产生真正的参数更新。Backprop 提供局部证据，optimizer 决定怎样依据这些证据行动。

全文用 \(L(\theta)\) 表示 objective，\(g_t\) 表示第 \(t\) 步得到的 gradient，\(\Delta\theta_t=\theta_{t+1}-\theta_t\) 表示选定的 update，\(\eta_t>0\) 表示 learning rate。

对 supervised learning，objective 通常是 per-example losses 的平均：

其中 \(\xi\) 是从 data distribution \(\mathcal D\) 中取得的一个样本，\(\ell(\theta;\xi)\) 衡量当前模型在这个样本上错得多严重，而 \(L\) 是我们真正关心的平均表现。若数据集有限，也可以直接把 expectation 理解成对全部样本求平均。

现在注意“问题”和“已知信息”之间的差距。问题是全局的——哪组参数的 loss 最小？——但第 \(t\) 步时，我们只站在一个位置 \(\theta_t\)。Forward pass 告诉我们 \(L(\theta_t)\)；backpropagation 告诉我们当前位置的 slope：

对 mini-batch 则返回它的一个估计。符号 \(\nabla_\theta\) 把每个参数的一条 partial derivative 收集起来：

因此，\(g_{t,i}\) 回答的是一个刻意局部的问题：若只把 \(\theta_i\) 增加一点点，loss 会以多快的速度改变？Loss value 与这个 vector 都没有告诉我们其他所有位置的 loss。Optimizer 必须只用这些有限的局部信息选择 \(\Delta\theta_t\)。

1.1 为什么这里会出现 Taylor Expansion？

假设我们正在考虑许多种候选更新。若对每个候选 \(\Delta\theta\) 都真正运行一次完整 network 并计算更新后的 loss，代价会高得无法接受。因此，我们需要一个便宜的局部模型来回答：

如果移动一个很小的 \(\Delta\theta\)，loss 大约会改变多少？

Derivative 的定义，正是描述函数在微小位移下怎样变化。把所有参数坐标的 derivatives 收集起来，就得到 first-order Taylor expansion：

第一项是已经知道的当前 loss；第二项的 dot product 是对 loss 变化量的预测：

每个参数的 slope 乘以该参数拟移动的距离，再把所有贡献相加。例如，\(g_t\) 的某个 component 为正，意味着只增加那个参数会在局部提高 loss；component 为负时则相反。

为什么只保留 first-order term？Higher-order terms 包含 curvature，确实能让预测更准确，但在 neural-network scale 上构造并求解这些项昂贵得多。而当 \(\Delta\theta\) 足够小时，higher-order terms 也会比 linear term 更快缩小。因此 Taylor expansion 不是为了证明而突然搬来的技巧；它是我们付得起的、最简单的附近 loss-surface 模型。

在优化这个 expression 之前，先把已知量与待选择量分开：

因此，真正等待选择 update direction 的是 \(\Delta\theta\)；已知 vector \(g_t\) 用来给每个候选方向打分。为了公平比较方向，先让每个 candidate 具有相同长度 \(\lVert\Delta\theta\rVert_2=r\)，再问哪个方向能使 \(g_t^\top\Delta\theta\) 尽可能负。Cauchy–Schwarz inequality 给出

只有当 \(\Delta\theta\) 与 \(g_t\) 方向恰好相反时才能达到 lower bound。因此，最佳 fixed-length candidate 是

这就回答了方向问题：\(g_t\) 指向局部上升最快的方向，\(-g_t\) 指向局部下降最快的方向。Plain gradient descent 通常不会显式选择固定 radius \(r\)，而是引入正 learning rate \(\eta_t\)，令

其中，负号选择方向，\(\eta_t\) 缩放距离。事实上，

所以固定 radius \(r\) 对应 \(\eta_t=r/\lVert g_t\rVert_2\)。代回 Taylor prediction，得到

例如，若 \(g_t=(3,4)\)，并允许 radius \(r=1\)，则 \(\lVert g_t\rVert_2=5\)，所以

它给出的 predicted change 是 \((3,4)^\top(-3/5,-4/5)=-5\)，这是所有 unit-length updates 中可能达到的最小值。若 \(g_t=0\)，first-order model 对每个方向都预测 zero change，因此只靠 gradient information 无法选出 descent direction。

所以 negative gradient 不只是某一个 downhill direction：在普通 Euclidean metric 下、所有固定长度的候选移动中，它能让 local linear model 预测的下降最大。Gradient descent 正是采用这个更新：

这个结论有两个限制。第一，它谈的是 local approximation，而不是完整 landscape，所以并不保证找到 global minimum。第二，它确定了方向，却还没有确定安全的移动距离。过大的 step 会离开 Taylor model 可靠的邻域，反而可能提高真实 loss。剩下的这个决定，就是 learning rate。

1.2 从方向到距离：Learning Rate

Learning rate 把 gradient direction 变成实际位移。它太小会浪费 iterations；太大则会让局部预测失效并产生 overshoot。Quadratic 是 degree-two function，也是具有 nonzero curvature 的最简单函数。一个一维 quadratic 足以让我们在没有 neural-network 细节的情况下看清这个 tradeoff：

Gradient descent 变为

这一个式子已经显示出三种区域：

• 若 \(0<\eta a<1\)，参数不改变符号，逐渐接近零；
• 若 \(1<\eta a<2\)，每一步都会跨过零，但振荡逐渐缩小；
• 若 \(\eta a>2\)，参数绝对值不断增长，训练发散。

这里的 \(a\) 衡量 curvature：\(a\) 越大，loss 向上弯曲得越急。三个区域说明，“gradient 指向下坡”还不够；step 必须小于 loss surface 显著弯曲的距离。在多维空间中，不同方向的 curvature 并不相同，于是我们马上会遇到下一个困难：一个 global learning rate 必须同时服务所有方向。

2. Stochastic Gradient Descent

上面得到的是最干净的 baseline，但它悄悄依赖两个便利条件。第一，\(g_t\) 能忠实描述我们真正关心的 objective；第二，用同一把 Euclidean 尺子衡量每个参数 coordinate 是合理的。Neural-network training 同时违背这两个假设：gradient 来自采样数据的估计，而 loss 在不同方向上的弯曲速度又相差很大。下面把两个困难分开引入，这样每一种 optimizer modification 都有明确的任务。

2.1 Mini-batch 使 gradient 变成估计量

Population objective 对所有可能的数据样本求平均，所以它的精确 gradient 本身也是一个 expectation。真实 data distribution 未知；即使只考虑已经存储的大型数据集，每次 update 前遍历全部样本也非常昂贵。Mini-batch 让一次更新变得可负担：对随机抽取、大小为 \(b\) 的集合 \(B_t\)，我们使用

为什么它是合理的替代？在通常的独立均匀采样下，每个样本都有正确的出现概率。把所有可能抽到的 mini-batches 都考虑进去，平均结果恰好恢复 full dataset gradient：

这里 \(\mid\theta_t\) 表示 expectation 只对随机抽到哪个 mini-batch 求平均，而 \(\theta_t\) 保持固定。

这里的 \(L\) 表示 finite-dataset average。这个等式并没有说 dataset gradient 等于未知的 population gradient；dataset 本身也只是对真实世界的一次采样。因此存在两层 sampling gap：mini-batch 近似 dataset，dataset 再近似 population。

Unbiased 描述的是第一层 gap——对所有假想 mini-batches 求平均后的性质。它并不表示眼前这个 batch 等于 full gradient。一个 batch 可能恰好多含容易样本，另一个多含困难样本，所以即使 \(\theta_t\) 相同，它们的 gradients 也可能不同。增大 batch size 通常会减少这种 sampling variation，却需要每次更新使用更多计算。

Tradeoff 至此已经清楚：full-batch gradient 要等到更稳定的方向才移动，但每次等待较久；mini-batch gradient 便宜且频繁，但每个 estimate 都带有不确定性。把这个 estimate 与 gradient-descent update 结合起来，就得到 SGD；本 section 稍后会写出完整 update。之后 Momentum 再利用历史信息，把持续存在的 signal 与 batch-to-batch fluctuation 区分开来。

2.2 Plain SGD 为什么需要更多 context？

即使拥有精确 full-batch gradient，一个 global learning rate 仍可能效率很低。Section 1 只用 first-order Taylor term 选择方向；要知道这个方向在多远的范围内仍然可信，就必须观察下一项：

Hessian \(H_t\) 是由 second derivatives 组成的 matrix。它记录 curvature，也就是 gradient 本身会怎样随移动而改变。在 smooth local minimum \(\theta^*\) 附近，gradient 约等于零，于是留下常见的 quadratic model：

\(H\) 的 eigenvector 是一个不会与其他方向混合的 quadratic direction；相应 eigenvalue \(\lambda_i\) 就是该方向的 curvature。Gradient descent 在这个方向上的行为，严格对应前面 \(a=\lambda_i\) 的一维例子。最大 curvature 限制稳定性，最小 curvature 则控制平坦方向上的前进速度。

Figure 1 使用的 objective 可以把问题具体化：

同样大小的位移发生在 \(y\) 方向时，对 loss 的影响是 \(x\) 方向的十二倍。为了在 \(y\) 方向稳定，必须有 \(\eta<2/12=1/6\)；可一旦遵守这个限制，在更平坦的 \(x\) 方向上，update 又必然不大。因此路径容易横跨狭长山谷来回反弹，却沿着山谷前进缓慢。更一般地，较大的 condition number

描述的正是这种尺度分离。

至此可以看到两种不同的信息缺失。当前 mini-batch gradient 没有告诉我们哪些 components 在下一个 batch 中仍会存在；raw gradient values 也没有告诉我们，一个很大的 component 究竟代表有用的移动方向，还是仅仅来自尺度不同、curvature 很大的方向。

这里的 coordinate 是 parameter vector 中的一个 scalar 位置。若

那么 \(g_{t,i}=\partial L/\partial\theta_i\) 就是参数 \(\theta_i\) 对应的 gradient component。“跨 coordinates”不是在比较不同 data examples，而是在比较同一个 parameter vector 的各个 components 及其历史。比如连续几个 steps 得到

第一个 coordinate 的 gradient magnitude 一直在 \(100\) 左右，第二个则在 \(0.01\) 左右。Plain SGD 对两者乘以同一个 learning rate，所以第一个 coordinate 的 update 会大数千倍。但 gradient 长期较大，并不能单独证明这个参数就应该移动得更远；差异也可能来自参数单位、network parameterization 或 curvature。因此，optimizer 可以积累两类 context：

• 跨 steps：某个 gradient component 是否一直指向同一方向，还是它的正负号在不断波动？
• 跨 coordinates：不同参数的 gradient magnitude 是否长期处于不同尺度，因此需要分别调整它们的 update scale？

Momentum 使用第一类历史信息。RMSProp 为每个 coordinate 分别记录典型的 squared-gradient magnitude，从而使用第二类信息；Adam 同时使用两类。

2.3 Optimizer 是带状态的更新规则

Stateless rule 只看到当前 \(g_t\)。它无法区分某个连续二十步都指向同一方向的 component，和另一个每一步都反转符号的 component；也不知道某个 coordinate 是否一贯很大。为了作出这些区分，optimizer 必须把过去的信息带到下一步。可以把这类 optimizer 抽象写成

其中 \(s_t\) 是 optimizer state，也就是对 gradient history 的紧凑摘要；\(F\) 更新这份记忆，\(U\) 再把记忆与当前 gradient 变成 parameter change。不带 momentum 的 SGD 没有持久的逐参数状态；Momentum 保存一个 moving average；Adam 保存两个。这些 buffers 会跨 training steps 保留。若想完全复现 resumed training 的 dynamics，就必须把它们一起 checkpoint。它们的显存成本在 LLM Optimization Basics: Memory 中单独讨论。

2.4 Plain SGD Algorithm

现在，定义 standard baseline 所需的两个 ingredients 已经齐全：gradient descent 说明拿到 gradient 后怎样更新，mini-batching 说明怎样便宜地估计这个 gradient。Stochastic gradient descent（SGD）把二者结合起来：

一次 SGD step 有明确的执行顺序：

1. 随机抽取 mini-batch \(B_t\)；
2. 运行 model，并平均该 batch 上的 per-example losses；
3. 用 backpropagation 计算 mini-batch gradient \(g_t\)；
4. 将它乘以 \(-\eta_t\)，再更新 parameters。

Stochastic 一词指的是第一步：batch 是随机的，因此 \(g_t\) 与最终 update 也是 random variable。一旦 \(B_t\)、\(\theta_t\) 与 \(\eta_t\) 固定，plain SGD rule 本身就是 deterministic。若 mini-batch estimator unbiased，并且在 conditional expectation 中 learning rate 固定，则

所以，一次具体 update 不一定降低 full objective，但它的 expected local direction 与 full-batch gradient descent 相同。计算上的好处是：一步只处理 \(b\) 个 examples，而不是整个 dataset；代价则是 sampling noise。

Plain SGD 不保存 moving average 或 per-coordinate scale：除了 parameters 本身以及外部 step counter 或 schedule，下一次 update 只使用当前 batch gradient。这使它成为显存需求低、容易理解的 baseline，同时也让上面指出的两个困难原封不动：batch fluctuations 直接进入 update，而且一个 global learning rate 仍要服务所有 coordinates。

Optimizer step 与 epoch 是不同单位。一个 step 消耗一个 effective batch，并更新一次参数；一个 epoch 大致处理一遍数据集。使用 gradient accumulation 时，parameters 保持不变，先把 \(K\) 个 microbatch gradients 合并，通常写作 \(g_t=K^{-1}\sum_{k=1}^K g_t^{(k)}\)，随后才执行一次 optimizer update。这样无需同时保存所有 examples 的 activations 就能扩大 effective batch，但它不会产生 \(K\) 个 optimizer steps。

2.5 SGD with Momentum：跨 Steps 的一致性

可以把 mini-batch gradient 非正式地拆成“持续存在的方向 + sampling fluctuation”。我们无法直接看到这两部分，但反复观察会提供线索：持续保持同一符号的 component 更像可靠 signal；快速交替正负的 components 则会在平均时互相抵消。

若简单平均过去所有 gradients，训练越久，optimizer 的反应就越迟钝，因为在完全不同参数位置算出的古老 gradients 与最新 gradient 仍有同等地位。Momentum 改用 exponential moving average（EMA）：赋予最新 gradient 权重 \(1-\beta\)，并把已有历史乘以 \(\beta\) 逐渐淡忘：

第一条式子更新记忆，第二条式子用这个平滑方向取代 raw current gradient。展开 recurrence，就能看到每个历史 gradient 的精确权重：

其中 \(m_0=0\)。只看当前 buffer \(m_t\) 内的显式加权和，越近的 gradient 具有越大的直接系数：一个距今 \(j\) 步的 gradient，其系数是 \((1-\beta)\beta^j\)。Characteristic averaging window 大致为 \(1/(1-\beta)\) steps：\(\beta=0.9\) 时约为 \(10\)，\(\beta=0.99\) 时约为 \(100\)。这只是尺度而不是突然截断的边界；更古老的 gradients 不会在某一步骤凭空消失。

这不等于说最新 data 对整条 training trajectory 的影响必然最大。在第一次 update 时，

此时 \(g_1\) 是唯一被观察到的方向，所以方向完全由它决定；在本文的 convention 下，系数 \(1-\beta\) 只会缩小 buffer magnitude。随后第一次 update 已经改变 \(\theta_1\)，而新的参数位置又会改变此后算出的所有 gradients。因此，即使早期 gradient 在后来 buffer \(m_t\) 中的直接系数已经衰减，它仍可能通过 parameters 对后续 trajectory 保留长久的间接影响。上面的展开式只描述前一种影响。

“Momentum” 这个名称来自与 velocity 的定性类比：update 会把一部分过去方向带到下一步，而不是只响应类似 force 的 current gradient。这个公式是 optimization rule，并不是对真实物理力学的逐项模拟。

为什么它能帮助狭长山谷中的优化？在陡峭方向上，相邻 gradients 经常改变符号，因此会在 moving average 中相互抵消；在平坦方向上，它们的符号长期一致，因此能通过 averaging 保留下来。Momentum 抑制了横跨山谷的来回振荡，同时保留沿山谷前进的 motion。在本文采用的 normalized EMA convention 下，额外 acceleration 是否出现取决于 learning rate 与完整 dynamics，而不是 buffer magnitude 无限制地增长。

3. Adam

3.1 Coordinate-wise Scaling

Momentum 处理的是 gradient evidence 在时间上是否一致，却仍以 raw gradient units 衡量每个 coordinate。回到狭长山谷：陡峭 coordinate 的 gradient 可能数值很大，即使参数已经接近谷底；平坦 coordinate 的 gradient 虽小，却可能是仍需大幅前进的方向。

这是否与“\(-g_t\) 是最优方向”矛盾？ 不矛盾。Section 1 证明的是一个更窄的结论：只使用 linear Taylor model，并让所有候选 updates 具有相同 Euclidean length 时，\(-g_t\) 能带来最大的即时预测下降。Gradient magnitude 衡量的是 slope——每单位参数变化会造成多少 loss 变化——而不是当前位置到 minimum 的距离，也没有说明这个 slope 在多远的范围内仍然成立。

一个一维 quadratic 会直接暴露缺少的信息：

同样观察到 \(g=1\)，当 \(a=1\) 时可能有 \(\theta=1\)；当 \(a=100\) 时却只有 \(\theta=0.01\)。两种情况下 gradient 的符号和大小完全相同，但当前位置到 minimum 的距离相差一百倍。只有再知道 curvature \(a\)，才能得到适合这个 quadratic 的 update：

在多维 quadratic 中，当 Hessian \(H\) positive definite 时，对应的 update 是 \(-H^{-1}g\)，它通常不与 \(-g\) 平行。因此，coordinate-wise scaling 不是在修补一个“忘记大小”的 gradient；它是在利用额外的 scale information，重新判断每个 gradient component 应该转换成多远的参数移动。Adaptive optimizers 并不知道精确 Hessian，而是把 historical squared-gradient statistics 当作一种实用的 diagonal scaling heuristic。

这里的 global multiplier 是所有 coordinates 共用的一个 scalar learning rate，并不是指“能够找到 global optimum 的 optimizer”。在 plain gradient descent 中，

因此每个 coordinate 都有 \(\Delta\theta_i/g_i=-\eta\)。改变 \(\eta\) 只能把整个 update 一起放大或缩小；它不能相对于 gradient 单独让陡峭 coordinate 更谨慎，同时让平坦 coordinate 更大胆。考虑

在 \((x,y)=(0.1,1)\) 处，gradient 是 \((10,1)\)，但到 minimum \((0,0)\) 的剩余距离在 \(x\) 方向只有 \(0.1\)，在 \(y\) 方向却有 \(1\)。共用 multiplier 给出

无论怎样选择 \(\eta\)，两个 coordinates 的 update ratio 都固定为 \(10:1\)。为了不在陡峭的 \(x\) 方向上震荡甚至发散，\(\eta\) 必须较小；同一个小数也会让平坦的 \(y\) 方向前进缓慢。Momentum 可以在时间上平滑符号变化，但它的 update \(-\eta m_t\) 仍然是把一个 scalar 乘到保留 raw gradient units 的整个 vector \(m_t\) 上。若要区分 coordinate scales，就需要一组逐坐标的 effective learning rates；等价地，需要为每个 coordinate 分别估计 scale。

Adaptive methods 因此为每个 coordinate 建立独立的近期尺度，再用 gradient 除以这个尺度。为了估计 magnitude 时不让正负 gradients 互相抵消，它们累计 \(g_t^2\)，而不是 \(g_t\)；随后开平方，把 squared quantity 恢复到 gradient units。于是“当前 gradient / 典型 gradient magnitude”中的 units 会互相消去；这个 ratio 在乘上 global learning rate 前是 dimensionless，即不再带 physical unit。

AdaGrad：累积 gradient squares

最直接的版本是 AdaGrad。从 \(v_0=0\) 开始，它把迄今见过的每个 squared gradient 都加起来：

这里 \(\odot\) 表示 element-wise multiplication（逐元素相乘），不是 dot product 或 matrix multiplication。若

则

因此 \(v_t\) 与 \(g_t\) shape 相同；coordinate \(v_{t,i}\) 只累积参数 \(\theta_i\) 过去得到的 squared gradients。随后更新

Update 中的平方根、加法与除法也都逐元素执行。只写 coordinate \(i\) 时，公式就是

因此 coordinate \(i\) 的 effective learning rate 是 \(\eta/(\sqrt{v_{t,i}}+\epsilon)\)。一个反复收到大 gradient 的 coordinate 会积累较大的 denominator，之后采用更小的 step；rare feature——只在很少样本中 active 的 feature——则在多数 steps 产生 zero gradient，并保留相对较大的 effective step。因此 AdaGrad 很适合 sparse data。小正数 \(\epsilon\) 避免某个 coordinate 尚未积累任何 scale 时发生除零。

它的弱点也直接写在 recurrence 中：\(v_t\) 只增不减。Effective learning rates 会不断缩小，在长期、nonstationary 的 neural-network training 中可能变得过小。

RMSProp：忘记过时的 gradient squares

AdaGrad 让训练第一步的 gradient 永久参与尺度估计；但 model 和 loss geometry 会在训练中改变，旧 scale 可能已经过时。RMSProp 只作一个关键修改：用 exponential moving average 替代永不遗忘的 cumulative sum：

Denominator 估计每个 coordinate 最近的 root-mean-square gradient。除以它可以减弱长期尺度差异：经过足够多 steps 后，把某个 coordinate 的 gradients 乘以正数，通常会让 numerator 与 denominator 同时乘以这个数，normalized update 近似不变。但 initialization、\(\epsilon\)、momentum、clipping 与 nonstationarity 都会破坏严格等价。

\(v_t\) 称为 uncentered second moment，因为 moment 是某个幂次的平均：first raw moment 平均 \(g\)，second raw moment 平均 \(g^2\)。它不是统计学 variance；后者是 \(\mathbb E[g^2]-\mathbb E[g]^2\)。

3.2 结合 Momentum 与 Coordinate-wise Scaling

Momentum 与 RMSProp 分别改造了 plain SGD 的两个独立部分。Momentum 用平滑方向替代 noisy numerator；RMSProp 保留当前 numerator，却除以近期的 coordinate-wise scale。Adam 同时采用两种想法：平滑分子中的方向，再用分母中的平滑 magnitude 对它归一化。

因此，它维护两个 buffers：

这里的 “first moment” 与 “second moment” 指 observed gradient stream 的平均：\(m_t\) 平均 \(g_t\)，\(v_t\) 平均 \(g_t^2\)。它们不是 training-data distribution 或 model predictions 的 moments。

两个 buffers 都从零开始。在最初几步，设想中的 averaging window 有一部分仍由零、而不是实际 observation 填充，因此数值会被拉向零。Adam 对这个 initialization bias 作修正：

更新为

现在，每一部分都有不同角色：

\(v_t\) 会不会影响 update direction？对单个 scalar coordinate 来说，它的 denominator 非负，因此不会改变由 \(m_t\) 决定的正负号，只会改变该 coordinate 的 magnitude。但对完整 parameter vector 来说，\(v_t\) 的不同 entries 会用不同倍数缩放各 coordinates，所以通常会旋转整个 update。忽略 \(\epsilon\) 时，例如

它与 \(-\hat m=(-1,-1)\) 并不平行。Numerator 提供每个 coordinate 的正负号；denominator 改变 coordinates 之间的相对 magnitudes。

下面的 pseudocode 按执行顺序写出一个完整 Adam step。将鼠标移到高亮符号上——或使用键盘 focus——即可查看该变量存放的内容以及它在 update 中的作用。

θ ← initial_parameters()
m ← zeros_like(θ)
v ← zeros_like(θ)
t ← 0

repeat，直到满足停止条件：
    t ← t + 1
    B ← sample_minibatch()
    g ← minibatch_gradient(θ, B)
    m ← β₁·m + (1−β₁)·g
    v ← β₂·v + (1−β₂)·(g ⊙ g)
    m̂ ← m / (1−β₁^t)
    v̂ ← v / (1−β₂^t)
    θ ← θ − ηₜ·m̂ / (sqrt(v̂) + ε)

return θ

3.3 为什么需要 Bias Correction？

假设前几步的 gradient 都是同一个 vector \(g\)。从零开始有

为什么 denominator 恰好是这个形式？截至第 \(t\) 步，EMA 中已经落在真实 observations 上的 weights 之和是 \(1-\beta^t\)，而不是一。除以这项，就是重新把已有 weights normalize 到和为一。等价地，在 constant-gradient example 中，分别除以 \(1-\beta_1^t\) 与 \(1-\beta_2^t\) 会恢复 \(g\) 与 \(g^2\)。History 逐渐填满后 correction 趋近一，因此它主要影响训练早期。

第一步还能看出一个重要结果。忽略 \(\epsilon\)，逐坐标有

Adam 的第一次更新主要依赖 nonzero gradient coordinates 的符号，而不是原始 magnitude。因此，Adam learning rate 与 SGD learning rate 不能按数值直接互换。

这里 \(\operatorname{sign}(a)\) 在 \(a>0\) 时为 \(+1\)，在 \(a<0\) 时为 \(-1\)，在 \(a=0\) 时为 \(0\)。因为实际公式中的 \(\epsilon\) 会阻止完全约分，所以对 nonzero coordinates 而言，这个结论是 approximation。

3.4 Adam 解决了什么，又没有解决什么？

当 gradients noisy、sparse，或各 coordinates 的尺度差异很大时，Adam 往往很有帮助；它通常比 plain SGD 更容易进入有效训练区间。但它的 denominator 只能独立缩放 coordinates，不能旋转 update 来消除参数之间的 off-diagonal interactions，也无法推断 local approximation 在多远范围内可信，更不保证找到更好的 minimum。较小的近期 \(v_{t,i}\) 还会放大 coordinate \(i\)：floor \(\epsilon\) 限制这种放大；learning-rate warmup 在开始时刻意使用较小 steps；gradient clipping 限制异常大的 gradient norm；足够的 numerical precision 则让这些小 statistics 仍可表示。

常见默认值 \(\beta_1=0.9\)、\(\beta_2=0.999\) 和较小的 \(\epsilon\) 只是起点，不是定律。更大的 \(\beta_2\) 使 scale estimate 拥有比 direction estimate 更长、更稳定的 averaging window。Batch size、loss normalization、model scale 或 precision 改变后，合适的 learning rate 乃至 moment constants 都可能改变。

4. AdamW

到这里为止，每一种 modification 都在尝试更有效地利用 data gradient。Weight decay 问的是另一个问题：除了拟合 training objective，我们是否还希望偏好 norm 较小的参数？这种偏好可以充当 regularization 或 norm control，虽然实际效果取决于 architecture。因为它恰好也在 update 时执行，所以很容易与 optimizer 本身混为一谈。

表达这种偏好的传统方式，是向 loss 加入一个 \(L_2\) penalty：

其 gradient 为

对 plain SGD，把这个 gradient 代入 update，再整理两项，可得

系数 \(1-\eta\lambda\) 会在这一步把每个被选中的参数按同一比例收缩。这就是为什么对 plain SGD，\(L_2\) regularization 与 weight decay 等价。

但对 Adam，把 \(\lambda\theta\) 塞进 gradient 后，penalty 会经过 first- and second-moment normalization；某个 coordinate 上的 penalty 会被 history-dependent scale 除掉，于是不同 coordinates 得到不同的 effective shrinkage。因此，coupled \(L_2\) regularization 不再等于 uniform weight decay。

AdamW 将 loss gradient 与 decay 分开：

Moments 根据 data gradient \(g_t\) 计算，而不是根据 \(g_t+\lambda\theta_t\) 计算。这就是 decoupled weight decay。

下面是最小的 array-level 实现。Production model 通常把 parameters 存成 parameter tree，即由名称和层级组织起来的一组 arrays；distributed training 还要通过 gradient reduction 合并多个 devices 各自计算的 gradients。这个例子省略这些结构以及 mixed precision 与 fused kernels，以便直接看到 update equation。

import jax.numpy as jnp

def adamw_step(params, grads, state, lr,
               beta1=0.9, beta2=0.999,
               eps=1e-8, weight_decay=0.01):
    t = state["t"] + 1
    m = beta1 * state["m"] + (1.0 - beta1) * grads
    v = beta2 * state["v"] + (1.0 - beta2) * jnp.square(grads)

    m_hat = m / (1.0 - beta1 ** t)
    v_hat = v / (1.0 - beta2 ** t)

    params = ((1.0 - lr * weight_decay) * params
              - lr * m_hat / (jnp.sqrt(v_hat) + eps))
    return params, {"t": t, "m": m, "v": v}

这里假设 \(t\) 从零开始计算 optimizer updates，而且 \(m,v\) 与 parameter array 形状相同。Production libraries 可能把 \(\epsilon\) 放在不同位置，使用 AMSGrad，采用 capturable step counter，或 fuse decay 与 update。所谓不同位置，是选择 \(\sqrt{\hat v}+\epsilon\) 还是 \(\sqrt{\hat v+\epsilon}\)；AMSGrad 用 running maximum 替代当前 second-moment denominator；capturable counter 把 \(t\) 存在 compiled accelerator graph 可访问的位置；fused kernel 则把多个 array operations 合在一起，以减少 memory traffic。这些选择保留了中心思想，却会影响精确 numerical reproduction。

5. Muon

Adam 的 scale estimate 是 coordinate-wise 的：即使把 weight matrix 拉平成一个长 vector，它仍会逐元素作同样的事。但 linear layer 中的 matrix 并非偶然以二维形式存储。若 \(W\) 把 activation \(h\) 映射为 \(Wh\)，它的 rows 与 columns 共同描述 input directions 和 output directions。一个 matrix update 可以在某一对方向上很强，在另一对方向上很弱。独立处理每个 entry 会丢掉这种结构。

Muon 先提出与 Momentum 相同的时间问题，再改变所得 matrix update 的几何结构。对 weight \(W\in\mathbb R^{m\times n}\)，其 gradient \(G_t\) 与 momentum \(M_t\) 具有相同 shape。首先形成 momentum matrix：

要理解下一步，先回顾 singular value decomposition 的含义。任意 matrix 都可写成

\(V\) 的 columns 给出彼此 orthogonal 的 input directions；\(U\) 的 columns 给出相应 output directions；\(\Sigma\) 对角线上的非负数——singular values——表示 update 沿每对 directions 作用得多强。若某个 singular value 远大于其他值，这一个 mode 就会支配 matrix update。

理想 orthogonalization 保留两组 directions，却把每个 nonzero singular value 换成一：

对 rectangular matrix，更准确的名称是 semi-orthogonal；关键结论相同：所有 nonzero singular modes 现在具有相同强度。这是一种 matrix-level normalization。它不是在断言真实 loss 的每个 singular direction 都有相同 curvature，而是根据 \(M_t\) 中已有的信息，选择一个更均衡的 candidate update。

若每一步都对每个 eligible weight 真正计算完整 SVD，代价会很高。Muon 会先 normalize 初始 matrix，再执行少量 Newton–Schulz-style polynomial iterations。每次 iteration 只使用 matrix multiplications，并把 singular values 推向一，从而在不显式构造 \(U\)、\(\Sigma\) 与 \(V\) 的情况下近似 \(UV^\top\)；之后再应用 recipe 指定的 dimension-dependent scale 与 learning rate。因此，实际 Muon configuration 不只有 \(UV^\top\) 这个符号；还必须说明 momentum convention、orthogonalization approximation、scaling、decay 与 parameter grouping。

这与 Adam 是两类不同的 adaptivity：

• Adam 根据各 scalar coordinate 的历史 squared gradients 分别缩放；
• Muon 根据 singular directions 重塑完整 matrix update。

为什么不把它用于所有 parameters？Bias 或 normalization gain 没有可以 orthogonalize 的二维 input–output structure；embedding matrix 负责把 discrete token IDs 映射为 hidden vectors，output head 则把 hidden vectors 映射回 token scores，它们的 shape、共享方式与 gradient statistics 都可能不同于内部 linear layers，因此同一种 scaling rule 未必合适。Muon 通常只用于符合条件的二维 hidden-layer weights，而 embeddings、output heads、normalization gains、biases 和其他 vector 或 scalar parameters 仍交给 AdamW。

这里介绍 Muon，是因为它能让 optimizer design space 更清楚，而不是断言它普遍优于 AdamW。Orthogonalization 会增加 matrix multiplications；当一个 matrix 分布在多个 devices 上时，这些运算还可能需要 devices 之间交换数据，因此 operation count 本身不能代表实际训练时间。公平比较应固定或明确报告总计算量或 wall-clock time、处理过的 tokens、learning-rate schedule，以及每种方法获得的 hyperparameter tuning 资源。

• Polyak, B. T. (1964). Some Methods of Speeding Up the Convergence of Iteration Methods. USSR Computational Mathematics and Mathematical Physics, 4(5), 1–17.
• Duchi, J., Hazan, E., & Singer, Y. (2011). Adaptive Subgradient Methods for Online Learning and Stochastic Optimization. JMLR, 12, 2121–2159.
• Tieleman, T. & Hinton, G. (2012). RMSProp, Lecture 6.5 of Neural Networks for Machine Learning.
• Kingma, D. P. & Ba, J. (2015). Adam: A Method for Stochastic Optimization. ICLR.
• Loshchilov, I. & Hutter, F. (2019). Decoupled Weight Decay Regularization. ICLR.
• Jordan, K. (2024). Muon: An Optimizer for Hidden Layers in Neural Networks.
• Liu, J. et al. (2025). Muon is Scalable for LLM Training.

• Polyak, B. T. (1964). Some Methods of Speeding Up the Convergence of Iteration Methods. USSR Computational Mathematics and Mathematical Physics, 4(5), 1–17。
• Duchi, J., Hazan, E., & Singer, Y. (2011). Adaptive Subgradient Methods for Online Learning and Stochastic Optimization. JMLR, 12, 2121–2159。
• Tieleman, T. & Hinton, G. (2012). RMSProp, Lecture 6.5 of Neural Networks for Machine Learning。
• Kingma, D. P. & Ba, J. (2015). Adam: A Method for Stochastic Optimization. ICLR。
• Loshchilov, I. & Hutter, F. (2019). Decoupled Weight Decay Regularization. ICLR。
• Jordan, K. (2024). Muon: An Optimizer for Hidden Layers in Neural Networks。
• Liu, J. et al. (2025). Muon is Scalable for LLM Training。