Foundations

RMS Normalization

EasyRaw PyTorchAbout 10 min1 public example plus private tests

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Your task

For each final-axis feature vector, divide every channel by the same square root of its mean squared value plus eps. Then apply channel scale and optional bias. The input is not centered, so its common offset still affects the output.

Function signature

rms_norm(
    x: torch.Tensor,
    scale: torch.Tensor,
    bias: torch.Tensor | None = None,
    eps: float = 1e-05,
) -> torch.Tensor

Worked example

RMS normalization preserves the mean direction

The first row has mean square $(3^2+4^2)/2=12.5$ ; eps=0.5 makes the denominator $\sqrt{13}$ . Its first output is $2\cdot3/\sqrt{13}-1\approx0.6641$ . The zero row returns the bias.

import torch

x = torch.tensor([[3.0, 4.0], [0.0, 0.0]], dtype=torch.float64, requires_grad=True)
scale = torch.tensor([2.0, 0.5], dtype=torch.float64, requires_grad=True)
bias = torch.tensor([-1.0, 2.0], dtype=torch.float64, requires_grad=True)
eps = 0.5

result = rms_norm(x=x, scale=scale, bias=bias, eps=eps)

Expected return

torch.tensor([[0.6641005886756874, 2.5547001962252294], [-1.0, 2.0]], dtype=torch.float64)

Parameters

Parameter	Shape and meaning
`x`torch.Tensor	$[\ldots, C]$ Feature vectors to normalize over their final channel axis.
`scale`torch.Tensor	$[C]$ Learned multiplier for each feature channel.
`bias`torch.Tensor \| NoneDefault `None`	$[C]$ Optional learned offset for each channel; omission means zero.
`eps`floatDefault `1e-05`	Strictly positive value added to the mean squared feature value, inside the square root.

Return

torch.Tensor · Exactly x.shape.

Return Normalized and affine-transformed feature vectors, with exactly x’s shape.

Required behavior

The final channel axis has positive width; every preceding axis indexes an independent vector. Empty leading axes are supported.
Compute $m=\mathrm{mean}_{-1}(x^2)$ , then $y=x/\sqrt{m+\varepsilon}\cdot\mathrm{scale}+\mathrm{bias}$ . Do not subtract the mean.
eps is a positive Python number, default 1e-5, added inside the square root. Scale has shape (C,); omitted bias means zero.
Constant rows are normalized by their RMS; the zero row returns bias. Optional bias in RMSNorm is an explicit educational extension.
Use CPU float32/float64 tensors with one shared floating dtype/device. Accept noncontiguous input layouts and do not mutate supplied tensors.
Use raw tensor operations, not a ready-made implementation of the requested primitive. Supplied parameters and random variates must not be reinitialized.

Why this matters

LayerNorm removes a row’s mean before scaling its spread; RMSNorm scales its magnitude without removing the mean.

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def rms_norm(
    x: torch.Tensor,  # (..., C)
    scale: torch.Tensor,  # (C,)
    bias: torch.Tensor | None = None,  # (C,)
    eps: float = 1e-05,
) -> torch.Tensor:
    """Return a tensor with shape x.shape."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement rms_norm")

Ready to implement it?

3 progressive hints, a solution walkthrough, and grading that checks values, shapes and gradients are waiting in the workspace.

Write and run your solution Sign in

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Related exercises

Sources and further reading

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.