LLM, piece by piece
← All chapters

22 · Content revision 8

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, θ0\theta_0 is the current value of the decay-group parameter decoder.output.weight, and g1g_1 is the accumulated token-mean loss gradient with respect to that same parameter:

θ0=[1,−2],g1=[0.2,−0.4].\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 η=0.1\eta=0.1, β1=β2=0.5\beta_1=\beta_2=0.5, ε=0.1\varepsilon=0.1, and λ=0.1\lambda=0.1. Floating-point arithmetic produces the represented decay product d1=fl⁡(ηλ)=0.010000d_1=\operatorname{fl}(\eta\lambda)=0.010000 and shrinkage factor ρ1=fl⁡(1−d1)=0.990000\rho_1=\operatorname{fl}(1-d_1)=0.990000. This direct worked step uses gradient scale α1=1\alpha_1=1, so the effective gradient is exactly g~1=g1\widetilde g_1=g_1. Predict the adaptive delta and the decay delta separately.

The zero-started moments are m1=[0.1,−0.2]m_1=[0.1,-0.2] and v1=[0.02,0.08]v_1=[0.02,0.08]. Both correction denominators equal 0.50.5, giving m^1=[0.2,−0.4]\hat m_1=[0.2,-0.4] and v^1=[0.04,0.16]\hat v_1=[0.04,0.16]. Therefore:

ηm^1v^1+ε≈[0.066667,−0.08],d1θ0=[0.01,−0.02].\eta\frac{\hat m_1}{\sqrt{\hat v_1}+\varepsilon} \approx[0.066667,-0.08], \qquad d_1\theta_0=[0.01,-0.02].

Subtract both from the old value to predict θ1≈[0.923333,−1.9]\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 gtg_t stored on the named parameter and a validated gradient transform. Its generic uniform form accepts a finite shared scale 0≤αt≤10\leq\alpha_t\leq1, including explicit zero. Chapter 33’s global-norm rule requires a mathematically positive clipping scale; if its ordinary represented scalar would round to zero, that rule uses the normalized form instead. A uniform transform forms the effective gradient

g~t=αtgt.\widetilde g_t=\alpha_t g_t.

The uniform scale is not a second learning rate: it changes only the gradient supplied to the moment recurrences. The optimizer reads gtg_t without changing the stored gradient. If a uniform scale would underflow to zero for very large finite gradients, the stable training path instead represents the same normalization structurally. Let c>0c>0 be the global-norm ceiling in the same units as the gradient norm. With s=max⁡i∣gi∣s=\max_i|g_i| and r=∑i(gi/s)2r=\sqrt{\sum_i(g_i/s)^2}, it supplies each coordinate as (gi/s) (c/r)(g_i/s)\,(c/r). If the positive real ratio c/rc/r rounds to zero in f64, it returns a typed failure before AdamW reads parameter values or optimizer state. This fallback changes neither raw gradient storage nor the separate decay branch. 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 11 preserves every gradient bit used by the calculation, as in this chapter’s worked step; Chapter 33 derives a scale below 11 when the complete gradient norm exceeds its ceiling.

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

g~t=0.25[0.8,−0.4]=[0.2,−0.1].\widetilde g_t=0.25[0.8,-0.4]=[0.2,-0.1].

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

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

mt=β1mt−1+(1−β1)g~t,vt=β2vt−1+(1−β2)g~t2.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:

m^t=mt1−β1t,v^t=vt1−β2t.\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, including the represented decay domain:

m^t=mt1−β1t,v^t=vt1−β2t,dt=fl⁡(ηλ),λ=0  ∨  0<dt<1,ρt=fl⁡(1−dt),θt=θt−1−ηm^tv^t+ε−dtθt−1\hat m_t=\frac{m_t}{1-\beta_1^t},\quad \hat v_t=\frac{v_t}{1-\beta_2^t},\quad d_t=\operatorname{fl}(\eta\lambda),\quad \lambda=0\;\lor\;0<d_t<1,\quad \rho_t=\operatorname{fl}(1-d_t),\quad \theta_t=\theta_{t-1}-\eta\frac{\hat m_t}{\sqrt{\hat v_t}+\varepsilon}-d_t\theta_{t-1}

