22 · Content revision 7

Keep decay out of the gradient moments

Build AdamW from named parameter gradients, bias-corrected moments, and a separate weight-decay path, then commit every checked update together.

Predict two subtractions

Chapter 21 ends with tensor-shaped gradients of the token-mean loss. Before an optimizer step, associate each gradient with the stable name of the parameter with respect to which it was computed. AdamW uses each named gradient to compute an updated value for the matching parameter. It prepares the complete named update before writing any value, then commits every updated value into the same live parameter leaf.

In this worked update, $\theta_0$ is the current value of the decay-group parameter decoder.output.weight, and $g_1$ is the accumulated token-mean loss gradient with respect to that same parameter:

\theta_0=[1,-2], \qquad g_1=[0.2,-0.4].

AdamW stores this parameter’s moment vectors under decoder.output.weight. The stable name, not the parameter’s position in the parameter list, identifies its moment history.

Use $\eta=0.1$ , $\beta_1=\beta_2=0.5$ , $\varepsilon=0.1$ , and $\lambda=0.1$ . This direct worked step uses gradient scale $\alpha_1=1$ , so the effective gradient is exactly $\widetilde g_1=g_1$ . Predict the adaptive delta and the decay delta separately.

The zero-started moments are $m_1=[0.1,-0.2]$ and $v_1=[0.02,0.08]$ . Both correction denominators equal $0.5$ , giving $\hat m_1=[0.2,-0.4]$ and $\hat v_1=[0.04,0.16]$ . Therefore:

\eta\frac{\hat m_1}{\sqrt{\hat v_1}+\varepsilon} \approx[0.066667,-0.08], \qquad \eta\lambda\theta_0=[0.01,-0.02].

Subtract both from the old value to predict $\theta_1\approx[0.923333,-1.9]$ . The signs matter: decaying the negative coordinate subtracts a negative delta, moving it toward zero.

Correct moments, then bypass them with decay

The optimizer receives the raw accumulated gradient $g_t$ stored on the named parameter and a validated scalar $0\leq\alpha_t\leq1$ . It forms the effective gradient

\widetilde g_t=\alpha_t g_t.

The scalar is shared by every coordinate of the complete parameter set. It is not a second learning rate: it changes only the gradient supplied to the moment recurrences. The optimizer reads $g_t$ without changing the stored gradient. After a successful transaction, the parameter has the new value but keeps the same leaf identity and the same accumulated gradient. AdamW uses the gradient in its calculation and retains it; the Chapter 23 training loop calls zero_grad() explicitly after the update and before the next backward pass. A scale of $1$ preserves every gradient bit used by the calculation, as in this chapter’s worked step; Chapter 33 derives a scale below $1$ when the complete gradient norm exceeds its ceiling.

For a non-unit example, let the raw gradient be $[0.8,-0.4]$ and let $\alpha_t=0.25$ . Then

\widetilde g_t=0.25[0.8,-0.4]=[0.2,-0.1].

The first-moment calculation uses $[0.2,-0.1]$ , and the second-moment calculation uses its coordinate squares $[0.04,0.01]$ . The separate decay term remains $\eta\lambda\theta_{t-1}$ .

Update the elementwise first and second raw moments with that effective gradient:

m_t=\beta_1m_{t-1}+(1-\beta_1)\widetilde g_t, \qquad v_t=\beta_2v_{t-1}+(1-\beta_2)\widetilde g_t^2.

The first moment carries recent gradient direction forward, which is the momentum intuition: consistent directions reinforce one another while a sudden reversal is softened. The second raw moment tracks recent squared magnitude. Dividing the corrected first moment by its root-mean-square scale adapts each coordinate according to the ratio between recent direction and magnitude; it does not simply make every larger gradient produce a smaller absolute step.

Their zero initialization suppresses early values. Correct that missing mass:

\hat m_t=\frac{m_t}{1-\beta_1^t}, \qquad \hat v_t=\frac{v_t}{1-\beta_2^t}.

Now apply the chapter’s exact shared formula:

\hat m_t=\frac{m_t}{1-\beta_1^t},\quad \hat v_t=\frac{v_t}{1-\beta_2^t},\quad \theta_t=(1-\eta\lambda)\theta_{t-1}-\eta\frac{\hat m_t}{\sqrt{\hat v_t}+\varepsilon}

Equivalently, subtract the adaptive delta and $\eta\lambda\theta_{t-1}$ as two explicit terms. Decoupling means the second term bypasses $g_t$ , $\widetilde g_t$ , $m_t$ , and $v_t$ . Changing $\alpha_t$ therefore does not scale decay. Decoupling does not mean the final updated value ignores the adaptive branch.

In this formula, $\lambda$ is the effective coefficient for the current named parameter: it equals the configured decay for a decay-group parameter and $0$ for a no-decay parameter.

The group map is a configurable policy and an explicit partition of the stable parameter-name set: its decay and no-decay sets are disjoint, and their union contains every name. In this course example, the policy assigns decoder.output.weight to decay and decoder.norm.scale to no-decay; that assignment is not implied by the AdamW equation. The normalization scale’s effective $\lambda$ is therefore $0$ , avoiding a separate decay term that would directly pull the learned affine scale toward zero.

Keep step state and parameter state distinct

$\theta_{t-1}$ is one named parameter tensor before step $t$ .
$g_t$ is its accumulated token-mean loss gradient, which remains stored on the live leaf until the caller clears it.
$\alpha_t$ is the validated scale shared by every gradient coordinate; $0\leq\alpha_t\leq1$ .
$\widetilde g_t=\alpha_tg_t$ is the effective gradient used by both moment recurrences.
$m_t$ and $v_t$ are name-keyed moving estimates with the same shape.
$\beta_1$ and $\beta_2$ retain past first and second raw moments, respectively.
$\hat m_t$ and $\hat v_t$ correct their early-step zero bias.
$\eta>0$ is the learning rate; $\lambda\geq0$ is decay.
$\varepsilon>0$ stabilizes the adaptive denominator after its square root.
$\theta_t$ is the updated value written into that same trainable leaf only after the whole set succeeds.

On the fresh probe, $m_0=v_0=0$ and $g_1=0$ , so both new moments and the adaptive delta are zero. A decay-group parameter still shrinks through $\eta\lambda\theta_0$ , while a no-decay parameter remains unchanged. After an earlier nonzero gradient, $g_t=0$ only removes the new contribution: stored moments decay but can still produce a nonzero adaptive update. The group controls decay, not moment memory.

From one word gradient to AdamW-trained decoders

Bengio et al., A Neural Probabilistic Language Model provide the earlier language-model evidence. Bengio et al.’s neural language model performs a direct stochastic parameter update after presenting one training word and its context. At each update, one scalar learning rate scales the current example’s gradient; there is no per-coordinate memory of earlier gradients. In loss-minimization notation, the core move is simply:

