25 · Content revision 5

Normalize scale without subtracting the mean

Implement last-axis RMSNorm, trace its input and gain gradients, and separate ideal scale invariance from epsilon-dominated behavior near zero.

Predict RMS, then test the epsilon boundary

Start with $x=[3,4]$ , learned gain $g=[1.5,0.5]$ , and $\varepsilon=10^{-5}$ . Before running the example, compute the mean square:

\frac{3^2+4^2}{2}=12.5.

The reciprocal RMS is about $0.282843$ . Predict the unscaled normalized vector, then multiply coordinatewise by $g$ . The fixture should produce approximately $[0.848528,1.131370]$ before the gain and $[1.272792,0.565685]$ after it.

Now multiply the input by positive factor $10$ . To test scale invariance, use the RMS-rescaled vector $\hat{x}$ before learned gain $g$ is applied. The final output after applying $g$ is not part of this comparison. Predict $\hat{x}$ first with $\varepsilon=0$ , then with $\varepsilon=10^{-5}$ . Repeat that pre-gain comparison for $x=[0.0003,0.0004]$ , whose mean square is smaller than the stabilizer.

Normalize the final feature axis

For one final-axis vector, RMSNorm is

\operatorname{RMSNorm}(x)=g\odot\frac{x}{\sqrt{\frac{1}{d}\sum_i x_i^2+\varepsilon}}

The statistic is shared by all $d$ coordinates in that vector. The learned gain then rescales individual features. No feature mean is subtracted, and no example contributes statistics to another example.

For a positive scalar $a$ and a nonzero vector, the epsilon-zero ideal obeys

\operatorname{RMSNorm}_{0}(ax)=\operatorname{RMSNorm}_{0}(x).

With finite epsilon, the denominator becomes $\sqrt{a^2\operatorname{RMS}(x)^2+\varepsilon}$ , so epsilon does not grow with the signal. The normalized mean square is

\operatorname{mean}(\hat{x}^2)= \frac{\operatorname{RMS}(x)^2} {\operatorname{RMS}(x)^2+\varepsilon}.

It approaches $1$ when the signal dominates epsilon and approaches $0$ when epsilon dominates the signal. This is why practical RMSNorm is approximately scale-invariant away from zero, not exactly invariant everywhere.

Keep scale, gain, and axes distinct

$x$ is one feature vector taken from the input’s final axis.
$d$ is that axis’s nonzero width and $i$ selects one coordinate.
$x_i$ is coordinate $i$ and $\sum_i x_i^2/d$ is the mean square of the coordinates in vector $x$ .
$\varepsilon\ge0$ stabilizes the reciprocal square root.
$\hat{x}$ is the RMS-rescaled vector before learned featurewise scaling.
$g\in\mathbb{R}^{d}$ is the learned gain and $\odot$ is elementwise multiplication.
The output has the same complete shape as the input.

RMSNorm does not force every coordinate to magnitude $1$ . It controls the vector’s aggregate root-mean-square scale. It also does not center the vector, clip values, or mix examples across a mini-batch.

From batch statistics to pre-RMSNorm language models

BatchNorm couples a training example to mini-batch statistics, while LayerNorm removes that cross-example dependency but still computes and subtracts a per-example feature mean before rescaling.

The earlier primary source is Ioffe and Szegedy, Batch Normalization. Ioffe and Szegedy make normalization part of the architecture and compute its training statistics for each mini-batch.

The second earlier source is Ba, Kiros, and Hinton, Layer Normalization. Ba, Kiros, and Hinton compute mean and variance across the summed inputs of one layer for one training case, avoiding dependencies between training cases.

RMSNorm removes mean subtraction and keeps RMS rescaling over the feature vector; LLaMA later used RMSNorm before each Transformer sublayer.

The first later source is Zhang and Sennrich, Root Mean Square Layer Normalization. Zhang and Sennrich remove the mean statistic, normalize by RMS, and retain the epsilon-free formulation’s positive rescaling invariance while giving up recentering invariance. The paper’s displayed formula has no epsilon. This implementation adds positive epsilon so zero and tiny vectors remain finite; the resulting near-zero behavior follows from that modified denominator.