The final expression subtracts the adaptive delta and dtθt−1d_t\theta_{t-1} as two explicit terms. The separately reported ρt\rho_t describes the corresponding shrinkage factor; the implementation retains the shown subtraction order so its floating-point evidence is reproducible. Decoupling means the second term bypasses gtg_t, g~t\widetilde g_t, mtm_t, and vtv_t. Neither the uniform nor normalized gradient transform affects 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 00 for a no-decay parameter. Zero is an explicit no-decay choice. For positive λ\lambda, the implementation rejects a nonfinite, underflowed-to-zero, unit, or greater-than-unit represented product before AdamW reads optimizer state or parameter values; this check is repeated for every scheduled learning rate.

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 00, avoiding a separate decay term that would directly pull the learned affine scale toward zero.

Keep step state and parameter state distinct

  • θt−1\theta_{t-1} is one named parameter tensor before step tt.
  • gtg_t is its accumulated token-mean loss gradient, which remains stored on the live leaf until the caller clears it.
  • αt\alpha_t is the conceptual uniform clipping scale. An ordinary positive represented scale is passed directly to the uniform transform; a scale that would round to zero is represented by the normalized transform instead. The generic AdamW API also accepts an explicit zero uniform scale.
  • g~t\widetilde g_t is the effective gradient produced by the validated uniform or normalized transform and used by both moment recurrences.
  • mtm_t and vtv_t are name-keyed moving estimates with the same shape.
  • β1\beta_1 and β2\beta_2 retain past first and second raw moments, respectively.
  • m^t\hat m_t and v^t\hat v_t correct their early-step zero bias.
  • η>0\eta>0 is the learning rate; λ=0\lambda=0 means no decay, while λ>0\lambda>0 requires the represented product dt=fl⁡(ηλ)d_t=\operatorname{fl}(\eta\lambda) to satisfy 0<dt<10<d_t<1.
  • dt=fl⁡(ηλ)d_t=\operatorname{fl}(\eta\lambda) is the represented decay product and ρt=fl⁡(1−dt)\rho_t=\operatorname{fl}(1-d_t) is the represented shrinkage factor; the fixture prints d1=0.010000d_1=0.010000 and ρ1=0.990000\rho_1=0.990000.
  • ε>0\varepsilon>0 stabilizes the adaptive denominator after its square root.
  • θt\theta_t is the updated value written into that same trainable leaf only after the whole set succeeds.

On the fresh probe, m0=v0=0m_0=v_0=0 and g1=0g_1=0, so both new moments and the adaptive delta are zero. A decay-group parameter still shrinks through d1θ0d_1\theta_0, while a no-decay parameter remains unchanged. After an earlier nonzero gradient, gt=0g_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:

θ←θ−ηg.\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 L2L_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 ut=μut−1+gtu_t=\mu u_{t-1}+g_t and θt=θt−1−ηut\theta_t=\theta_{t-1}-\eta u_t. Coupled L2L_2 feeds gt+λθt−1g_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 αt=1\alpha_t=1 example, that input is gtg_t; after Chapter 33 applies clipping, it is g~t\widetilde g_t. Here utu_t is the retained update velocity and 0≤μ<10\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−β1t1-\beta_1^t and 1−β2t1-\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 L2L_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 77B to 6565B parameters. The course’s optimizer writes each checked update into the existing named parameter leaf; Chapter 33 supplies a scheduled learning rate, a validated uniform or normalized gradient transform, 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 77B to 6565B parameters. Their recipe uses β1=0.9\beta_1=0.9, β2=0.95\beta_2=0.95, weight decay 0.10.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 L2L_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 and represented-decay domains before any optimizer state or parameter value is read. Zero decay is explicit; positive decay must produce a finite represented product strictly between zero and one, including when a scheduled learning rate is supplied:

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,
            });
        }
        let decay_product = learning_rate * weight_decay;
        if weight_decay != 0.0
            && (!decay_product.is_finite() || decay_product <= 0.0 || decay_product >= 1.0)
        {
            return Err(AdamWError::InvalidDecayProduct {
                learning_rate,
                weight_decay,
                product: decay_product,
            });
        }
        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
    }

    /// The represented product that controls decoupled shrinkage.
    pub fn decay_product(self) -> f64 {
        self.learning_rate * self.weight_decay
    }

    /// The represented factor corresponding to the decoupled decay product.
    pub fn shrinkage_factor(self) -> f64 {
        1.0 - self.decay_product()
    }
}

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 an ordinary uniform αt\alpha_t while still returning only the step number. The structured method step_with_learning_rate_and_gradient_transform additionally accepts the normalized fallback. The four earlier ordinary and traced entry points behave exactly as αt=1\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 g~t\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,
            AdamWGradientTransform::uniform(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,
            AdamWGradientTransform::uniform(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,
            AdamWGradientTransform::uniform(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> {
        self.step_with_learning_rate_and_gradient_transform(
            parameters,
            learning_rate,
            AdamWGradientTransform::uniform(gradient_scale),
        )
    }

    /// Applies one scheduled rate and one validated gradient transformation.
    ///
    /// Both forms affect only Adam's moment input. The normalized form prevents
    /// the complete transform from collapsing solely because its equivalent
    /// uniform scalar underflowed; individual tiny coordinates may still round
    /// to zero. Neither form changes raw gradients or the decoupled
    /// weight-decay branch.
    pub fn step_with_learning_rate_and_gradient_transform(
        &mut self,
        parameters: &[NamedParameter],
        learning_rate: f64,
        gradient_transform: AdamWGradientTransform,
    ) -> Result<u64, AdamWError> {
        let step_config = self.config.with_learning_rate(learning_rate)?;
        self.step_with_config(parameters, step_config, gradient_transform, 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,
            AdamWGradientTransform::uniform(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 and its represented decay product while constructing its per-step configuration. It then enters the shared path, which validates the uniform or normalized gradient transform 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 β1t\beta_1^t and β2t\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 neither gradient transform is applied to it.

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_transform: AdamWGradientTransform,
        mut observer: O,
    ) -> Result<O::Output, AdamWError> {
        let gradient_transform = gradient_transform.validate()?;
        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_transform,
                },
                &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 scalar configuration, represented-decay, and uniform or normalized gradient-transform 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,
    },
    InvalidDecayProduct {
        learning_rate: f64,
        weight_decay: f64,
        product: f64,
    },
    InvalidGradientScale {
        value: f64,
    },
    InvalidGradientNormalization {
        divisor: f64,
        multiplier: 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::InvalidDecayProduct {
                learning_rate,
                weight_decay,
                product,
            } => write!(
                formatter,
                "positive weight decay requires represented learning_rate * weight_decay in (0,1), got {learning_rate} * {weight_decay} = {product}"
            ),
            Self::InvalidGradientScale { value } => write!(
                formatter,
                "gradient scale must be finite in the closed interval [0,1], got {value}"
            ),
            Self::InvalidGradientNormalization {
                divisor,
                multiplier,
            } => write!(
                formatter,
                "normalized gradient transform requires a finite positive divisor and a finite multiplier in (0,divisor], got divisor {divisor} and multiplier {multiplier}"
            ),
            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−1210^{-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 α1=1\alpha_1=1, so its input-gradient field represents both the raw g1g_1 and the effective g~1\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}|decay_product={:.6}|shrinkage_factor={:.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.config.decay_product(),
        evidence.config.shrinkage_factor(),
        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=1t=1
Learning rate
η=0.100000\eta=0.100000
Moment rates
β1=0.500000\beta_1=0.500000 β2=0.500000\beta_2=0.500000
Denominator stabilizer
ε=0.100000\varepsilon=0.100000
Weight decay
λ=0.100000\lambda=0.100000
Represented decay product
dt=0.010000d_t=0.010000
Represented shrinkage factor
ρt=0.990000\rho_t=0.990000

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−β1t=0.5000001-\beta_1^t=0.500000 1−β2t=0.5000001-\beta_2^t=0.500000

Named parameter
decoder.output.weight
Parameter group
Apply decay decay
Shape
[2]
Value before θt−1\theta_{t-1}
[1,−2]\left[1,-2\right]
Accumulated gradient gtg_t
[0.2,−0.4]\left[0.2,-0.4\right]
Adaptive path
First moment mtm_t
[0.1,−0.2]\left[0.1,-0.2\right]
Second raw moment vtv_t
[0.02,0.08]\left[0.02,0.08\right]
Corrected first moment m^t\hat m_t
[0.2,−0.4]\left[0.2,-0.4\right]
Corrected second raw moment v^t\hat v_t
[0.04,0.16]\left[0.04,0.16\right]
Adaptive path

Adaptive delta

[0.066667,−0.08]\left[0.066667,-0.08\right]
Value before θt−1\theta_{t-1} [1,−2]\left[1,-2\right]
Decay bypass · Apply decay

Decay delta

[0.01,−0.02]\left[0.01,-0.02\right]
−
Subtract both deltas

Updated value θt\theta_t

[0.923333,−1.9]\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 ηm^t/(v^t+ε)\eta\hat m_t/(\sqrt{\hat v_t}+\varepsilon) − Decay bypass dtθt−1d_t\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
θt−1=[0.5]\theta_{t-1}=\left[0.5\right]
Accumulated gradient
gt=[0]g_t=\left[0\right]
Corrected first moment
m^t=[0]\hat m_t=\left[0\right]
Corrected second raw moment
v^t=[0]\hat v_t=\left[0\right]
Adaptive delta
ηm^tv^t+ε=[0]\frac{\eta\hat m_t}{\sqrt{\hat v_t}+\varepsilon}=\left[0\right]
Decay delta
d1θ0=[0]d_1\theta_0=\left[0\right]
Updated value
θt=[0.5]\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)=12(x2+4y2)q(x,y)=\frac12(x^2+4y^2) diag⁡(H)=[1,4]\operatorname{diag}(H)=\left[1,4\right]

SGD trajectory
  1. Optimizer step t=0t=0 [1,1]\left[1,1\right]
  2. Optimizer step t=1t=1 [0.9,0.6]\left[0.9,0.6\right]
  3. Optimizer step t=2t=2 [0.81,0.36]\left[0.81,0.36\right]
  4. Optimizer step t=3t=3 [0.729,0.216]\left[0.729,0.216\right]
  5. Optimizer step t=4t=4 [0.6561,0.1296]\left[0.6561,0.1296\right]
AdamW trajectory
  1. Optimizer step t=0t=0 [1,1]\left[1,1\right]
  2. Optimizer step t=1t=1 [0.899091,0.892439]\left[0.899091,0.892439\right]
  3. Optimizer step t=2t=2 [0.799889,0.786278]\left[0.799889,0.786278\right]
  4. Optimizer step t=3t=3 [0.702629,0.681677]\left[0.702629,0.681677\right]
  5. Optimizer step t=4t=4 [0.60758,0.578823]\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
d1θ0=0.030000d_1\theta_0=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)=12(x2+4y2)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

  1. Derive m1m_1, v1v_1, m^1\hat m_1, and v^1\hat v_1 for the fixed vector.
  2. Predict the two deltas and θ1\theta_1 before rounding to six decimals.
  3. Start a fresh optimizer for decay-group value 33 with g1=0g_1=0 and predict the new value. Then explain why a later zero gradient can still have an adaptive contribution after nonzero history.
  4. Explain why reversing two parameter records must not reverse their moments.
  5. Predict what changes if the second candidate overflows after the first candidate has been calculated successfully.
  6. 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.
  7. Order plain SGD, momentum, Adam with coupled L2L_2, and AdamW, and state which terms enter optimizer memory.
  8. Explain why the output weight receives decay while the normalization scale does not.
  9. Compare the fixed SGD and AdamW paths on the unequal-curvature quadratic.
  10. For η=0.1\eta=0.1 and λ=0.1\lambda=0.1, state the represented decay product and shrinkage factor. What must happen when positive decay yields a represented product that is zero, equal to one, greater than one, or nonfinite?
Check the predictions
  1. m1=[0.1,−0.2]m_1=[0.1,-0.2], v1=[0.02,0.08]v_1=[0.02,0.08], m^1=[0.2,−0.4]\hat m_1=[0.2,-0.4], and v^1=[0.04,0.16]\hat v_1=[0.04,0.16].
  2. The adaptive delta is approximately [0.066667,−0.08][0.066667,-0.08], decay is [0.01,−0.02][0.01,-0.02], and θ1≈[0.923333,−1.9]\theta_1\approx[0.923333,-1.9].
  3. With fresh zero moments, the adaptive delta stays zero, decay is 0.030.03, and the value becomes 2.972.97. After nonzero history, stored moments can keep the adaptive delta nonzero even when the current gradient is zero.
  4. The state map uses stable names; slice order is presentation only.
  5. Nothing commits: both parameter values, both gradients, both node identities, and all optimizer state remain unchanged.
  6. 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.
  7. Plain SGD uses only the current gradient; momentum remembers a decaying direction; coupled L2L_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.
  8. 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 dtθt−1d_t\theta_{t-1} from it. It assigns decoder.norm.scale to no-decay, so that parameter’s effective λ\lambda is 00; this avoids a decay term that directly pulls the learned normalization scale toward zero.
  9. SGD reduces the high-curvature coordinate much faster; AdamW scales coordinates through their moment history and also applies a separate decay contribution.
  10. d1=0.010000d_1=0.010000 and ρ1=0.990000\rho_1=0.990000. Positive decay is accepted only when the represented product is finite and strictly between zero and one; underflow to zero, one, greater than one, and nonfinite products are typed failures before AdamW reads parameter values or optimizer state.

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.