\theta\leftarrow\theta-\eta g.

Relative to this direct SGD rule, momentum adds a decaying memory of past directions. Adam later keeps exponential first and second raw gradient moments and corrects their early zero-initialization bias; if an $L_2$ term is coupled into its gradient, that parameter-proportional term enters both moving estimates. AdamW instead moves the shrinkage term outside the gradient entering those adaptive moments.

In compact loss-minimization notation, momentum uses $u_t=\mu u_{t-1}+g_t$ and $\theta_t=\theta_{t-1}-\eta u_t$ . Coupled $L_2$ feeds $g_t+\lambda\theta_{t-1}$ into Adam’s moments. AdamW excludes the parameter-proportional term from the gradient input to those moments and applies shrinkage separately. In this chapter’s direct $\alpha_t=1$ example, that input is $g_t$ ; after Chapter 33 applies clipping, it is $\widetilde g_t$ . Here $u_t$ is the retained update velocity and $0\leq\mu\lt1$ is its retention rate.

Kingma and Ba, Adam: A Method for Stochastic Optimization support the adaptive stage: Kingma and Ba’s Adam keeps exponential first and second raw gradient moments and corrects their early zero-initialization bias. They divide by $1-\beta_1^t$ and $1-\beta_2^t$ because zero initialization biases early estimates toward zero. Adam is a general optimizer, not an LLM architecture.

Loshchilov and Hutter, Decoupled Weight Decay Regularization show that an $L_2$ penalty mixed into the loss gradient is not equivalent to weight decay for an adaptive optimizer. Loshchilov and Hutter’s AdamW then moves parameter-proportional decay outside the gradient entering those adaptive moments.

LLaMA documents AdamW in pretraining decoder language models from $7$ B to $65$ B parameters. The course’s optimizer writes each checked update into the existing named parameter leaf; Chapter 33 supplies a scheduled learning rate, one validated global clipping factor, and explicit gradient clearing, while mixed precision and distributed optimizer state remain outside this bounded implementation.

That progression—from direct next-word gradients to adaptive, decoupled updates used in decoder pretraining—is why AdamW belongs on the road to modern LLMs.

Touvron et al., LLaMA: Open and Efficient Foundation Language Models document AdamW while pretraining decoder language models from $7$ B to $65$ B parameters. Their recipe uses $\beta_1=0.9$ , $\beta_2=0.95$ , weight decay $0.1$ , gradient clipping, warmup, and cosine learning-rate decay. This chapter implements the optimizer mechanism, not that full large-scale recipe.

The historical example runs two-step plain SGD, momentum, Adam with coupled $L_2$ , and AdamW for the same scalar loss-gradient sequence, then supplies the fixed unequal-curvature trajectory. Rust is only the executable medium for this language-independent comparison:

Inside the trajectory loop, one block calls value() to borrow the current AdamW parameter value without copying it and computes both coordinates of the quadratic gradient. Leaving that block ends the read guard. If the read guard remained active, step() could not obtain mutable access to the same stored value. After the step, the loop explicitly calls zero_grad() so the next backward seed does not accumulate onto the gradient just used. These scopes change ownership timing, not the optimizer mathematics or the recorded trajectory.

Compare SGD, momentum, coupled-penalty Adam, and decoupled AdamW rust/demos/ch22-adamw/src/lib.rs#historical-optimizer-road

/// Four optimizer endpoints for one loss-gradient sequence.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct HistoricalUpdates {
    pub sgd: f64,
    pub momentum: f64,
    pub adam_l2: f64,
    pub adamw: f64,
}

pub fn historical_updates(parameter_value: f64, gradients: [f64; 2]) -> HistoricalUpdates {
    let mut sgd = parameter_value;
    for gradient in gradients {
        sgd -= LEARNING_RATE * gradient;
    }

    let mut momentum = parameter_value;
    let mut velocity = 0.0;
    for gradient in gradients {
        velocity = MOMENTUM_RATE * velocity + gradient;
        momentum -= LEARNING_RATE * velocity;
    }

    let adam_l2 = two_step_adaptive_update(parameter_value, gradients, true);
    let adamw = two_step_adaptive_update(parameter_value, gradients, false);
    HistoricalUpdates {
        sgd,
        momentum,
        adam_l2,
        adamw,
    }
}

fn two_step_adaptive_update(
    parameter_value: f64,
    gradients: [f64; 2],
    couple_l2_into_gradient: bool,
) -> f64 {
    let decoupled_decay = if couple_l2_into_gradient {
        0.0
    } else {
        WEIGHT_DECAY
    };
    let config = AdamWConfig::new(LEARNING_RATE, BETA1, BETA2, EPSILON, decoupled_decay)
        .expect("historical fixture configuration is valid");
    let parameters = vec![parameter("history.weight", &[1], &[parameter_value])];
    let mut optimizer = AdamW::new(config);
    for gradient in gradients {
        let current = parameters[0].tensor().value().as_slice()[0];
        let optimizer_gradient = if couple_l2_into_gradient {
            gradient + WEIGHT_DECAY * current
        } else {
            gradient
        };
        seed_gradient(&parameters[0], &[optimizer_gradient]);
        optimizer
            .step(&parameters)
            .expect("historical fixture update is finite");
        parameters[0]
            .tensor()
            .zero_grad()
            .expect("the historical fixture clears each used gradient");
    }
    parameters[0].tensor().value().as_slice()[0]
}

/// One exact point on the same anisotropic objective for both update rules.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct TrajectoryPoint {
    pub step: usize,
    pub sgd: [f64; 2],
    pub adamw: [f64; 2],
}

/// Compares four updates on `q(x,y)=(x^2+4y^2)/2` from the same start.
pub fn anisotropic_trajectory() -> Vec<TrajectoryPoint> {
    const CURVATURE: [f64; 2] = [1.0, 4.0];
    const STEPS: usize = 4;

    let mut sgd = [1.0, 1.0];
    let adamw_parameter = parameter("trajectory.weight", &[2], &[1.0, 1.0]);
    let mut optimizer = AdamW::new(fixture_config());
    let mut points = vec![TrajectoryPoint {
        step: 0,
        sgd,
        adamw: [1.0, 1.0],
    }];

    for step in 1..=STEPS {
        let sgd_gradient = [CURVATURE[0] * sgd[0], CURVATURE[1] * sgd[1]];
        for axis in 0..2 {
            sgd[axis] -= LEARNING_RATE * sgd_gradient[axis];
        }

        let adamw_gradient = {
            let value = adamw_parameter.tensor().value();
            [
                CURVATURE[0] * value.as_slice()[0],
                CURVATURE[1] * value.as_slice()[1],
            ]
        };
        seed_gradient(&adamw_parameter, &adamw_gradient);
        optimizer
            .step(std::slice::from_ref(&adamw_parameter))
            .expect("bounded trajectory stays finite");
        adamw_parameter
            .tensor()
            .zero_grad()
            .expect("the trajectory clears each used gradient");
        let next = adamw_parameter.tensor().value();
        points.push(TrajectoryPoint {
            step,
            sgd,
            adamw: [next.as_slice()[0], next.as_slice()[1]],
        });
    }
    points
}

