Optimizer update

An optimizer update changes the adjustable parameters of an optimization problem according to an algorithm's current information. Evaluating a loss function, computing its gradient and changing the parameters are distinct operations. A general update may also change an optimizer state, such as stored information from previous steps:Here $g_k$ is the supplied gradient and ${\eta }_k$ a learning rate. The PyTorch interface makes this separation explicit with gradient computation followed by optimizer.step().^[ref]

For plain gradient descent with no extra state, ${\theta }_{k+1}={\theta }_k-\eta g_k$. In an original scalar example, choose $f(\theta )={\theta }^2/2$, ${\theta }_0=2$ and $\eta =0.1$. The gradient is 2, so the update gives ${\theta }_1=1.8$ and lowers the loss from 2 to 1.62. Recompute the gradient at the new parameter value to obtain ${\theta }_2=1.62$. Reusing the old gradient 2 instead would give 1.6. This explicitly locates where the algorithm obtains its information.^[ref]

An update is not automatically an improvement: using $\eta =3$ for the first step above gives ${\theta }_1=-4$ and loss 8. With gradient accumulation, several gradient calculations instead contribute before a single optimizer update; their weighting sets the effective batch size. Track the update count separately from the number of processed groups, particularly when applying a learning rate schedule.

Weight decay supplies a parameter-shrinkage contribution to an optimizer update, with a coefficient commonly denoted $\lambda$. Two implementations must be distinguished. Coupled L2 regularization adds the penalty $\lambda \Vert \theta {\Vert }^2/2$ to the objective function, replacing its data gradient $g$ by $g+\lambda \theta$. Decoupled weight decay applies parameter shrinkage separately from the optimizer's processing of the data gradient. These operations are equivalent for plain gradient descent with matched coefficients, but not in general for an adaptive gradient method.^[ref]

Use the PyTorch convention here: a decoupled update shrinks parameters by $1-\eta \lambda$, where $\eta$ is the learning rate. A source writing $1-\lambda$ instead has absorbed the step-size factor into its decay coefficient. For an unmodified gradient step,In AdamW, the decay contribution does not enter the stored first- or second-moment estimates of the gradient.^[ref]

An original fixed-scaling example exposes the difference. Let $\theta =(2,2{)}^T$, $g=(1,1{)}^T$, $\eta =0.1$, $\lambda =0.2$ and diagonal matrix $A=\mathrm{diag}(1,10)$. Applying the same scaling to the entire coupled gradient givesSeparating uniform parameter shrinkage givesThe second coordinate receives different shrinkage. This is an illustrative fixed matrix, not a full AdamW simulation. With history-dependent scaling, coupling the penalty into the gradient can also alter the optimizer state. State the update equation and coefficient convention when reporting a decay setting; the label alone does not specify either.

Decoupled weight decay applies a parameter-shrinkage contribution separately from the optimizer's processing of the data gradient. Write $u_t$ for the processed data-gradient direction, including any algorithm-specific optimizer state. Using the coefficient convention documented for PyTorch AdamW,The shrinkage term uses the pre-update parameter and is not inserted into the gradient moments.^[ref]

This differs from coupled L2 regularization, which processes $g_t+\lambda {\theta }_{t-1}$ as a single gradient input. Decoupled weight decay was introduced for adaptive optimization in the AdamW work; the original paper also writes decay coefficients with the learning-rate factor absorbed differently, so its coefficient must be translated before comparing equations.^[ref]

An original limiting case isolates the shrinkage. Suppose the processed direction is zero at every step and $\eta ,\lambda$ are constant. Repeated substitution givesWith $\eta =0.1$, $\lambda =0.2$ and ${\theta }_0=1$, one update gives 0.98 and 50 updates give approximately 0.36417. For varying learning rates, the corresponding factor isThus the same numerical decay coefficient can have a different accumulated effect under a different learning rate schedule. In this zero-direction model, $0<\eta \lambda <1$ gives sign-preserving shrinkage; other factors have different behaviour.

