We can't find the internet
Attempting to reconnect
Something went wrong!
Attempting to reconnect
← Normalization step 4 of 6
RMSNorm
RMSNorm rescales a vector by its root mean square (RMS). Square its entries, average those squares, add epsilon, then take a square root. Divide the original entries by that value and apply a learned scale gamma.
Unlike LayerNorm, RMSNorm does not subtract the mean. For [1, 1]
with gamma [1, 1] and epsilon 0, it returns [1, 1]; centering that
row would instead produce zeros. Its output need not have mean zero.
Your task
Implement rmsnorm(x, gamma, eps=1e-6):
xis a floating-point tensor of shape(N, D), with positive dimensions.gammais a floating-point scale tensor of shape(D,).epsis nonnegative. If it is zero, every row has nonzero mean square.- Return a tensor of shape
(N, D).
For each row, compute denominator = sqrt(mean(x ** 2) + eps) over the
D features, then return (x / denominator) * gamma. Preserve the reduced
axis so each row's denominator broadcasts over its features. Do not
subtract the mean, add a bias, or normalize across the batch.
For x = [[0.001, -0.001]], gamma = [1, 2], and eps = 1e-6,
the mean square is 1e-6 and the denominator is sqrt(2e-6).
The output is approximately [[0.707107, -1.414214]]. Epsilon has a
visible effect at this scale: adding it after the square root gives a
different answer. With positive epsilon, an all-zero row returns zeros.
The diagram shows which entries share statistics. This exercise has one feature per column and no spatial axes. With RMSNorm selected, each row shares a denominator; the diagram does not display the resulting values.
Loading visualization…
Reference: Root Mean Square Layer Normalization.
Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.
Decide which values share a mean and variance
Many normalization methods subtract a mean and divide by a standard deviation:
$$\hat x=\frac{x-\mu}{\sqrt{\sigma^2+\epsilon}}.$$
The small positive epsilon keeps the denominator away from zero. The main
difference between methods is which values are used to compute the mean and
variance.
- Batch normalization usually computes one set of statistics per channel, combining examples and spatial positions during training. Its usual inference behavior uses running statistics collected during training.
- Layer normalization computes statistics over selected feature dimensions within each example. Different examples do not share statistics.
- Group normalization divides channels into groups. Within each example, it computes statistics over each group's channels and spatial positions.
For the feature vector [1, 3], the mean is 2 and population variance is 1.
Subtracting the mean gives [-1, 1]; dividing by sqrt(1 + epsilon) leaves values
close to [-1, 1]. Computing statistics over the batch axis instead would combine
values from different examples and produce a different operation.
RMSNorm rescales without centering
RMSNorm divides by the root mean square, with epsilon added for stability:
$$\hat x=\frac{x}{\sqrt{\operatorname{mean}(x^2)+\epsilon}}.$$
It does not subtract the mean, so its output need not have mean zero. For example,
[1, 1] stays close to [1, 1] rather than becoming [0, 0].
Normalization layers may apply learnable scale and bias after normalizing. Their shapes must broadcast over the intended axes. These extra parameters mean that the final output need not have mean zero or variance one.
Match the formula and the source of the statistics
Check whether epsilon goes inside or outside the square root. Also check whether the variance divides by the number of values (population variance) or one fewer (sample variance). These choices change the answer.
For feature preprocessing, estimate statistics on training data and reuse them on validation and test data. Using test data to estimate them leaks information into training. An individual exercise may instead ask you to use only the input it supplies. Follow that contract and handle constant features as specified.
Your code, saved here
Sign in to access your accepted solutions and saved draft history. You can keep reading the lesson without an account.
Sign in to your accountStuck?
JAX reference solution
Sign in to track your attempts toward unlocking the reference solution.