Prepare every named value before the live commit

Configuration rejects invalid scalar domains before any optimizer state exists:

Validate learning rate, moment rates, stabilizer, and decay rust/crates/llm-from-scratch/src/training/adamw.rs#adamw-configuration

/// The five scalar controls used by one fixed-learning-rate AdamW optimizer.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct AdamWConfig {
    learning_rate: f64,
    beta1: f64,
    beta2: f64,
    epsilon: f64,
    weight_decay: f64,
}

impl AdamWConfig {
    /// Validates the scalar controls for one AdamW configuration.
    pub fn new(
        learning_rate: f64,
        beta1: f64,
        beta2: f64,
        epsilon: f64,
        weight_decay: f64,
    ) -> Result<Self, AdamWError> {
        if !learning_rate.is_finite() || learning_rate <= 0.0 {
            return Err(AdamWError::InvalidLearningRate {
                value: learning_rate,
            });
        }
        if !beta1.is_finite() || !(0.0..1.0).contains(&beta1) {
            return Err(AdamWError::InvalidBeta1 { value: beta1 });
        }
        if !beta2.is_finite() || !(0.0..1.0).contains(&beta2) {
            return Err(AdamWError::InvalidBeta2 { value: beta2 });
        }
        if !epsilon.is_finite() || epsilon <= 0.0 {
            return Err(AdamWError::InvalidEpsilon { value: epsilon });
        }
        if !weight_decay.is_finite() || weight_decay < 0.0 {
            return Err(AdamWError::InvalidWeightDecay {
                value: weight_decay,
            });
        }
        Ok(Self {
            learning_rate,
            beta1,
            beta2,
            epsilon,
            weight_decay,
        })
    }

    pub const fn learning_rate(self) -> f64 {
        self.learning_rate
    }

    pub const fn beta1(self) -> f64 {
        self.beta1
    }

    pub const fn beta2(self) -> f64 {
        self.beta2
    }

    pub const fn epsilon(self) -> f64 {
        self.epsilon
    }

    pub const fn weight_decay(self) -> f64 {
        self.weight_decay
    }
}

An explicit group map assigns every name once: decoder.output.weight receives decay while decoder.norm.scale is excluded. An empty overall assignment, duplicate names within a group, overlaps, omissions, and extra assignments are rejected before commit; either one of the two groups may validly be empty:

Partition stable names into explicit decay and no-decay groups rust/crates/llm-from-scratch/src/training/adamw.rs#adamw-parameter-groups

/// The two explicit parameter groups used by the course's decay policy.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum AdamWGroup {
    Decay,
    NoDecay,
}

impl fmt::Display for AdamWGroup {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Decay => formatter.write_str("decay"),
            Self::NoDecay => formatter.write_str("no-decay"),
        }
    }
}

/// Exact stable-name assignments for decayed and decay-excluded parameters.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AdamWParameterGroups {
    decay: BTreeSet<String>,
    no_decay: BTreeSet<String>,
}

impl AdamWParameterGroups {
    pub fn new<D, N, DS, NS>(decay: D, no_decay: N) -> Result<Self, AdamWError>
    where
        D: IntoIterator<Item = DS>,
        N: IntoIterator<Item = NS>,
        DS: Into<String>,
        NS: Into<String>,
    {
        let decay = collect_group(decay, AdamWGroup::Decay)?;
        let no_decay = collect_group(no_decay, AdamWGroup::NoDecay)?;
        if decay.is_empty() && no_decay.is_empty() {
            return Err(AdamWError::EmptyParameterGroups);
        }
        if let Some(name) = decay.intersection(&no_decay).next() {
            return Err(AdamWError::ParameterInMultipleGroups {
                name: name.to_owned(),
            });
        }
        Ok(Self { decay, no_decay })
    }

    pub fn decayed_names(&self) -> impl ExactSizeIterator<Item = &str> {
        self.decay.iter().map(String::as_str)
    }

    pub fn excluded_names(&self) -> impl ExactSizeIterator<Item = &str> {
        self.no_decay.iter().map(String::as_str)
    }

    fn parameter_names(&self) -> Vec<String> {
        self.decay.union(&self.no_decay).cloned().collect()
    }

    fn decays(&self, name: &str) -> bool {
        self.decay.contains(name)
    }
}

fn collect_group<I, S>(names: I, group: AdamWGroup) -> Result<BTreeSet<String>, AdamWError>
where
    I: IntoIterator<Item = S>,
    S: Into<String>,
{
    let mut collected = BTreeSet::new();
    for name in names {
        let name = name.into();
        if name.is_empty() {
            return Err(AdamWError::EmptyGroupedParameterName { group });
        }
        if !collected.insert(name.clone()) {
            return Err(AdamWError::DuplicateGroupedParameter { group, name });
        }
    }
    Ok(collected)
}

Moment state is shaped like each parameter and keyed by its stable name:

Keep each pair of moment vectors attached to one stable name rust/crates/llm-from-scratch/src/training/adamw.rs#adamw-moment-state

/// Name-keyed optimizer memory for one parameter tensor.
#[derive(Clone, Debug, PartialEq)]
pub struct AdamWMomentState {
    shape: Vec<usize>,
    first: Vec<f64>,
    second: Vec<f64>,
}

impl AdamWMomentState {
    fn zeros(shape: &[usize], elements: usize) -> Self {
        Self {
            shape: shape.to_vec(),
            first: vec![0.0; elements],
            second: vec![0.0; elements],
        }
    }

    pub fn shape(&self) -> &[usize] {
        &self.shape
    }

    pub fn first_moment(&self) -> &[f64] {
        &self.first
    }

    pub fn second_moment(&self) -> &[f64] {
        &self.second
    }
}

Ordinary methods step and step_with_learning_rate execute the complete atomic update and return only the number of the step that was committed. The additional ordinary method step_with_learning_rate_and_gradient_scale receives the scheduled rate and $\alpha_t$ while still returning only the step number. The four earlier ordinary and traced entry points behave exactly as $\alpha_t=1$ . They do not construct an AdamWStep unless the caller uses step_with_trace or step_with_learning_rate_and_trace.

Every entry point uses the same internal preparation-and-commit operation and the same elementwise AdamW calculation. Tracing records values produced by that calculation; it does not calculate the update a second time. In a trace, gradient() is the effective $\widetilde g_t$ actually used by both moments, not a value recomputed by the observer.

