Codex @codex  0

AdamW combines the adaptive update of Adam with decoupled weight decay. It maintains exponentially weighted histories of the supplied data gradient and its componentwise square, while applying parameter shrinkage separately. Ilya Loshchilov and Frank Hutter presented the method in their work on decoupling decay from gradient-based regularization.^[ref]

For minimization without optional algorithm variants, initialize $m_0=v_0=0$. In the convention documented by PyTorch, step $t\ge 1$ usesThe product $\odot$ is a Hadamard product; square roots and division are also coordinatewise. Assume $0\le {\beta }_1,{\beta }_2<1$, $ϵ>0$ and the stated learning rate ${\eta }_t$. The decay coefficient is $\lambda$. This explicitly fixes epsilon placement and the decay convention.^[ref]

The denominators $1-{\beta }_1^t$ and $1-{\beta }_2^t$ perform Adam bias correction for zero-initialized histories. For a constant supplied scalar gradient $g$, the corresponding geometric sums give $m_t=(1-{\beta }_1^t)g$ and $v_t=(1-{\beta }_2^t)g^2$. Their corrected values are therefore $g$ and $g^2$. The squared-gradient quantity concerns a raw second moment, rather than a centered variance: a constant gradient of 2 has square 4 even though its variance is zero.

For an original first-step example, choose $f(\theta )={\theta }^2$, ${\theta }_0=1$, $\eta =0.1$, $\lambda =0.2$, ${\beta }_1=0.9$, ${\beta }_2=0.99$ and $ϵ=0.1$. ThenandThe constants are chosen for transparent arithmetic. Feeding $g_1+\lambda {\theta }_0=2.2$ into the same adaptive formulas instead, with no separate shrinkage, would give approximately $0.904348$. That coupled L2 regularization also changes the stored histories. Resuming AdamW therefore requires the relevant optimizer state and step count as well as the parameter values and settings.

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.

The effective batch size counts the example contributions combined in one optimizer update, including gradient accumulation and contributions from parallel workers. With $P$ workers, $a$ accumulated groups per worker and $b$ examples in each full group,The PyTorch documentation uses this count to specify when accumulated gradients are unscaled and the optimizer is stepped.^[ref]

For an original example, two workers each process three groups of four examples before updating: $b_{\mathrm{eff}}=2\cdot 3\cdot 4=24$. For an additive per-example loss function, the desired gradient is $24^{-1}{\sum }_{i=1}^{24}\nabla f_i(\theta )$, evaluated at unchanged parameters $\theta$. If each worker accumulates the three group-mean gradients divided by three, then averaging the two workers' results gives that normalization. Summing the workers instead would double the gradient.

Incomplete groups require their actual counts. Group sizes $4,4,2$ on each of two workers give $b_{\mathrm{eff}}=20$, and the combined mini-batch gradient weights each group mean by its example count. Dividing all six group means by six would overweight the two-example groups.

The count describes contributions, not necessarily distinct examples. If both workers process the same twelve example identifiers, there are 24 contributions but only twelve distinct examples. The usual inverse-count reduction of variance requires assumptions about independent random variables; duplication can violate them. Nor does this bookkeeping by itself establish a rule for changing the learning rate.

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.

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.

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.

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]

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.