MSA and pair representations

Masked outer-product mean

MediumRaw PyTorchAbout 20 min2 public examples plus private tests

AlphaFold2OpenFoldAlphaFold3Boltz-1ProtenixRoseTTAFold All-AtomESMFold2

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

Your task

Convert MSA features into directed residue-pair features. Average projected outer products over sequences where both residues are valid, then return the pair update.

Function signature

outer_product_mean(
    m: torch.Tensor,
    m_scale: torch.Tensor,
    m_bias: torch.Tensor,
    parameters: dict,
    masks: dict,
    count_eps: float = 0.001,
    eps: float = 1e-05,
) -> torch.Tensor

Worked example

Count supporting sequences for each pair

With eps=3, LayerNorm gives [-0.5, 0.5]. The left projection is $2$ , the right projection is $3$ , so each supporting sequence contributes outer product $6$ . Pair (0, 0) has two supporting rows: $2\times6[2,-1]+[1,2]=[25,-10]$ . Divide by $2+0.5$ to obtain [10, -4]. The output bias is inside this division. A pair with no support returns zero.

import torch

m = torch.tensor([[[[-1.0, 1.0], [-1.0, 1.0], [-1.0, 1.0]], [[-1.0, 1.0], [-1.0, 1.0], [-1.0, 1.0]], [[-1.0, 1.0], [-1.0, 1.0], [-1.0, 1.0]]]], dtype=torch.float64, requires_grad=True)
m_scale = torch.tensor([1.0, 1.0], dtype=torch.float64, requires_grad=True)
m_bias = torch.tensor([0.0, 0.0], dtype=torch.float64, requires_grad=True)
parameters = {
    'w_left': torch.tensor([[1.0], [3.0]], dtype=torch.float64, requires_grad=True),
    'b_left': torch.tensor([1.0], dtype=torch.float64, requires_grad=True),
    'w_right': torch.tensor([[2.0], [4.0]], dtype=torch.float64, requires_grad=True),
    'b_right': torch.tensor([2.0], dtype=torch.float64, requires_grad=True),
    'w_out': torch.tensor([[[2.0, -1.0]]], dtype=torch.float64, requires_grad=True),
    'b_out': torch.tensor([1.0, 2.0], dtype=torch.float64, requires_grad=True),
}
masks = {
    'msa_mask': torch.tensor([[[True, True, False], [True, False, True], [False, False, True]]], dtype=torch.bool),
}
count_eps = 0.5
eps = 3.0

result = outer_product_mean(m=m, m_scale=m_scale, m_bias=m_bias, parameters=parameters, masks=masks, count_eps=count_eps, eps=eps)

Expected return

torch.tensor([[[[10.0, -4.0], [8.666666666666666, -2.6666666666666665], [8.666666666666666, -2.6666666666666665]], [[8.666666666666666, -2.6666666666666665], [8.666666666666666, -2.6666666666666665], [0.0, 0.0]], [[8.666666666666666, -2.6666666666666665], [0.0, 0.0], [10.0, -4.0]]]], dtype=torch.float64)

Parameters

Parameter	Shape and meaning
`m`torch.Tensor	$[B, S, N, \mathrm{Cm}]$ MSA representations to normalize independently at each sequence/residue position.
`m_scale`torch.Tensor	$[\mathrm{Cm}]$ Learned channel LayerNorm scale.
`m_bias`torch.Tensor	$[\mathrm{Cm}]$ Learned channel LayerNorm bias.
`parameters`dict	Supplied model parameters, keyed as shown below.
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`count_eps`floatDefault `0.001`	Positive offset added once to each residue pair’s joint valid-sequence count. Divide the projected numerator, including its bias, by count + count_eps; empty pairs return zero.
`eps`floatDefault `1e-05`	Positive value added to final-axis population variance inside each LayerNorm square root. It is not a mask/count denominator or an output offset.

Return

torch.Tensor · $[B, N, N, \mathrm{Cz}]$

Return [B,N,N,Cz] directed pair update; zero where joint support is empty.

Required behavior

Implement a masked outer-product mean with learned normalization and affine projections. $\operatorname{LN}(x; \mathrm{scale},\mathrm{bias}) = \mathrm{scale}\cdot (x-\mathrm{mean\_last}(x))/\sqrt{\mathrm{mean\_last}((x-\mathrm{mean\_last}(x))^{2})+\varepsilon}+\mathrm{bias}.$
Use population variance, over channels only; $\varepsilon$ is inside the square root. All dimension sizes are positive, all floating inputs are finite and share a dtype/device, and all learned tensors are supplied (no parameter initialization).
Set $x=\operatorname{LN}(m; m_{\mathrm{scale}},m_{\mathrm{bias}})$ , $a_{b,s,i,a}=\mathrm{msa\_mask}_{b,s,i}\left(\sum_c x_{b,s,i,c}\,w_{\mathrm{left},c,a}+b_{\mathrm{left},a}\right)$ , and $r_{b,s,j,d}=\mathrm{msa\_mask}_{b,s,j}\left(\sum_c x_{b,s,j,c}\,w_{\mathrm{right},c,d}+b_{\mathrm{right},d}\right)$ .
Sum $O_{b,i,j,a,d}=\sum_s a_{b,s,i,a}\,r_{b,s,j,d}$ , retaining BOTH projected channel axes. The joint count $c_{b,i,j}=\sum_s\mathrm{msa\_mask}_{b,s,i}\,\mathrm{msa\_mask}_{b,s,j}$ depends on the residue pair.
For $c_{b,i,j}>0$ : $\mathrm{update}_{b,i,j,o}=\frac{\sum_{a,d} O_{b,i,j,a,d}\,w_{\mathrm{out},a,d,o}+b_{\mathrm{out},o}}{c_{b,i,j}+\mathrm{count\_eps}}.$
For $c_{b,i,j}=0$ , return exactly zero, with finite zero gradients for that pair. The OUTPUT BIAS IS INSIDE the numerator and is divided by $\mathrm{count}+\mathrm{count\_eps}$ . Use additive count_eps, not $\operatorname{clamp}(\mathrm{count},1)$ , and do not symmetrize.
Flattening $(a,d)$ is row-major, with $d$ varying fastest, if implemented using a matrix projection. Return the update only, not a residual sum. Config: $\varepsilon$ (positive float, default $10^{-5}$ ) is the normalization $\epsilon$ ; count_eps (positive float, default $10^{-3}$ ) is added after the joint-count reduction.

Why this matters

Two independent feature projections let each aligned sequence contribute an outer product to a residue pair; the pair’s shared observations determine its denominator.

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 outer_product_mean(
    m: torch.Tensor,  # (B, S, N, Cm)
    m_scale: torch.Tensor,  # (Cm,)
    m_bias: torch.Tensor,  # (Cm,)
    parameters: dict,
    masks: dict,
    count_eps: float = 0.001,
    eps: float = 1e-05,
) -> torch.Tensor:
    """Return a tensor with shape (B, N, N, Cz)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement outer_product_mean")

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.

Sources and further reading

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