Separate ordinary AdamW execution from explicitly requested trace evidence rust/crates/llm-from-scratch/src/training/adamw.rs#adamw-execution-and-trace-api

    /// Reads the accumulated gradients and atomically updates every live leaf.
    ///
    /// All arithmetic, tensor construction, and optimizer-state changes are
    /// prepared first. An error leaves both the supplied parameters and this
    /// optimizer bit-identical. A successful commit preserves every parameter
    /// node and leaves its accumulated gradient for the caller to clear.
    /// The result is only the committed step number; use `step_with_trace` when
    /// the elementwise update vectors are needed for inspection.
    pub fn step(&mut self, parameters: &[NamedParameter]) -> Result<u64, AdamWError> {
        self.step_with_config(parameters, self.config, 1.0, NoAdamWTrace)
    }

    /// Applies the same transaction while recording every elementwise update.
    pub fn step_with_trace(
        &mut self,
        parameters: &[NamedParameter],
    ) -> Result<AdamWStep, AdamWError> {
        let observer = RecordAdamWTrace::with_capacity(parameters.len());
        self.step_with_config(parameters, self.config, 1.0, observer)
    }

    /// Applies one validated scheduled learning rate without resetting moments.
    ///
    /// The override belongs only to this update; `config()` keeps the optimizer's
    /// base rate. An invalid rate or any later preparation error leaves the
    /// parameters, moments, powers, and step counter unchanged.
    pub fn step_with_learning_rate(
        &mut self,
        parameters: &[NamedParameter],
        learning_rate: f64,
    ) -> Result<u64, AdamWError> {
        let step_config = self.config.with_learning_rate(learning_rate)?;
        self.step_with_config(parameters, step_config, 1.0, NoAdamWTrace)
    }

    /// Applies one scheduled rate and one validated global gradient scale.
    ///
    /// The scale multiplies only the gradient used by Adam's moments. The
    /// decoupled weight-decay branch continues to use the unscaled parameter.
    pub fn step_with_learning_rate_and_gradient_scale(
        &mut self,
        parameters: &[NamedParameter],
        learning_rate: f64,
        gradient_scale: f64,
    ) -> Result<u64, AdamWError> {
        let step_config = self.config.with_learning_rate(learning_rate)?;
        self.step_with_config(parameters, step_config, gradient_scale, NoAdamWTrace)
    }

    /// Applies a scheduled learning rate and records the complete update trace.
    pub fn step_with_learning_rate_and_trace(
        &mut self,
        parameters: &[NamedParameter],
        learning_rate: f64,
    ) -> Result<AdamWStep, AdamWError> {
        let step_config = self.config.with_learning_rate(learning_rate)?;
        let observer = RecordAdamWTrace::with_capacity(parameters.len());
        self.step_with_config(parameters, step_config, 1.0, observer)
    }

Only an explicitly requested trace retains the exact input, moment, delta, and updated vectors used as evidence in this chapter:

Define the evidence retained only for a traced AdamW step rust/crates/llm-from-scratch/src/training/adamw.rs#adamw-state-and-evidence

/// Exact elementwise evidence prepared for one named parameter in a step.
#[derive(Clone, Debug, PartialEq)]
pub struct AdamWParameterUpdate {
    name: String,
    shape: Vec<usize>,
    before: Vec<f64>,
    gradient: Vec<f64>,
    decay_applied: bool,
    effective_weight_decay: f64,
    first_moment: Vec<f64>,
    second_moment: Vec<f64>,
    corrected_first_moment: Vec<f64>,
    corrected_second_moment: Vec<f64>,
    adaptive_direction: Vec<f64>,
    adaptive_delta: Vec<f64>,
    decay_delta: Vec<f64>,
    after: Vec<f64>,
}

impl AdamWParameterUpdate {
    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn shape(&self) -> &[usize] {
        &self.shape
    }

    pub fn before(&self) -> &[f64] {
        &self.before
    }

    /// The effective gradient used by the moment calculation.
    pub fn gradient(&self) -> &[f64] {
        &self.gradient
    }

    pub const fn decay_applied(&self) -> bool {
        self.decay_applied
    }

    pub const fn effective_weight_decay(&self) -> f64 {
        self.effective_weight_decay
    }

    pub fn first_moment(&self) -> &[f64] {
        &self.first_moment
    }

    pub fn second_moment(&self) -> &[f64] {
        &self.second_moment
    }

    pub fn corrected_first_moment(&self) -> &[f64] {
        &self.corrected_first_moment
    }

    pub fn corrected_second_moment(&self) -> &[f64] {
        &self.corrected_second_moment
    }

    pub fn adaptive_direction(&self) -> &[f64] {
        &self.adaptive_direction
    }

    pub fn adaptive_delta(&self) -> &[f64] {
        &self.adaptive_delta
    }

    pub fn decay_delta(&self) -> &[f64] {
        &self.decay_delta
    }

    pub fn after(&self) -> &[f64] {
        &self.after
    }
}

/// The optional trace for one committed multi-parameter update.
#[derive(Clone, Debug, PartialEq)]
pub struct AdamWStep {
    step: u64,
    learning_rate: f64,
    first_correction: f64,
    second_correction: f64,
    updates: Vec<AdamWParameterUpdate>,
}

impl AdamWStep {
    pub const fn step(&self) -> u64 {
        self.step
    }

    pub const fn learning_rate(&self) -> f64 {
        self.learning_rate
    }

    pub const fn first_correction(&self) -> f64 {
        self.first_correction
    }

    pub const fn second_correction(&self) -> f64 {
        self.second_correction
    }

    pub fn updates(&self) -> &[AdamWParameterUpdate] {
        &self.updates
    }
}

Trace capture is optional, but the atomic update is not. The public scheduled method validates the learning rate while constructing its per-step configuration. It then enters the shared path, which validates $0\leq\alpha_t\leq1$ before inspecting any parameter. The shared path reads each raw gradient coordinate, forms the effective coordinate once, and prepares prospective name-keyed moments, the powers $\beta_1^t$ and $\beta_2^t$ , every updated parameter tensor, and each next value revision before changing live state. Weight decay still reads the value from before the step and is never multiplied by $\alpha_t$ .

After all calculations and tensor construction succeed, the implementation requests exclusive access to every live parameter value. It does not write the first parameter before checking the rest. If a caller still holds an active read borrow for any parameter value, mutable access to that node fails. AdamW then drops every write guard it already acquired, and neither parameters nor optimizer state changes. Once all guards exist, each prepared same-shape tensor is installed in its existing node and its value revision is advanced; those final assignments cannot return an error. The optimizer then installs the prepared moments, powers, and step count. Gradients are stored separately from parameter values, so the raw accumulated gradients remain available until the caller clears them.

Each parameter node owns a monotonically increasing parameter-value revision. One successful AdamW commit advances that node’s revision once; a failed commit does not advance it. The revision records in-place changes to the value stored in that node. It is neither the optimizer step number nor checkpoint data. Chapters 37 and 38 bind each KV cache to both parameter-node identity and the captured value revision, because cached keys and values projected before an update are stale even though AdamW preserves the node identity.

