MSA and pair representations

MSA pair-weighted averaging

MediumRaw PyTorchAbout 20 min3 public examples plus private tests

AlphaFold3OpenFold3Boltz-1Boltz-2ProtenixESMFold2

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

Your task

Use pair features to choose residue weights shared across MSA rows. Combine each sequence’s projected values with those weights, gate the result, and return the MSA update.

Function signature

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

Worked example

Missing MSA values do not change the shared weights

Inspect one valid query and one head after normalization and projection. Pair logits over two residues are [0,log(3)], with both pair_mask entries True, so the returned shared weights are [1/4,3/4]. Row A has projected values [2,6], giving mixture 5. Row B has the same values but its second residue is absent under msa_mask, giving mixture (1/4)×2+(3/4)×0=0.5. Do not renormalize Row B to [1,0]. With query gate=1/2, scalar output weight=2 and bias=1, the two valid-query updates are 6 and 1.5.

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

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Compute an AF3-style gated MSA UPDATE and expose its pair-only weights. All dimensions $B,S,N,C_m,C_z,H,D$ are positive; finite floating inputs share dtype/device. True means valid in both Boolean masks. No query or key projection of $m$ exists in this exercise.
First independently layer-normalize each $m$ $[b,s,i,:]$ and $z$ $[b,i,j,:]$ over their last feature axis using their own learned scale/shift. Use population variance and denominator $\sqrt{\mathrm{variance}+\varepsilon}$ , calling the results mhat and zhat.
Pair logits $L[b,h,i,j]=\sum_{c}\,\widehat{z}[b,i,j,c]\cdot \mathrm{w\_pair}[c,h]$ . Stable softmax over key token $j$ , restricted solely by pair_mask $[b,i,j]$ , gives weights $[B,H,N,N]$ . These weights have no $S$ axis and no MSA-derived QK term. An all-False pair row gives all-zero weights; there is no softmax epsilon.
Project values $V[b,s,j,h,d]=\sum_{c}\,\widehat{m}[b,s,j,c]\cdot \mathrm{w\_value}[c,h,d]$ , then zero $V$ where msa_mask $[b,s,j]$ is False. $A[b,s,i,h,d]=\sum_{j}\,\mathrm{weights}[b,h,i,j]\cdot V[b,s,j,h,d]$ .
Do not re-normalize weights for each MSA row: an absent row-specific value contributes zero while the shared denominator remains unchanged. Gate $G=\sigma(\widehat{m}\,\mathrm{w\_gate}+\mathrm{b\_gate})$ has shape $[B,S,N,H,D]$ and multiplies $A$ after aggregation.
Flatten ( $H,D$ ), then compute $\mathrm{update}=\operatorname{flatten}(G\cdot A)\,\mathrm{w\_out}+\mathrm{b\_out}$ , with no residual addition. Force update $[b,s,i]$ to exact zero after the output bias if msa_mask $[b,s,i]$ is False OR there is no $j$ with pair_mask $[b,i,j]$ AND msa_mask $[b,s,j]$ .
Returned weights depend only on pair features, their normalization/projection, and pair_mask. msa_mask affects values and final query eligibility, not the pair-only softmax denominator. eps is positive, defaults to 1e-5, and is used inside both normalization square roots.
Projections w_pair and w_value are bias-free; gate/output biases are supplied. No dropout or hidden weights.

Why this matters

Here the pair representation decides which residues matter. Each MSA row supplies its own values, but it does not get to change the pair-derived attention distribution when some values are missing.

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_pair_weighted_averaging(
    activations: dict,
    parameters: dict,
    masks: dict,
    eps: float = 1e-05,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    update: (B, S, N, Cm).
    weights: (B, H, N, N). (B,H,N,N) pair-only probabilities, shared unchanged across S"""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement msa_pair_weighted_averaging")

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.