Foundations · INTEGRATION

Adaptive LayerNorm and RMS condition modulation

MediumRaw PyTorchAbout 20 min3 public examples plus private tests

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

Your task

Transform each valid feature vector using the selected normalization rule. Plain mode centers and scales channels; AF3 mode uses normalized conditioning to set sigmoid gains and shifts; RMS mode uses uncentered features with condition-derived one-plus-scale and shift. Return the transformed vector, with invalid positions zero.

Function signature

adaptive_layer_normalization(
    x: torch.Tensor,
    valid: torch.Tensor,
    parameters: dict,
    condition: torch.Tensor | None = None,
    cond_scale: torch.Tensor | None = None,
    eps: float = 1e-05,
    mode: str = 'af3',
) -> torch.Tensor

Worked example

Conditioning changes both gain and shift

In AF3 mode, take one valid position with x=[1,3], condition=[1,-1], illustrative eps=3, condition scale [1,1], scale matrix 2*I, zero scale bias, and shift matrix I. Normalizing x gives [-0.5,0.5]; normalizing condition gives [0.5,-0.5]. Thus the scale logits are [1,-1], the gains are $[\sigma(1),\sigma(-1)]\approx[0.731059,0.268941]$ , and the shift is [0.5,-0.5].

The output is $[-0.5\sigma(1)+0.5,\ 0.5\sigma(-1)-0.5]\approx[0.134471,-0.365529]$ . Both gain and shift depend on the condition. Here $\sigma(z)=1/(1+e^{-z})$ .

Parameters

Parameter	Shape and meaning
`x`torch.Tensor	$[\ldots, C]$ $C\ge 1$ , arbitrary nonempty leading dimensions.
`valid`torch.Tensor	$[\ldots]$ `True` means `valid`; `False` payloads may be `NaN`/`Inf`.
`condition`torch.Tensor \| NoneDefault `None`	$[\ldots, S]$ Conditional modes only: floating $[\ldots,S]$ , same leading shape, $S\ge 1$ .
`cond_scale`torch.Tensor \| NoneDefault `None`	$[S]$ af3 only: learned $[S]$ multiplier for normalized `condition`, no `condition` `bias`.
`parameters`dict	Supplied model parameters, keyed as shown below.
`eps`floatDefault `1e-05`	Positive stabilizer inside all population-variance or mean-square normalization roots; it is not added to sigmoid gains.
`mode`strDefault `'af3'`	Choose "plain" for affine LayerNorm, "af3" for sigmoid-gain adaptive LayerNorm, or "rms" for SiLU-conditioned shift and 1+scale RMS modulation.

Return

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

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

Required behavior

Let $\operatorname{LN}(z)=(z-\mathrm{mean\_last}(z))/\sqrt{\mathrm{mean\_last}((z-\mathrm{mean\_last}(z))^{2})+\varepsilon}$ , population variance and no implicit affine.
mode='af3' (default): $c=\operatorname{LN}(\mathrm{condition})\cdot \mathrm{cond\_scale}$ , $\mathrm{output}=\operatorname{LN}(x)\cdot \operatorname{sigmoid}(c\mathbin{@}w_{\mathrm{scale}}+b_{\mathrm{scale}})+c\mathbin{@}w_{\mathrm{shift}}$ . The condition LN has learned scale but no shift; $x$ LN has neither.
The scale gate is sigmoid, not $1+\mathrm{linear}$ .
mode='rms': $[\mathrm{shift},\mathrm{scale}]=\operatorname{SiLU}(\mathrm{condition})$ contracted with w_mod $[S,2,C]$ , then $\mathrm{output}=x/\sqrt{\mathrm{mean\_last}(x^{2})+\varepsilon}\cdot (1+\mathrm{scale})+\mathrm{shift}$ . This is the shift/scale portion of ESMFold2 atom modulation, with no residual gate in this primitive.
mode='plain': $\mathrm{output}=\operatorname{LN}(x)\cdot \mathrm{scale}+\mathrm{bias}$ ; condition/modulation tensors are absent. Shapes are given below; all present floating tensors share dtype/device. eps is positive, default $10^{-5}$ , inside every root; ESMFold2 atom RMS uses $10^{-6}$ in the source and can be selected explicitly.
$\operatorname{SiLU}(z)=z\cdot \operatorname{sigmoid}(z)$ , computed without overflow. First replace invalid $x$ /condition values by zero; zero the final result after modulation/affine bias. Valid entries and weights are finite; invalid payloads may be NaN/Inf. A zero vector has finite derivatives.
No running statistics, residual sum, dropout, random initialization or input mutation. Every mode is last-channel only. Numerical domain: $\varepsilon$ must be finite and remain strictly positive when represented in the input dtype.
Active inputs and parameters must be bounded so the required arithmetic intermediates are finite and within that dtype's representable range: normalization sums/means, centered differences, squares, variance-plus- $\varepsilon$ , projections, residual/statistic updates, and intermediates for the graded first-order derivatives.
Finite inputs alone do not guarantee this; overflow-sized reductions or squares are outside the contract. Representable subnormal $\operatorname{sigmoid}/\operatorname{SiLU}$ tails are in scope and must not be prematurely discarded.

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 adaptive_layer_normalization(
    x: torch.Tensor,  # (..., C)
    valid: torch.Tensor,  # (...)
    parameters: dict,
    condition: torch.Tensor | None = None,  # (..., S)
    cond_scale: torch.Tensor | None = None,  # (S,)
    eps: float = 1e-05,
    mode: str = 'af3',
) -> torch.Tensor:
    """Return a tensor with shape (..., C)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement adaptive_layer_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.