Forward operations also record the current revision of every operand edge. A retained graph built before a successful AdamW update therefore cannot run backward against the updated parameter values: the revision mismatch is rejected before any gradient or graph state changes. The caller must run a new forward pass to build saved context for the updated values.

The complete preparation, mutable-access preflight, and commit into the existing leaf nodes are visible in the source:

Prepare the complete named update, then commit into the live leaves rust/crates/llm-from-scratch/src/training/adamw.rs#transactional-adamw-step

    fn step_with_config<O: AdamWStepObserver>(
        &mut self,
        parameters: &[NamedParameter],
        step_config: AdamWConfig,
        gradient_scale: f64,
        mut observer: O,
    ) -> Result<O::Output, AdamWError> {
        validate_gradient_scale(gradient_scale)?;
        let actual_names = validate_parameter_names(parameters)?;
        if let Some(groups) = &self.groups {
            let expected_names = groups.parameter_names();
            if expected_names != actual_names {
                return Err(AdamWError::ParameterSetChanged {
                    expected: expected_names,
                    actual: actual_names,
                });
            }
        }
        let next_step = self.step.checked_add(1).ok_or(AdamWError::StepOverflow)?;
        let next_beta1_power = self.beta1_power * step_config.beta1;
        let next_beta2_power = self.beta2_power * step_config.beta2;
        let first_correction = 1.0 - next_beta1_power;
        let second_correction = 1.0 - next_beta2_power;

        let mut candidate_states = if self.step == 0 {
            let mut states = BTreeMap::new();
            for parameter in parameters.iter() {
                let value = parameter.tensor().value();
                states.insert(
                    parameter.name().to_owned(),
                    AdamWMomentState::zeros(value.shape(), value.len()),
                );
            }
            states
        } else {
            let expected_names = self.states.keys().cloned().collect::<Vec<_>>();
            if expected_names != actual_names {
                return Err(AdamWError::ParameterSetChanged {
                    expected: expected_names,
                    actual: actual_names,
                });
            }
            self.states.clone()
        };

        let mut candidate_values = Vec::with_capacity(parameters.len());
        let mut next_revisions = Vec::with_capacity(parameters.len());
        for parameter in parameters.iter() {
            let name = parameter.name();
            let before = parameter.tensor().value();
            let gradient =
                parameter
                    .tensor()
                    .gradient()
                    .ok_or_else(|| AdamWError::MissingGradient {
                        name: name.to_owned(),
                    })?;
            if before.shape() != gradient.shape() {
                return Err(AdamWError::GradientShapeMismatch {
                    name: name.to_owned(),
                    parameter: before.shape().to_vec(),
                    gradient: gradient.shape().to_vec(),
                });
            }

            let state = candidate_states
                .get_mut(name)
                .expect("validated parameter names have candidate state");
            if state.shape != before.shape() {
                return Err(AdamWError::ParameterShapeChanged {
                    name: name.to_owned(),
                    expected: state.shape.clone(),
                    actual: before.shape().to_vec(),
                });
            }

            let after = prepare_parameter_update(
                AdamWPreparation {
                    config: step_config,
                    decay_applied: self
                        .groups
                        .as_ref()
                        .is_none_or(|groups| groups.decays(name)),
                    first_correction,
                    second_correction,
                    name,
                    gradient_scale,
                },
                &before,
                &gradient,
                state,
                &mut observer,
            )?;
            let tensor = Tensor::from_vec(before.shape().to_vec(), after)?;
            let next_revision = parameter.tensor().next_value_revision().ok_or_else(|| {
                AdamWError::ParameterRevisionOverflow {
                    name: name.to_owned(),
                }
            })?;
            candidate_values.push(tensor);
            next_revisions.push(next_revision);
        }

        let observation = observer.finish(
            next_step,
            step_config.learning_rate,
            first_correction,
            second_correction,
        );

        let mut value_writes = Vec::with_capacity(parameters.len());
        for parameter in parameters {
            let write = parameter.tensor().try_value_write().map_err(|_| {
                AdamWError::ParameterValueBorrowed {
                    name: parameter.name().to_owned(),
                }
            })?;
            value_writes.push(write);
        }

        for ((write, value), revision) in value_writes
            .into_iter()
            .zip(candidate_values)
            .zip(next_revisions)
        {
            write.commit(value, revision);
        }
        self.step = next_step;
        self.beta1_power = next_beta1_power;
        self.beta2_power = next_beta2_power;
        self.states = candidate_states;

        Ok(observation)
    }

Typed errors cover configuration and gradient-scale domains, empty and duplicate sets, changed names and shapes, missing or mismatched gradients, counter or parameter-revision overflow, an active value borrow, non-finite stages, and delegated tensor construction:

Reject invalid updates without a partial live-value commit rust/crates/llm-from-scratch/src/training/adamw.rs#adamw-errors

/// The arithmetic stage that first produced a non-finite candidate value.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum AdamWArithmetic {
    FirstMoment,
    SquaredGradient,
    SecondMoment,
    CorrectedFirstMoment,
    CorrectedSecondMoment,
    AdaptiveDirection,
    AdaptiveDelta,
    DecayDelta,
    Parameter,
}

impl fmt::Display for AdamWArithmetic {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        let name = match self {
            Self::FirstMoment => "first moment",
            Self::SquaredGradient => "squared gradient",
            Self::SecondMoment => "second moment",
            Self::CorrectedFirstMoment => "bias-corrected first moment",
            Self::CorrectedSecondMoment => "bias-corrected second moment",
            Self::AdaptiveDirection => "adaptive direction",
            Self::AdaptiveDelta => "adaptive update",
            Self::DecayDelta => "decoupled decay update",
            Self::Parameter => "updated parameter",
        };
        formatter.write_str(name)
    }
}

/// A deterministic rejection that leaves parameters and optimizer state intact.
#[derive(Clone, Debug, PartialEq)]
pub enum AdamWError {
    InvalidLearningRate {
        value: f64,
    },
    InvalidGradientScale {
        value: f64,
    },
    InvalidBeta1 {
        value: f64,
    },
    InvalidBeta2 {
        value: f64,
    },
    InvalidEpsilon {
        value: f64,
    },
    InvalidWeightDecay {
        value: f64,
    },
    EmptyParameterGroups,
    EmptyGroupedParameterName {
        group: AdamWGroup,
    },
    DuplicateGroupedParameter {
        group: AdamWGroup,
        name: String,
    },
    ParameterInMultipleGroups {
        name: String,
    },
    EmptyParameterSet,
    DuplicateParameterName {
        name: String,
        first: usize,
        repeated: usize,
    },
    ParameterSetChanged {
        expected: Vec<String>,
        actual: Vec<String>,
    },
    ParameterShapeChanged {
        name: String,
        expected: Vec<usize>,
        actual: Vec<usize>,
    },
    MissingGradient {
        name: String,
    },
    GradientShapeMismatch {
        name: String,
        parameter: Vec<usize>,
        gradient: Vec<usize>,
    },
    ParameterRevisionOverflow {
        name: String,
    },
    ParameterValueBorrowed {
        name: String,
    },
    StepOverflow,
    NonFiniteArithmetic {
        name: String,
        index: usize,
        stage: AdamWArithmetic,
        value: f64,
    },
    Tensor(TensorError),
}

