Foundations

Conditioned RMS Normalization

MediumRaw PyTorchAbout 20 min1 public example plus private tests

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

Your task

Scale each valid feature vector by its root mean square without subtracting a mean, then use the conditioning vector to change each channel’s gain and shift. The gain is 1 plus the projected scale; return the modulated vector with the same shape as x and zero at invalid positions.

Function signature

conditioned_rms_normalization(
    x: torch.Tensor,
    valid: torch.Tensor,
    condition: torch.Tensor,
    w_mod: torch.Tensor,
    eps: float = 1e-05,
) -> torch.Tensor

Worked example

One condition controls two different channels

For one valid position with x=[3,4], mean square is 12.5. Illustrative eps=3.5 gives denominator $\sqrt{12.5+3.5}=4$ , so RMS normalization gives [0.75,1]. Condition 1 becomes $s=\operatorname{SiLU}(1)=1/(1+e^{-1})\approx0.731059$ . The modulation tensor [[[1,0],[0,1]]] maps this to shift [s,0] in slot 0 and scale offset [0,s] in slot 1.

The output is $[0.75+s,1(1+s)]\approx[1.481059,1.731059]$ . No channel mean is subtracted, and the scale offset is added to 1 without a sigmoid.

Parameters

Parameter	Shape and meaning
`x`torch.Tensor	$[\ldots, C]$ Features to normalize independently at each position.
`valid`torch.Tensor	$[\ldots]$ True marks a usable position; invalid x and condition may contain nonfinite padding.
`condition`torch.Tensor	$[\ldots, S]$ Conditioning features with the same leading position axes as x.
`w_mod`torch.Tensor	$[S, 2, C]$ Map SiLU(condition) to shift in slot 0 and scale in slot 1.
`eps`floatDefault `1e-05`	Positive stabilizer inside sqrt(mean of squared x channels + eps), before condition-dependent gain and shift.

Return

torch.Tensor · $[\ldots, C]$

Return Floating [...,C], normalized/modulated features, not a residual sum; invalid rows zero.

Required behavior

valid is boolean over all leading position axes. Sanitize invalid x and condition before reductions, and zero every invalid output after any shift or bias.
Compute $m=SiLU(condition)W_{mod}$ with shift in slot 0 and scale in slot 1. Return $RMSNorm(x)\odot(1+m_1)+m_0$ ; do not center x or apply sigmoid to the scale.
Normalize over the final channel axis; eps defaults to $10^{-5}$ and is added inside the square root. Supplied weights use input-channel-first orientation.
Use CPU float32/float64 tensors with one shared floating dtype/device. Preserve strided 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.
Here $\operatorname{RMSNorm}(x)=x/\sqrt{\operatorname{mean}_{last}(x^2)+eps}$ and $\operatorname{SiLU}(c)=c\,\operatorname{sigmoid}(c)$ . Compute the mean only over C channels; do not subtract a mean.

Why this matters

Normalization controls feature magnitude; conditioning changes which normalized channels the next block emphasizes.

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 conditioned_rms_normalization(
    x: torch.Tensor,  # (..., C)
    valid: torch.Tensor,  # (...)
    condition: torch.Tensor,  # (..., S)
    w_mod: torch.Tensor,  # (S, 2, C)
    eps: float = 1e-05,
) -> torch.Tensor:
    """Return a tensor with shape (..., C)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement conditioned_rms_normalization")

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.