Foundations

AlphaFold 3 Adaptive Normalization

MediumRaw PyTorchAbout 20 min2 public examples plus private tests

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

Your task

Normalize each valid feature vector, then use its conditioning vector to choose a sigmoid gain and an additive shift for every output channel. The returned vector has x’s shape, and invalid positions return zero.

Function signature

af3_adaptive_normalization(
    x: torch.Tensor,
    valid: torch.Tensor,
    condition: torch.Tensor,
    cond_scale: torch.Tensor,
    w_scale: torch.Tensor,
    b_scale: torch.Tensor,
    w_shift: torch.Tensor,
    eps: float = 1e-05,
) -> 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]$ Floating [...,C], C>=1, arbitrary nonempty leading dimensions.
`valid`torch.Tensor	$[\ldots]$ Bool [...], True=valid; False payloads may be NaN/Inf.
`condition`torch.Tensor	$[\ldots, S]$ Supplied conditioning vector at the same positions as x.
`cond_scale`torch.Tensor	$[S]$ Learned [S] multiplier for normalized condition, no condition bias.
`w_scale`torch.Tensor	$[S, C]$ Projection from normalized conditioning channels to feature scale logits.
`b_scale`torch.Tensor	$[C]$ Scale-logit bias.
`w_shift`torch.Tensor	$[S, C]$ Bias-free shift projection.
`eps`floatDefault `1e-05`	Positive stabilizer inside all population-variance or mean-square normalization roots; it is not added to sigmoid gains.

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.
Let $c=LN(condition)\odot cond\_scale$ . Return $LN(x)\odot sigmoid(cW_{scale}+b_{scale})+cW_{shift}$ . Both LN operations use population variance without extra learned affine parameters; the scale gate is sigmoid, with no added one.
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 $LN(v)=(v-mean(v))/\sqrt{mean((v-mean(v))^2)+eps}$ , reducing only the final axis, and $sigmoid(z)=1/(1+e^{-z})$ .

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 af3_adaptive_normalization(
    x: torch.Tensor,  # (..., C)
    valid: torch.Tensor,  # (...)
    condition: torch.Tensor,  # (..., S)
    cond_scale: torch.Tensor,  # (S,)
    w_scale: torch.Tensor,  # (S, C)
    b_scale: torch.Tensor,  # (C,)
    w_shift: torch.Tensor,  # (S, 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 af3_adaptive_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.