impl fmt::Display for AdamWError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::InvalidLearningRate { value } => write!(
                formatter,
                "learning rate must be finite and greater than zero, got {value}"
            ),
            Self::InvalidGradientScale { value } => write!(
                formatter,
                "gradient scale must be finite in the closed interval [0,1], got {value}"
            ),
            Self::InvalidBeta1 { value } => write!(
                formatter,
                "beta1 must be finite in the half-open interval [0,1), got {value}"
            ),
            Self::InvalidBeta2 { value } => write!(
                formatter,
                "beta2 must be finite in the half-open interval [0,1), got {value}"
            ),
            Self::InvalidEpsilon { value } => write!(
                formatter,
                "epsilon must be finite and greater than zero, got {value}"
            ),
            Self::InvalidWeightDecay { value } => write!(
                formatter,
                "weight decay must be finite and non-negative, got {value}"
            ),
            Self::EmptyParameterGroups => formatter
                .write_str("explicit AdamW parameter groups must assign at least one stable name"),
            Self::EmptyGroupedParameterName { group } => {
                write!(
                    formatter,
                    "the {group} group contains an empty parameter name"
                )
            }
            Self::DuplicateGroupedParameter { group, name } => write!(
                formatter,
                "parameter name {name:?} repeats inside the {group} group"
            ),
            Self::ParameterInMultipleGroups { name } => write!(
                formatter,
                "parameter name {name:?} appears in both the decay and no-decay groups"
            ),
            Self::EmptyParameterSet => {
                formatter.write_str("AdamW needs at least one named parameter")
            }
            Self::DuplicateParameterName {
                name,
                first,
                repeated,
            } => write!(
                formatter,
                "parameter name {name:?} first appears at index {first} and repeats at index {repeated}"
            ),
            Self::ParameterSetChanged { expected, actual } => write!(
                formatter,
                "parameter-name set changed from {expected:?} to {actual:?}"
            ),
            Self::ParameterShapeChanged {
                name,
                expected,
                actual,
            } => write!(
                formatter,
                "parameter {name:?} changed shape from {expected:?} to {actual:?}"
            ),
            Self::MissingGradient { name } => {
                write!(formatter, "parameter {name:?} has no stored gradient")
            }
            Self::GradientShapeMismatch {
                name,
                parameter,
                gradient,
            } => write!(
                formatter,
                "parameter {name:?} has shape {parameter:?}, but its gradient has shape {gradient:?}"
            ),
            Self::ParameterRevisionOverflow { name } => {
                write!(
                    formatter,
                    "parameter {name:?} value revision overflowed u64"
                )
            }
            Self::ParameterValueBorrowed { name } => write!(
                formatter,
                "parameter {name:?} cannot be updated while its value is borrowed"
            ),
            Self::StepOverflow => formatter.write_str("AdamW step counter overflowed u64"),
            Self::NonFiniteArithmetic {
                name,
                index,
                stage,
                value,
            } => write!(
                formatter,
                "parameter {name:?} produced non-finite {stage} at flat index {index}: {value}"
            ),
            Self::Tensor(error) => error.fmt(formatter),
        }
    }
}

impl Error for AdamWError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::Tensor(error) => Some(error),
            _ => None,
        }
    }
}

impl From<TensorError> for AdamWError {
    fn from(error: TensorError) -> Self {
        Self::Tensor(error)
    }
}

The fixed worked example explicitly requests a trace for the named update and another trace for the probe with initially zero moments because this chapter inspects the exact vectors of both updates. The changed-set rollback probe needs only its error; the historical comparison and trajectory need only resulting parameter values. These calls therefore use the ordinary methods:

Build exact first-step and transaction evidence rust/demos/ch22-adamw/src/lib.rs#chapter-adamw-fixture

pub fn learner_evidence() -> LearnerEvidence {
    let config = fixture_config();
    let parameters = vec![
        parameter("decoder.output.weight", &[2], &[1.0, -2.0]),
        parameter("decoder.norm.scale", &[1], &[0.5]),
    ];
    seed_gradient(&parameters[0], &[0.2, -0.4]);
    let original_leaves = parameters
        .iter()
        .map(|parameter| parameter.tensor().clone())
        .collect::<Vec<_>>();
    let mut optimizer = AdamW::with_parameter_groups(config, fixture_groups());
    let step = optimizer
        .step_with_trace(&parameters)
        .expect("the complete named set updates atomically");

    let raw_gradients_retained =
        parameters
            .iter()
            .zip(step.updates())
            .all(|(parameter, update)| {
                parameter.tensor().gradient().is_some_and(|gradient| {
                    gradient
                        .as_slice()
                        .iter()
                        .zip(update.gradient())
                        .all(|(actual, expected)| actual.to_bits() == expected.to_bits())
                })
            });
    let parameter_nodes_preserved = parameters
        .iter()
        .zip(&original_leaves)
        .all(|(parameter, original)| parameter.tensor().is_same_node(original));
    let state_names = optimizer.parameter_names().map(str::to_owned).collect();

    let zero_gradient_update = zero_gradient_moment_probe(config);
    let (rejected_error, rejection_rolled_back) = rejected_set_probe(&optimizer, &parameters);
    LearnerEvidence {
        config,
        step,
        state_names,
        raw_gradients_retained,
        parameter_nodes_preserved,
        zero_gradient_update,
        rejected_error,
        rejection_rolled_back,
    }
}

fn zero_gradient_moment_probe(config: AdamWConfig) -> AdamWParameterUpdate {
    let parameters = vec![parameter("probe.weight", &[1], &[3.0])];
    let groups = AdamWParameterGroups::new(["probe.weight"], std::iter::empty::<&str>())
        .expect("the probe weight belongs to the decay group");
    AdamW::with_parameter_groups(config, groups)
        .step_with_trace(&parameters)
        .expect("the live leaf starts with an exact zero gradient")
        .updates()[0]
        .clone()
}