The update order is also part of the definition. For $\theta =1$, $u=2$, $\eta =0.1$ and $\lambda =0.2$, the stated formula gives $0.98-0.20=0.78$. Applying shrinkage after first taking the gradient step would instead give $(1-0.20)0.98=0.784$. Both operations may be described informally as decay plus a step, but they are different recurrences. An optimization checkpoint must retain the relevant state and settings to reproduce the chosen one.

Coupled L2 regularization adds a quadratic parameter penalty to the objective function and sends its gradient through the same optimizer update as the data gradient. With coefficient $\lambda \ge 0$, defineSince the gradient of the squared Euclidean norm divided by two is $\theta$, a supplied data gradient $g=\nabla f(\theta )$ becomesThe factor $1/2$ fixes the convention: a penalty written as $\lambda \Vert \theta {\Vert }_2^2$ instead contributes $2\lambda \theta$.

For plain gradient descent, this substitution givesIt is therefore equivalent to decoupled weight decay with the matched coefficient in this particular update. The equivalence need not survive coordinate-dependent processing of the gradient.^[ref]

For an original fixed-scaling example, let $\theta =(1,1{)}^T$, $g=0$, $\eta =0.1$, $\lambda =0.2$ and diagonal matrix $D=\mathrm{diag}(1,0.1)$. Applying $D$ to the whole coupled gradient givesThe same quadratic penalty produces unequal coordinate shrinkage after this scaling. Uniform decoupled weight decay would instead give $(0.98,0.98{)}^T$. The example isolates a fixed scaling operation, rather than claiming to simulate a complete adaptive gradient method.

Coupling can also change the optimizer state. For example, a momentum buffer formed from $\tilde{g}$ stores contributions from $\lambda \theta$ along with the data gradient. If an adaptive accumulator uses squared gradient entries, it receives $(g_i+\lambda {\theta }_i{)}^2$, including a cross term $2\lambda g_i{\theta }_i$. Subsequent updates can therefore differ even when the next data gradient is the same. State both the penalty convention and the optimizer's gradient-processing rule when describing an L2 setting.

The optimizer state is information retained by an optimization algorithm beyond the current parameter values. It can include a momentum buffer and algorithm-specific counters or running estimates. Consequently, two runs with identical parameters and current gradients can produce different next optimizer updates if their stored states differ.

For an original scalar example, minimize $f(\theta )={\theta }^2/2$ with the recurrenceThis is the momentum form with dampening, decay and the nesterov flag disabled in the PyTorch stochastic gradient descent implementation.^[ref] At ${\theta }_k=1$, $\eta =0.1$ and $\mu =0.9$, a stored $b_k=0$ gives $b_{k+1}=1$ and ${\theta }_{k+1}=0.9$. A stored $b_k=2$ instead gives $b_{k+1}=2.8$ and ${\theta }_{k+1}=0.72$. The current parameter and gradient were identical; the difference came from history.

An optimization checkpoint that saves only the parameters cannot reproduce that next step. In PyTorch, Optimizer.state_dict() stores per-parameter state and parameter-group settings such as learning rate, but not the parameter values themselves. Its parameter identifiers associate state with ordered parameter groups and are not a check of semantic parameter identity during loading.^[ref] Save and restore the corresponding parameters separately and keep the state-to-parameter association consistent. Reproducing an entire training run may additionally require the data order and other changing state.

An optimization checkpoint is a saved description of an ongoing optimization algorithm from which its optimizer updates can resume. It combines the current parameters with the relevant optimizer state. The practical criterion is optimizer restart equivalence: after restoring the saved information and supplying the same subsequent inputs, the chosen update rule should produce the same subsequent states.

For a deterministic scalar example, minimize the quadratic function $f(\theta )={\theta }^2/2$ with gradient $g_k={\theta }_k$, learning rate $\eta =0.1$ and momentum buffer recurrence $b_{k+1}=0.9b_k+g_k$, followed by ${\theta }_{k+1}={\theta }_k-0.1b_{k+1}$. Start with $({\theta }_0,b_0)=(1,0)$. After one step the saved pair is $({\theta }_1,b_1)=(0.9,1)$. Restoring both gives $b_2=1.8$ and ${\theta }_2=0.72$. Restoring only ${\theta }_1$ and initializing a new momentum buffer to zero gives $b_2=0.9$ and ${\theta }_2=0.81$. Equal saved parameters can therefore lead to unequal continuations. These numbers are an original calculation for the declared recurrence.