The second later source is Touvron et al., LLaMA. Touvron et al. normalize the input of each Transformer sublayer and identify RMSNorm as the normalization function.

A pre-normalized decoder feeds a controlled-scale residual-stream vector into each attention or feed-forward branch while leaving the identity path outside that normalization operation.

The executable contrast holds anchor $[1,3]$ fixed and changes only its batch companion. BatchNorm changes the anchor output; LayerNorm and RMSNorm remain per-example. LayerNorm centers the anchor, while RMSNorm keeps a nonzero output mean:

Compare normalization axes and centering behavior rust/demos/ch25-rmsnorm/src/lib.rs#historical-normalization-contrast

fn batch_normalize_anchor(anchor: &[f64], companion: &[f64], epsilon: f64) -> Vec<f64> {
    anchor
        .iter()
        .zip(companion)
        .map(|(&anchor, &companion)| {
            let mean = (anchor + companion) / 2.0;
            let variance = ((anchor - mean).powi(2) + (companion - mean).powi(2)) / 2.0;
            (anchor - mean) / (variance + epsilon).sqrt()
        })
        .collect()
}

fn layer_normalize(values: &[f64], epsilon: f64) -> Vec<f64> {
    let mean = values.iter().sum::<f64>() / values.len() as f64;
    let variance = values
        .iter()
        .map(|value| (value - mean).powi(2))
        .sum::<f64>()
        / values.len() as f64;
    values
        .iter()
        .map(|value| (value - mean) / (variance + epsilon).sqrt())
        .collect()
}

fn historical_evidence() -> Result<HistoryEvidence, FixtureError> {
    let anchor = [1.0, 3.0];
    let batch_anchor_a = batch_normalize_anchor(&anchor, &[5.0, 7.0], EPSILON);
    let batch_anchor_b = batch_normalize_anchor(&anchor, &anchor, EPSILON);
    let layer_norm = layer_normalize(&anchor, EPSILON);
    let rms_norm = normalized(&anchor, IDEAL_EPSILON)?;
    let rms_mean = rms_norm.iter().sum::<f64>() / rms_norm.len() as f64;
    Ok(HistoryEvidence {
        batch_anchor_a,
        batch_anchor_b,
        layer_norm,
        rms_norm,
        rms_mean,
    })
}

Compose RMSNorm from cumulative tape operations

RmsNorm owns one rank-one NamedParameter and no bias. It validates epsilon, the gain name and width, input rank, and final-axis width. The forward pass then composes the existing multiply, retained last-axis mean, scalar addition, logarithm, exponential, broadcast, and multiply operations. Both the input and gain therefore remain connected to the cumulative reverse-mode tape:

Normalize the final axis with existing differentiable operations rust/crates/llm-from-scratch/src/nn/rmsnorm.rs#rmsnorm-layer

/// Inspectable tensors from one composed RMSNorm forward pass.
#[derive(Clone, Debug)]
pub struct RmsNormForward {
    mean_square: TensorValue,
    inverse_rms: TensorValue,
    normalized: TensorValue,
    output: TensorValue,
}

impl RmsNormForward {
    pub fn mean_square(&self) -> &TensorValue {
        &self.mean_square
    }

    pub fn inverse_rms(&self) -> &TensorValue {
        &self.inverse_rms
    }

    pub fn normalized(&self) -> &TensorValue {
        &self.normalized
    }

    pub fn output(&self) -> &TensorValue {
        &self.output
    }

    pub fn into_output(self) -> TensorValue {
        self.output
    }
}

/// One learned gain applied after final-axis root-mean-square rescaling.
#[derive(Clone, Debug)]
pub struct RmsNorm {
    epsilon: f64,
    feature_width: usize,
    parameters: NamedParameters,
}