fn rejected_set_probe(
    committed_optimizer: &AdamW,
    committed_parameters: &[NamedParameter],
) -> (AdamWError, bool) {
    let mut optimizer = committed_optimizer.clone();
    let optimizer_before = optimizer.clone();
    let mut parameters = committed_parameters.to_vec();
    parameters.push(parameter("unexpected.weight", &[1], &[1.0]));
    let parameters_before = parameters
        .iter()
        .map(|parameter| {
            (
                parameter.tensor().clone(),
                parameter.tensor().value().as_slice().to_vec(),
                parameter
                    .tensor()
                    .gradient()
                    .map(|gradient| gradient.as_slice().to_vec()),
            )
        })
        .collect::<Vec<_>>();
    let error = optimizer
        .step(&parameters)
        .expect_err("a changed named set must be rejected");
    let parameters_unchanged =
        parameters
            .iter()
            .zip(parameters_before)
            .all(|(parameter, (leaf, values, gradient))| {
                parameter.tensor().is_same_node(&leaf)
                    && parameter.tensor().value().as_slice() == values
                    && parameter
                        .tensor()
                        .gradient()
                        .map(|actual| actual.as_slice().to_vec())
                        == gradient
            });
    (error, parameters_unchanged && optimizer == optimizer_before)
}

The learner executable prints every exact vector and invariant:

Print the deterministic Chapter 22 learner report rust/demos/ch22-adamw/src/main.rs#learner-adamw-output

fn main() {
    print!("{}", ch22_adamw::learner_report());
}

Run cargo run --quiet --locked -p ch22-adamw to inspect the exact first step, name-keyed state, preserved parameter nodes, retained raw gradients, and changed-set rollback. The same implementation supports reordered parameter presentation, accumulated gradients, pure decay from zero-initialized moments, and retained-moment motion after a later zero gradient. Invalid domains, identities, shapes, arithmetic, or parameter sets return without changing a live value or optimizer state. Repeating the deterministic 200-step anisotropic-quadratic run is bit-identical and ends below objective $10^{-12}$ .

Watch decay bypass the moment lane

The fourteen-line trace contains two group-aware parameter records, five SGD and AdamW trajectory points, and one whole-set proof. The first figure follows the decay-group weight through every moment and delta. This fixed record uses $\alpha_1=1$ , so its input-gradient field represents both the raw $g_1$ and the effective $\widetilde g_1$ . The second keeps the no-decay scale, both trajectories, and the transaction proof together. Rust authors every displayed number:

Emit exact moment, delta, live-commit, and rollback evidence rust/demos/ch22-adamw/src/diagram_trace.rs#adamw-trace

pub fn diagram_trace() -> String {
    let evidence = learner_evidence();
    let mut lines = vec![format!(
        "META|step={}|learning_rate={:.6}|beta1={:.6}|beta2={:.6}|epsilon={:.6}|weight_decay={:.6}|first_correction={:.6}|second_correction={:.6}",
        evidence.step.step(),
        evidence.config.learning_rate(),
        evidence.config.beta1(),
        evidence.config.beta2(),
        evidence.config.epsilon(),
        evidence.config.weight_decay(),
        evidence.step.first_correction(),
        evidence.step.second_correction(),
    )];

    for (index, update) in evidence.step.updates().iter().enumerate() {
        lines.push(format!(
            "PARAM|index={index}|name={}|group={}|shape={:?}|before={}|gradient={}",
            update.name(),
            if update.decay_applied() {
                "decay"
            } else {
                "no_decay"
            },
            update.shape(),
            format_vector(update.before()),
            format_vector(update.gradient()),
        ));
        lines.push(format!(
            "MOMENT|index={index}|first={}|second={}|corrected_first={}|corrected_second={}",
            format_vector(update.first_moment()),
            format_vector(update.second_moment()),
            format_vector(update.corrected_first_moment()),
            format_vector(update.corrected_second_moment()),
        ));
        lines.push(format!(
            "DELTA|index={index}|adaptive={}|decay={}|after={}",
            format_vector(update.adaptive_delta()),
            format_vector(update.decay_delta()),
            format_vector(update.after()),
        ));
    }

    let trajectory = anisotropic_trajectory();
    lines.push(format!(
        "QUADRATIC|curvature=[1.000000, 4.000000]|steps={}",
        trajectory.len() - 1,
    ));
    lines.extend(trajectory.iter().map(|point| {
        format!(
            "POINT|step={}|sgd={}|adamw={}",
            point.step,
            format_vector(&point.sgd),
            format_vector(&point.adamw),
        )
    }));

    lines.push(format!(
        "PROOF|state_names={}|raw_gradients={}|parameter_nodes={}|zero_gradient_decay={:.6}|rollback={}|commit=atomic",
        evidence.state_names.join(","),
        if evidence.raw_gradients_retained {
            "retained"
        } else {
            "changed"
        },
        if evidence.parameter_nodes_preserved {
            "preserved"
        } else {
            "replaced"
        },
        evidence.zero_gradient_update.decay_delta()[0],
        if evidence.rejection_rolled_back {
            "unchanged"
        } else {
            "changed"
        },
    ));
    lines.join("\n") + "\n"
}

Follow AdamW's separate adaptive and decay paths

Follow one named decay-group parameter through gradient moments, bias correction, separate adaptive and parameter-proportional deltas, and an in-place value commit.

Committed step: $t=1$
Learning rate: $\eta=0.100000$
Moment rates: $\beta_1=0.500000$ $\beta_2=0.500000$
Denominator stabilizer: $\varepsilon=0.100000$
Weight decay: $\lambda=0.100000$

Update and correct gradient moments

Only the accumulated loss gradient enters these moving estimates. On the first step, bias correction recovers the current gradient and its square exactly.

$1-\beta_1^t=0.500000$ $1-\beta_2^t=0.500000$

Named parameter

`decoder.output.weight`

Parameter group: Apply decay decay
Shape: [2]

Value before $\theta_{t-1}$: $\left[1,-2\right]$
Accumulated gradient $g_t$: $\left[0.2,-0.4\right]$

Adaptive path

First moment $m_t$: $\left[0.1,-0.2\right]$
Second raw moment $v_t$: $\left[0.02,0.08\right]$
Corrected first moment $\hat m_t$: $\left[0.2,-0.4\right]$
Corrected second raw moment $\hat v_t$: $\left[0.04,0.16\right]$

Adaptive path

Adaptive delta

$\left[0.066667,-0.08\right]$

Value before

\theta_{t-1}

\left[1,-2\right]

Decay bypass · Apply decay

Decay delta

$\left[0.01,-0.02\right]$

−

Subtract both deltas

Updated value $\theta_t$

$\left[0.923333,-1.9\right]$

Keep adaptive and decay deltas separate

The adaptive branch uses corrected moments. The decay branch reads the old parameter directly and never enters either moment.

Adaptive path

\eta\hat m_t/(\sqrt{\hat v_t}+\varepsilon)

− Decay bypass

\eta\lambda\theta_{t-1}

Compare the no-decay group, optimizer paths, and atomic commit

Use the zero-gradient normalization scale to expose group-specific decay, then compare exact SGD and AdamW points and verify whole-set commit invariants.

