MSA and pair representations

MSA column attention

MediumRaw PyTorchAbout 20 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

Update each MSA position by attending across sequences at that residue. Keep residue columns independent and return the gated, projected MSA update.

Function signature

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

Worked example

Mix sequences at one fixed residue

At one valid query sequence and one fixed residue, suppose the normalized/projected one-head tensors have D=V=1, query q=1, key scalars [0,log(3),100], value scalars [2,6,99], and the column’s key mask [True,True,False]. The first two attention weights are 1/4 and 3/4; the masked sequence gets zero. Their value mixture is 2/4+18/4=5. With gate=1/2, scalar output weight=2 and bias=1, this query’s update is 6. This computation mixes sequences at this residue; features at other residue columns do not enter.

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 column-attention update, zero at invalid queries.

Required behavior

Implement full gated MSA COLUMN attention across $S$ . $\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). Let $x=\operatorname{LN}(m;\,\mathrm{m\_scale},\mathrm{m\_bias})$ .
Project $q$ and $k$ to $[B,S,N,H,D]$ and $v$ to $[B,S,N,H,V]$ using bias-free w_q, w_k, w_v; $H,D,V$ are inferred from those tensors. At fixed batch $b$ , residue $n$ , head $h$ and query sequence $s$ , logits over key sequence $t$ are $L_{b,n,h,s,t}=(\sum_{d}\,q_{b,s,n,h,d}\cdot k_{b,t,n,h,d})/\sqrt{D}$ .
Apply a stable masked softmax over $t$ , using $\mathrm{msa\_mask}_{b,t,n}=\mathrm{True}$ for valid keys. All-invalid columns have exactly zero attention weights; no residue-to-residue mixing occurs. $u_{b,s,n,h,v}=\sum_{t}\,\mathrm{attention}_{b,n,h,s,t}\cdot v_{b,t,n,h,v}$ .
$g_{b,s,n,h,v}=\sigma(\sum_{c}\,x_{b,s,n,c}\cdot \mathrm{w\_gate}_{c,h,v}+\mathrm{b\_gate}_{h,v})$ . Return $\mathrm{update}_{b,s,n,c}=\mathrm{msa\_mask}_{b,s,n}\cdot (\sum_{h,v}\,u_{b,s,n,h,v}\cdot g_{b,s,n,h,v}\cdot \mathrm{w\_out}_{h,v,c}+\mathrm{b\_out}_{c})$ . The output affine bias precedes final query masking.
Return only the update, without residual or dropout. Config: eps (positive float, default $1\times10^{-5}$ ). This is full attention, not pooled-query column-global attention.

Why this matters

An MSA column compares alternative sequences at one aligned position. Each query sequence chooses its own mixture of the other sequence representations; residue positions never mix in this operation.

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_attention(
    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_column_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

google-deepmind/alphafold · modules.py

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