impl RmsNorm {
    /// Creates a gain vector initialized to one.
    pub fn new(
        gain_name: impl Into<String>,
        feature_width: usize,
        epsilon: f64,
    ) -> Result<Self, RmsNormError> {
        validate_epsilon(epsilon)?;
        if feature_width == 0 {
            return Err(RmsNormError::EmptyFeatureWidth);
        }
        let gain_name = gain_name.into();
        validate_name(&gain_name)?;
        let mut values = Vec::new();
        values.try_reserve_exact(feature_width).map_err(|_| {
            RmsNormError::GainAllocationFailed {
                elements: feature_width,
            }
        })?;
        values.resize(feature_width, 1.0);
        let gain = NamedParameter::from_tensor(
            gain_name,
            Tensor::from_vec(vec![feature_width], values).map_err(InitializationError::from)?,
        )?;
        Self::from_gain(gain, epsilon)
    }

    /// Builds the layer from one externally named rank-one trainable gain.
    pub fn from_gain(gain: NamedParameter, epsilon: f64) -> Result<Self, RmsNormError> {
        validate_epsilon(epsilon)?;
        let shape = gain.tensor().shape();
        if shape.len() != 1 {
            return Err(RmsNormError::GainRank { shape });
        }
        let feature_width = shape[0];
        if feature_width == 0 {
            return Err(RmsNormError::EmptyFeatureWidth);
        }
        Ok(Self {
            epsilon,
            feature_width,
            parameters: NamedParameters::try_new(vec![gain])?,
        })
    }

    /// Normalizes the final feature axis and returns only the scaled output.
    pub fn forward(&self, input: &TensorValue) -> Result<TensorValue, RmsNormError> {
        self.forward_with_intermediates(input)
            .map(RmsNormForward::into_output)
    }

    /// Normalizes the final feature axis and preserves each teaching value.
    pub fn forward_with_intermediates(
        &self,
        input: &TensorValue,
    ) -> Result<RmsNormForward, RmsNormError> {
        let shape = input.shape();
        let Some(&actual_width) = shape.last() else {
            return Err(RmsNormError::InputRankZero);
        };
        if actual_width != self.feature_width {
            return Err(RmsNormError::InputWidthMismatch {
                expected: self.feature_width,
                actual: actual_width,
            });
        }
        let feature_axis = shape.len() - 1;
        let squared = input
            .mul(input)
            .map_err(autodiff_error(RmsNormStage::Square))?;
        let mean_square = squared
            .mean_axis(feature_axis, true)
            .map_err(autodiff_error(RmsNormStage::MeanSquare))?;
        if self.epsilon == 0.0 {
            for (row, value) in mean_square.value().as_slice().iter().enumerate() {
                if *value == 0.0 {
                    return Err(RmsNormError::ZeroEnergyRow { row });
                }
            }
        }
        let epsilon = scalar_constant(self.epsilon, RmsNormStage::EpsilonConstant)?;
        let stabilized = mean_square
            .add(&epsilon)
            .map_err(autodiff_error(RmsNormStage::Stabilize))?;
        let log_mean_square = stabilized
            .log()
            .map_err(autodiff_error(RmsNormStage::LogMeanSquare))?;
        let exponent = scalar_constant(-0.5, RmsNormStage::ExponentConstant)?;
        let scaled_log = log_mean_square
            .mul(&exponent)
            .map_err(autodiff_error(RmsNormStage::ScaleLog))?;
        let inverse_rms = scaled_log
            .exp()
            .map_err(autodiff_error(RmsNormStage::ReciprocalRoot))?;
        let normalized = input
            .mul(&inverse_rms)
            .map_err(autodiff_error(RmsNormStage::Normalize))?;
        let output = normalized
            .mul(self.gain().tensor())
            .map_err(autodiff_error(RmsNormStage::ApplyGain))?;

        Ok(RmsNormForward {
            mean_square,
            inverse_rms,
            normalized,
            output,
        })
    }

    pub fn gain(&self) -> &NamedParameter {
        &self.parameters.as_slice()[0]
    }

    pub fn parameters(&self) -> &[NamedParameter] {
        self.parameters.as_slice()
    }