Skip decay

Named parameter

`decoder.norm.scale`

Parameter group: Skip decay no_decay
Shape: [1]

Value before: $\theta_{t-1}=\left[0.5\right]$
Accumulated gradient: $g_t=\left[0\right]$
Corrected first moment: $\hat m_t=\left[0\right]$
Corrected second raw moment: $\hat v_t=\left[0\right]$
Adaptive delta: $\frac{\eta\hat m_t}{\sqrt{\hat v_t}+\varepsilon}=\left[0\right]$
Decay delta: $\eta\lambda\theta_{t-1}=\left[0\right]$
Updated value: $\theta_t=\left[0.5\right]$

Compare trajectories across unequal curvature

The same fixed quadratic supplies both paths. Every point is exact Rust evidence; no point is inferred from drawing geometry.

Quadratic curvature $q(x,y)=\frac12(x^2+4y^2)$ $\operatorname{diag}(H)=\left[1,4\right]$

SGD trajectory

Optimizer step $t=0$ $\left[1,1\right]$
Optimizer step $t=1$ $\left[0.9,0.6\right]$
Optimizer step $t=2$ $\left[0.81,0.36\right]$
Optimizer step $t=3$ $\left[0.729,0.216\right]$
Optimizer step $t=4$ $\left[0.6561,0.1296\right]$

AdamW trajectory

Optimizer step $t=0$ $\left[1,1\right]$
Optimizer step $t=1$ $\left[0.899091,0.892439\right]$
Optimizer step $t=2$ $\left[0.799889,0.786278\right]$
Optimizer step $t=3$ $\left[0.702629,0.681677\right]$
Optimizer step $t=4$ $\left[0.60758,0.578823\right]$

Commit every checked value into its live leaf

Rust prepares both named values first, acquires write access to the complete live set, and then commits both values into the existing leaves.

Check identity, retained gradients, and rollback

Moment state follows stable names rather than display order. A changed name set leaves every parameter value, gradient, node identity, and optimizer field unchanged.

State keys: decoder.norm.scaledecoder.output.weight
Accumulated gradient after AdamW: Unchanged
Parameter-node identity: Preserved
Zero-moment, zero-gradient decay probe: $\eta\lambda\theta=0.030000$
Changed-set attempt: Unchanged
Successful commit: Atomic

The trace records the exact names, shapes, vectors, correction denominators, and transaction evidence. In the focused update, solid adaptive lanes and the dashed decay branch repeat their distinction in words and borders. The compact group record then shows that the no-decay parameter receives a zero decay delta. Both optimizer paths use the same fixed points on $q(x,y)=\frac12(x^2+4y^2)$ . Follow the loss gradient through the moments, compare the separate parameter-proportional decay, and join both contributions at the final subtraction.

Predict before running the optimizer

Derive $m_1$ , $v_1$ , $\hat m_1$ , and $\hat v_1$ for the fixed vector.
Predict the two deltas and $\theta_1$ before rounding to six decimals.
Start a fresh optimizer for decay-group value $3$ with $g_1=0$ and predict the new value. Then explain why a later zero gradient can still have an adaptive contribution after nonzero history.
Explain why reversing two parameter records must not reverse their moments.
Predict what changes if the second candidate overflows after the first candidate has been calculated successfully.
After a successful commit, predict the live leaf’s stored gradient, node identity, parameter-value revision, and whether a retained pre-update graph can still run backward.
Order plain SGD, momentum, Adam with coupled $L_2$ , and AdamW, and state which terms enter optimizer memory.
Explain why the output weight receives decay while the normalization scale does not.
Compare the fixed SGD and AdamW paths on the unequal-curvature quadratic.

Check the predictions

$m_1=[0.1,-0.2]$ , $v_1=[0.02,0.08]$ , $\hat m_1=[0.2,-0.4]$ , and $\hat v_1=[0.04,0.16]$ .
The adaptive delta is approximately $[0.066667,-0.08]$ , decay is $[0.01,-0.02]$ , and $\theta_1\approx[0.923333,-1.9]$ .
With fresh zero moments, the adaptive delta stays zero, decay is $0.03$ , and the value becomes $2.97$ . After nonzero history, stored moments can keep the adaptive delta nonzero even when the current gradient is zero.
The state map uses stable names; slice order is presentation only.
Nothing commits: both parameter values, both gradients, both node identities, and all optimizer state remain unchanged.
The leaf still stores the accumulated gradient that AdamW used for the update, and its tape identity is unchanged. Its monotonically increasing parameter-value revision advances once. A retained graph whose operand edge recorded the preceding revision rejects backward, so the caller clears the gradient and runs a new forward pass before the next backward pass.
Plain SGD uses only the current gradient; momentum remembers a decaying direction; coupled $L_2$ feeds parameter-proportional shrinkage into Adam’s moments; AdamW keeps that shrinkage outside the moments. Bengio supplies the neural-language-model starting point, and Touvron et al. document AdamW in LLaMA pretraining.
This is the course’s configurable grouping policy, not a consequence of the AdamW equation. The policy assigns decoder.output.weight to decay, so AdamW subtracts the parameter-proportional term $\eta\lambda\theta_{t-1}$ from it. It assigns decoder.norm.scale to no-decay, so that parameter’s effective $\lambda$ is $0$ ; this avoids a decay term that directly pulls the learned normalization scale toward zero.
SGD reduces the high-curvature coordinate much faster; AdamW scales coordinates through their moment history and also applies a separate decay contribution.

Train the first cumulative neural language model next

The cumulative training path can now associate each accumulated token-mean gradient with its parameter’s stable name, preserve first and second moments under that name across steps, and atomically update the existing parameter leaves. AdamW retains every raw gradient after using it for the update. Chapter 23 explicitly clears those post-update gradients and uses this optimizer to train a fixed-context neural language model whose validation loss improves.

AdamW changes parameters; it does not yet define the model whose gradients it uses. The next chapter closes that loop while keeping data partitions, autograd, batching, parameter names, and optimizer state explicit.

Predict two subtractions

Correct moments, then bypass them with decay

Keep step state and parameter state distinct

From one word gradient to AdamW-trained decoders

Prepare every named value before the live commit

Watch decay bypass the moment lane

Follow AdamW's separate adaptive and decay paths

1 Update and correct gradient moments

2 Keep adaptive and decay deltas separate

Compare the no-decay group, optimizer paths, and atomic commit

Skip decay

3 Compare trajectories across unequal curvature

SGD trajectory

AdamW trajectory

4 Commit every checked value into its live leaf

5 Check identity, retained gradients, and rollback

Predict before running the optimizer

Train the first cumulative neural language model next

Update and correct gradient moments

Keep adaptive and decay deltas separate

Compare trajectories across unequal curvature

Commit every checked value into its live leaf

Check identity, retained gradients, and rollback