MSA and pair representations

Global-column MSA attention with shared keys and values

HardRaw PyTorchAbout 30 min3 public examples plus private tests

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

Your task

Compute one shared attention summary for each MSA column, then let every sequence gate that summary separately. Return the per-sequence updates and shared sequence-attention weights.

Function signature

msa_column_global_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
    eps: float = 1e-05,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

Shared summary, different row gates

Consider one residue column and one head after normalization and projection, with head width D=1. Let its shared query be 1, sequence keys [0,log(3)], and shared projected values [2,6]; both rows are valid. Sequence attention is [1/4,3/4] and the shared value summary is (1/4)·2+(3/4)·6=5. Suppose the two rows’ projected gate logits are [0,log(3)], giving gates [1/2,3/4]. The gated values are 2.5 and 3.75. With w_out=[[2,4]] and b_out=[1,-1], the row updates are [6,9] and [8.5,14]. Both rows read the same attention summary, but their own gates change the output.

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

tuple[torch.Tensor, torch.Tensor]

Return these tensors in the stated order.

Required behavior

Return a gated global-column MSA update and its shared attention weights. All $B,S,N,C,H,D,V$ positive and all floats finite. m $[B,S,N,C]$ uses alignment row $S$ then residue $N$ ; msa_mask $[B,S,N]$ True means usable.
Zero masked $m$ before LayerNorm, normalize the last $C$ channels with population variance and $\varepsilon>0$ default $1\times10^{-5}$ inside $\sqrt{\mathrm{var}+\varepsilon}$ , using norm_scale/norm_bias $[C]$ , then zero masked normalized vectors again.
For each $b,n$ , average normalized $x$ over valid alignment rows: $\mathrm{mean}=\sum_{s}\,x/\operatorname{count}$ , exactly zero when $\operatorname{count}=0$ . $Q[b,n,h,d]=\mathrm{mean}\,\mathrm{w\_q}[C,H,D]$ .
$K[b,s,n,d]=x\,\mathrm{w\_k}[C,D]$ , $V[b,s,n,v]=x\,\mathrm{w\_v}[C,V]$ have no HEAD AXIS: the keys/values are shared by all heads. $\mathrm{Logits}[b,n,h,s]=\operatorname{dot}(Q,K)/\sqrt{D}$ . Softmax stably over valid alignment rows $S$ only, all-masked rows return zero weights $[B,N,H,S]$ .
$A[b,n,h,v]=\sum_{s}\,\mathrm{weights}\cdot V$ is one shared output per residue/head, not a separate output query for each row. Each row has gate $G[b,s,n,h,v]=\sigma(x\,\mathrm{w\_gate}[C,H,V]+\mathrm{b\_gate}[H,V])$ .
Multiply $A$ (broadcast over $S$ ) by $G$ , flatten $H,V$ in head-major order, and project with $\mathrm{w\_out}[H\cdot V,C]+\mathrm{b\_out}[C]$ . Return update $[B,S,N,C]$ , forced to exact zero on masked rows after bias. No residual addition or dropout; query/key/value projections have no biases.
An all-valid mask reproduces the global-mean/shared-KV algebra in the original RFdiffusion block. The explicit empty-row and ragged-mask policy is a teaching extension to finite-mask-bias or unpadded source implementations.

Why this matters

Instead of producing a separate query for every sequence, this operator first averages the valid normalized rows. That pooled query selects a column summary; row-specific gates decide how much of it each sequence receives.

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_column_global_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
    eps: float = 1e-05,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    update: (B, S, N, C).
    weights: (B, N, H, S)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement msa_column_global_attention")

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.