    pub const fn feature_width(&self) -> usize {
        self.feature_width
    }

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

Positive epsilon maps an all-zero row to finite zeros. With zero epsilon, a row whose computed mean square is zero is rejected before the logarithm, including a nonzero floating-point row whose squares underflow. Typed errors keep that boundary visible:

Keep configuration and numeric boundaries explicit rust/crates/llm-from-scratch/src/nn/rmsnorm.rs#rmsnorm-errors

/// A rejected RMSNorm configuration, input, or delegated tape operation.
#[derive(Clone, Debug, PartialEq)]
pub enum RmsNormError {
    InvalidEpsilon {
        value: f64,
    },
    EmptyFeatureWidth,
    GainRank {
        shape: Vec<usize>,
    },
    GainAllocationFailed {
        elements: usize,
    },
    InputRankZero,
    InputWidthMismatch {
        expected: usize,
        actual: usize,
    },
    ZeroEnergyRow {
        row: usize,
    },
    Initialization(InitializationError),
    Autodiff {
        stage: RmsNormStage,
        source: TensorAutodiffError,
    },
}

impl fmt::Display for RmsNormError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::InvalidEpsilon { value } => write!(
                formatter,
                "RMSNorm epsilon must be finite and nonnegative, got {value:?}"
            ),
            Self::EmptyFeatureWidth => {
                formatter.write_str("RMSNorm feature width must be greater than zero")
            }
            Self::GainRank { shape } => {
                write!(
                    formatter,
                    "RMSNorm gain must have rank 1, got shape {shape:?}"
                )
            }
            Self::GainAllocationFailed { elements } => write!(
                formatter,
                "could not reserve storage for {elements} RMSNorm gain values"
            ),
            Self::InputRankZero => formatter.write_str("RMSNorm input must have at least one axis"),
            Self::InputWidthMismatch { expected, actual } => write!(
                formatter,
                "RMSNorm input feature width must be {expected}, got {actual}"
            ),
            Self::ZeroEnergyRow { row } => write!(
                formatter,
                "RMSNorm epsilon is zero but feature row {row} has zero mean square"
            ),
            Self::Initialization(source) => source.fmt(formatter),
            Self::Autodiff { stage, source } => {
                write!(formatter, "RMSNorm {stage}: {source}")
            }
        }
    }
}

impl Error for RmsNormError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::Initialization(source) => Some(source),
            Self::Autodiff { source, .. } => Some(source),
            _ => None,
        }
    }
}

impl From<InitializationError> for RmsNormError {
    fn from(error: InitializationError) -> Self {
        Self::Initialization(error)
    }
}

The worked fixture runs the primary reverse pass with $\bar y=[1,-2]$ . It finds $\bar x\approx[0.407293,-0.305470]$ and $\bar g\approx[0.848528,-2.262741]$ , verifies independent rows for input shape $[2,2]$ , and checks a zero-sized outer batch. The stable parameter name decoder.block.0.attention_norm.gain belongs only to AdamWParameterGroups’ no-decay set. That grouping is an optimizer policy, not part of the RMSNorm formula:

Build forward, backward, epsilon, shape, and optimizer evidence rust/demos/ch25-rmsnorm/src/lib.rs#rmsnorm-fixture

pub fn learner_evidence() -> Result<LearnerEvidence, FixtureError> {
    let primary = primary_once()?;
    let replay = primary_once()?;
    let ideal_scale = scale_evidence("ideal", &INPUT_VALUES, IDEAL_EPSILON)?;
    let production_scale = scale_evidence("production", &INPUT_VALUES, EPSILON)?;
    let near_zero_scale = scale_evidence("near-zero", &TINY_VALUES, EPSILON)?;

    let zero_layer = layer(&[1.0, 1.0], EPSILON)?;
    let zero_output = zero_layer
        .forward(&TensorValue::constant(tensor(&[2], &[0.0, 0.0]))?)?
        .value()
        .as_slice()
        .to_vec();
    let batch_output = layer(&GAIN_VALUES, EPSILON)?
        .forward(&TensorValue::constant(tensor(
            &[2, 2],
            &[3.0, 4.0, 0.0, 5.0],
        ))?)?
        .value_snapshot();

    let groups = AdamWParameterGroups::new([] as [&str; 0], [GAIN_NAME])?;
    let no_decay = groups.decayed_names().count() == 0
        && groups.excluded_names().collect::<Vec<_>>() == [GAIN_NAME];
    let (input_checks, gain_checks, gradcheck_passed) = gradient_evidence(&primary)?;

    Ok(LearnerEvidence {
        replay_bitwise: primary_matches(&primary, &replay),
        primary,
        ideal_scale,
        production_scale,
        near_zero_scale,
        zero_output,
        batch_output,
        history: historical_evidence()?,
        errors: error_evidence()?,
        input_checks,
        gain_checks,
        gradcheck_passed,
        no_decay,
    })
}