In PyTorch, a general training checkpoint commonly serializes separate model and optimizer dictionaries, together with the training position. The model dictionary contains learned parameters and registered model buffers; the optimizer dictionary contains its own changing state and group settings.^[ref] Restoring the optimizer state-to-parameter association also matters: the optimizer loader pairs stored identifiers with parameters in group order, without independently establishing that they denote the same mathematical parameters.^[ref]

The required checkpoint boundary depends on the computation being resumed. A changing learning rate schedule, the next position in the training data, and the state of a pseudorandom number generator can influence future updates. Matching them is part of computational reproducibility; a saved parameter vector alone does not specify the whole training process.

Optimizer state-to-parameter association identifies which parameter each saved piece of optimizer state belongs to. An optimization checkpoint can contain every stored number and still resume incorrectly if those numbers are attached to different parameters. Equal shapes and data types do not establish mathematical identity.

The PyTorch optimizer dictionary separates per-parameter state from parameter-group settings. Its loader pairs saved parameter identifiers with the current parameters in group order, without independently verifying the intended association.^[ref]

For an original two-parameter example, use $b^{\prime }=0.5b+g$ and ${\theta }^{\prime }=\theta -0.1b^{\prime }$. Let the named parameters be $a=1$ and $c=10$, their current gradients both be 1, and their saved momentum buffers be $b_a=2$ and $b_c=20$. Correct restoration givesIf a reversed parameter ordering attaches the same saved buffers to the wrong scalar parameters, the result is insteadBoth buffers had valid scalar shapes, so a shape check would not expose this error.

Consistently permuting the parameters, gradients and state is merely a change of ordering; permuting only the state changes the algorithm's continuation. Record a stable parameter schema and group order with the checkpoint, and check any intentional remapping against that schema. Parameter values and optimizer dictionaries are distinct saved objects, so restoring one does not automatically restore the other's meaning.^[ref]

Optimizer restart equivalence means that restoring an optimization checkpoint reproduces the uninterrupted optimizer's parameter sequence for a specified subsequent computation. State the optimizer update rule, continuation inputs and intended notion of equality. Exact mathematical equality and agreement within a numerical tolerance are different claims.

For a deterministic algorithm modeled by $S_{k+1}=F(S_k,u_k)$, equality of the restored full state at step $K$, together with the same update rule and future inputs $u_K,u_{K+1},\dots$, is sufficient. One application of $F$ gives equal next states; mathematical induction then gives equal continuations. This is a sufficient condition for the chosen observations, not a claim that every internal field must always affect them.

For an original exact example, take $f(\theta )={\theta }^2/2$ andFrom $({\theta }_0,b_0)=(2,0)$, one step gives $({\theta }_1,b_1)=(3/2,2)$. Restoring that pair gives $b_2=5/2$ and ${\theta }_2=7/8$, exactly as uninterrupted execution. Resetting the buffer while restoring ${\theta }_1$ gives $b_2=3/2$ and ${\theta }_2=9/8$. Thus a parameter-only save is not restart-equivalent for this recurrence.

Practical checkpoints must also retain the optimizer state-to-parameter association and the relevant settings, update position, learning rate schedule, data position and pseudorandom number generator state.^[ref] Fixing a random seed alone does not specify all of this continuation. Floating-point arithmetic and nondeterministic operations can additionally affect exact reproduction across environments; the primary documentation does not guarantee identical results across releases or platforms.^[ref]

A momentum buffer stores an accumulated history of gradients for an optimization algorithm. In one common convention, a parameter vector ${\theta }_k$, supplied gradient $g_k$, learning rate ${\eta }_k$ and constant coefficient $0\le \mu <1$ obey the recurrence relationThe buffer belongs to the optimizer state: it remembers earlier inputs even when the current parameters are unchanged. This is the undampened momentum convention used by PyTorch stochastic gradient descent with weight decay and the Nesterov accelerated gradient method disabled; its documented first buffer is the first gradient, agreeing with this recurrence from $b_0=0$. Other conventions absorb the learning rate into the stored buffer, so a buffer value needs its update convention to be interpreted.^[ref]

