MSA and pair representations

MSA row attention with pair bias

HardRaw PyTorchAbout 30 min2 public examples plus private tests

AlphaFold2AlphaFold-MultimerOpenFold

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

Your task

For each sequence and query residue, mix value features from allowed residues in that same sequence. Its weights combine query–key similarity with a directed pair bias shared across all sequences. Gate and project the mixture to return one MSA update vector at that sequence/residue.

Function signature

msa_row_pair_bias(
    activations: dict,
    parameters: dict,
    masks: dict,
    eps: float = 1e-05,
) -> torch.Tensor

Worked example

How the pair bias changes one residue’s mixture

Consider one valid query after the specified normalization and projections, with H=D=V=1. Its scaled QK scores over three residues are [1,0,-1], its pair biases are [0,log(2),100], and the key mask is [True,True,False]. The included exponential weights are e and 2, so attention is approximately [0.576117,0.423883,0]. For projected values [2,6,99], the mixture is (2e+12)/(e+2) ≈ 3.695532. A query gate of 1/2, scalar output weight 2 and output bias 1 produce update ≈ 4.695532. Repeat the same pair biases for another MSA sequence, but recompute its QK scores and values.

Parameters

Parameter	Shape and meaning
`activations`dict	Supplied intermediate tensors.
`parameters`dict	Supplied model parameters, keyed as shown below.
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`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, S, N, \mathrm{Cm}]$

Return [B,S,N,Cm] floating gated row-attention update; masked queries exactly zero.

Required behavior

Implement full gated MSA ROW attention with pair bias. $\operatorname{LN}(x;\,\mathrm{scale},\mathrm{bias})\,=\,\mathrm{scale}\cdot (x-\operatorname{mean}_{\mathrm{last}}(x))/\sqrt{\operatorname{mean}_{\mathrm{last}}((x-\operatorname{mean}_{\mathrm{last}}(x))^2)+\varepsilon}+\mathrm{bias}$ .
Use population variance, over channels only; eps 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;\,\mathrm{m\_scale},\mathrm{m\_bias})$ and $p=\operatorname{LN}(z;\,\mathrm{z\_scale},\mathrm{z\_bias})$ .
For each batch $b$ , sequence $s$ , head $h$ and residues $i,j$ : $q_{b,s,i,h,d}=\sum_{c}\,x_{b,s,i,c}\cdot \mathrm{w\_q}_{c,h,d}$ ; $k$ and $v$ use w_k/w_v. Pair bias $\beta_{b,h,i,j}=\sum_{c}\,p_{b,i,j,c}\cdot \mathrm{w\_pair}_{c,h}$ is shared across $s$ , not across batches.
Logits $L_{b,s,h,i,j}=(\sum_{d}\,q_{b,s,i,h,d}\cdot k_{b,s,j,h,d})/\sqrt{D}+\beta_{b,h,i,j}$ . Normalize $\exp(L)$ over VALID key residues $j$ only; use a max shift for stability. An all-invalid key row has exactly zero weights and no NaNs. Never mix sequence rows.
Mix values: $t_{b,s,i,h,v}=\sum_{j}\,\mathrm{attention}_{b,s,h,i,j}\cdot v_{b,s,j,h,v}$ . Gate $g_{b,s,i,h,v}=\sigma(\sum_{c}\,x_{b,s,i,c}\cdot \mathrm{w\_gate}_{c,h,v}+\mathrm{b\_gate}_{h,v})$ , with $\sigma(t)=1/(1+\exp(-t))$ .
Return $\mathrm{update}_{b,s,i,c}=\mathrm{msa\_mask}_{b,s,i}\cdot (\sum_{h,v}\,t_{b,s,i,h,v}\cdot g_{b,s,i,h,v}\cdot \mathrm{w\_out}_{h,v,c}+\mathrm{b\_out}_{c})$ . $Q$ / $K$ / $V$ and pair-bias projections have no affine bias. Final query masking occurs after output bias. Return only the update, not $m+\mathrm{update}$ .
Config: eps (positive float, default $1\times10^{-5}$ ); no other options.

Why this matters

The MSA row supplies sequence-specific queries, keys and values. The pair representation adds structural context to every row's attention logits without mixing information across different sequences.

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

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

google-deepmind/alphafold · modules.py

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