Central differences check both input coordinates and both gain coordinates with step $10^{-6}$ and tolerance $2\times10^{-6}$ :

Check both differentiable inputs numerically rust/demos/ch25-rmsnorm/src/lib.rs#rmsnorm-gradcheck

fn gradient_evidence(primary: &PrimaryEvidence) -> Result<(usize, usize, bool), FixtureError> {
    let upstream = &primary.upstream;
    let input_report = sampled_tensor_gradient_check(
        &mut primary.input.clone(),
        &primary.input_gradient.view(),
        STEP,
        TOLERANCE,
        INPUT_VALUES.len(),
        |probe| {
            dot_output(
                &layer(&GAIN_VALUES, EPSILON).expect("fixture layer is valid"),
                probe,
                upstream,
            )
        },
    )?;
    let gain_report = sampled_tensor_gradient_check(
        &mut primary.gain.clone(),
        &primary.gain_gradient.view(),
        STEP,
        TOLERANCE,
        GAIN_VALUES.len(),
        |probe| {
            dot_output(
                &RmsNorm::from_gain(
                    NamedParameter::from_tensor(GAIN_NAME, probe.clone())
                        .expect("probe gain is finite"),
                    EPSILON,
                )
                .expect("probe gain shape is valid"),
                &primary.input,
                upstream,
            )
        },
    )?;
    Ok((
        input_report.checks.len(),
        gain_report.checks.len(),
        input_report.passed && gain_report.passed,
    ))
}

Run the checked RMSNorm fixture rust/demos/ch25-rmsnorm/src/main.rs

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let evidence = ch25_rmsnorm::learner_evidence()?;
    print!("{}", ch25_rmsnorm::render_report(&evidence));
    Ok(())
}

Run cargo run --quiet --locked -p ch25-rmsnorm to print the same checked forward, backward, scale, and boundary values.

Compare rescaling across the epsilon boundary

The diagram brings the checked primary values, gradients, scale comparisons, zero and batch behavior, historical outputs, and rejected requests into one view:

Follow one feature vector through RMSNorm

Read exact Rust-authored values from input through RMS rescaling and learned gain, then compare scale, history, gradient, and rejected-boundary evidence.

Solid input border
Dashed RMS-rescaled border
Double gain-scaled border

Rescale one final-axis feature vector

Input

$x_{0}=3.000000$
$x_{1}=4.000000$

Mean square

$m_2=[12.500000]$

Reciprocal RMS

$r^{-1}=[0.282843]$

$\varepsilon=0.000010$

RMS-rescaled vector

$\hat x_{0}=0.848528$
$\hat x_{1}=1.131370$

Learned gain

$g_{0}=1.500000$
$g_{1}=0.500000$

Gain-scaled output

$y_{0}=1.272792$
$y_{1}=0.565685$

Bar lengths are local visual guides; exact values remain the authority.

Separate the ideal from finite epsilon

Epsilon-zero ideal

$\varepsilon=0.000000,\quad a=10.000000$

Original vector: $\hat x=[0.848528,1.131371]$
Input multiplied by ten: $\widehat{ax}=[0.848528,1.131371]$
Maximum absolute difference: $\Delta_{\mathrm{max}}=0.000000000000000222$

Ordinary production input

$\varepsilon=0.000010,\quad a=10.000000$

Original vector: $\hat x=[0.848528,1.131370]$
Input multiplied by ten: $\widehat{ax}=[0.848528,1.131371]$
Maximum absolute difference: $\Delta_{\mathrm{max}}=0.000000448$

Epsilon-dominated input

$\varepsilon=0.000010,\quad a=10.000000$

Original vector: $\hat x=[0.094281,0.125708]$
Input multiplied by ten: $\widehat{ax}=[0.632456,0.843274]$
Maximum absolute difference: $\Delta_{\mathrm{max}}=0.717566$