Repeated substitution gives the exponentially weighted gradient historyTo verify this identity, substitute its expression for $b_k$ into $b_{k+1}=\mu b_k+g_k$: multiplication by $\mu$ raises every old exponent by one, and the added $g_k$ is precisely the new exponent-zero term. The formula holds at $k=0$ because its sum is empty, completing the mathematical induction.

For an original scalar example with $b_0=0$ and a constant supplied gradient $g_j=1$, the geometric series yields $b_k=(1-{\mu }^k)/(1-\mu )$. At $\mu =0.9$, $b_{10}\approx 6.5132$ and the limit is 10. The buffer is an unnormalized sum: multiplying it by $1-\mu$ changes its scale, and would require a corresponding learning rate adjustment to preserve these parameter updates.

The same memory explains why momentum can temporarily oppose the current gradient. If $b_k=10$, $\mu =0.9$ and the new scalar gradient is $g_k=-1$, then $b_{k+1}=8$. With a positive learning rate, the parameter still decreases, although the present negative gradient alone would call for an increase. Persistence can carry motion through rapidly alternating gradients, but this example also shows why the rule need not decrease the objective function at every step.

A momentum buffer can retain enough earlier history that it points against the current gradient. For the undampened rulethe parameter step is $\eta d_k$ for positive learning rate $\eta$. If $g_k^T{b}_{k+1}<0$, then the directional derivative in direction $d_k$ is $g_k^T{d}_k>0$. The actual update direction is therefore locally uphill for a differentiable objective function, even though gradient descent would use the opposite direction. This history-dependent recurrence is the non-Nesterov momentum convention documented for stochastic gradient descent with dampening and weight decay disabled.^[ref]

For an original scalar example, let $f(\theta )={\theta }^2/2$, current $\theta =1$, stored buffer $b=-3$, coefficient $\mu =1/2$ and $\eta =0.1$. The new gradient is 1, butThe objective increases from $f(1)=0.5$ to $f(1.05)=0.55125$. This is a fully specified finite step, not only a first-order prediction.

More generally,An old buffer sufficiently opposed to the new gradient makes this negative. At $\mu =0$, that mechanism disappears. The example explains why a momentum method need not decrease the objective at every step; it does not imply that retaining gradient history always harms convergence. A descent direction test concerns the direction at the current point, while a convergence claim concerns an entire sequence.

An exponentially weighted gradient history assigns geometrically decreasing coefficients to older supplied gradients. For constant $0\le \mu <1$ and momentum buffer recurrence $b_{k+1}=\mu b_k+g_k$,Each additional update multiplies an old contribution by $\mu$. The formula follows by substituting the expression for $b_k$ into the recurrence and adding the newest exponent-zero term. The same undampened recurrence, including its initialization convention, appears in the primary stochastic gradient descent documentation.^[ref]

For an original example, choose $b_0=0$, $\mu =1/2$ and gradients $g_0=2$, $g_1=-2$, $g_2=4$. Successive buffers are 2, −1, 3.5. Directly expanding the third buffer givesThe newest contribution has coefficient 1; the oldest has coefficient 1/4. The buffer is a weighted sum, not yet a weighted mean. At $b_0=0$, its coefficient sum is the geometric seriesHere $W_3=1.75$, so the normalized weighted mean is $b_3/W_3=2$, rather than the ordinary mean $4/3$.

The coefficients also give a concrete memory scale. In the infinite geometric sequence, the fraction assigned to the newest $m$ gradients is $1-{\mu }^m$. At $\mu =0.9$, 29 gradients account for at least 95% of that coefficient mass, because $0.9^{29}<0.05$ whereas $0.9^{28}>0.05$. This is a statement about coefficients, not a guarantee that gradients at different times are independent or that the optimizer averages a stationary process.