The ordinary difference is tiny; the near-zero difference is visible.

Compare where each statistic comes from

The fixed anchor is paired with two different batch companions.
Gain-scaled output	First batch companion	Second batch companion
BatchNorm	$[-0.999999,-0.999999]$	$[0.000000,0.000000]$
LayerNorm	$[-0.999995,0.999995]$	$[-0.999995,0.999995]$
RMSNorm	$[0.447214,1.341641]$	$[0.447214,1.341641]$

RMSNorm output mean: $\operatorname{mean}(\hat x)=0.894427$

Inspect gradients, shapes, and boundaries

Reverse-mode evidence

$\bar y=[1.000000,-2.000000]$

Input gradient: $\bar x=[0.407293,-0.305470]$
Gain gradient: $\bar g=[0.848528,-2.262741]$

Stable gain parameter

decoder.block.0.attention_norm.gain

$g\in\mathbb{R}^{2}$

Optimizer grouping: no_decay=true

All-zero input

$x=[0.000000,0.000000]$

$y=[0.000000,0.000000]$

Accepted · finite=true

Independent final-axis rows

$\operatorname{shape}(y)=[2,2]$

$y_{0,:}=[1.272792,0.565685]$

$y_{1,:}=[0.000000,0.707106]$

axis=last

Rejected boundaries

Rejected rank-zero The input needs at least one axis.
Rejected width-mismatch The final feature width must match the gain width.
Rejected zero-energy-epsilon-zero Zero epsilon cannot normalize a row whose mean square is zero.

Rust-authored checks

$\operatorname{mean}(\hat x^2)=0.999999$

Reverse-mode evidence

input_checks=2

gain_checks=2

$\tau=0.000002$

Accepted

gradcheck=true

replay=bitwise

Read the exact values from left to right: first the input, then the shared RMS factor, the RMS-rescaled vector, the learned gain, and the final output. Solid, dashed, and double borders distinguish those three vector states without relying on color. The scale cards then show why multiplying an ordinary vector by ten barely changes the finite-epsilon result, while the same operation moves a tiny vector far out of the epsilon-dominated regime.

Predict before reading the evidence

Compute the mean square and reciprocal RMS for $x=[3,4]$ .
Apply $g=[1.5,0.5]$ to the normalized vector.
For a nonzero vector and $\varepsilon=0$ , predict the pre-gain RMS-rescaled vector $\hat{x}$ after multiplying the input by $10$ ; do not apply learned gain $g$ .
Explain why a positive $\varepsilon$ changes a near-zero vector much more.
Predict the output for an all-zero row with positive epsilon and zero epsilon.
Decide whether RMSNorm subtracts the feature mean or mixes batch examples.
Predict the output shape and gain-gradient shape for input $[B,T,d]$ .
Explain why no-decay assignment is optimizer policy, not part of the formula.
Place RMSNorm relative to the identity and learned paths in a pre-normalized residual block.

Check the predictions

The mean square is $12.5$ and reciprocal RMS is about $0.282843$ .
The output is approximately $[1.272792,0.565685]$ .
When $\varepsilon=0$ , positive input scaling cancels in $\hat{x}$ , so the pre-gain RMS-rescaled vector is unchanged before learned gain $g$ is applied.
Finite epsilon does not scale with the signal, so it dominates the tiny vector’s denominator.
Positive epsilon returns $[0,0]$ ; zero epsilon rejects a row whose mean square is zero before the logarithm.
Neither: it uses only one example’s final feature axis and does not subtract its mean.
The output shape is $[B,T,d]$ and the gain-gradient shape is $[d]$ .
The optimizer decides which named parameters receive decay; RMSNorm’s equation does not.
The learned branch receives RMSNorm’s output, while the identity path carries the unchanged residual stream.

Project normalized features into query, key, and value tensors next

The cumulative decoder now has differentiable last-axis RMSNorm for each learned residual branch and a stable gain that training can exclude from weight decay. Chapter 26 turns those normalized features into separate bias-free query, key, and value projections, with their attention dimensions made explicit.