MSA and pair representations

Sequence-Weighted MSA Attention

HardRaw PyTorchAbout 30 min1 public example plus private tests

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

Your task

Weight MSA sequences using target row zero, then compute a residue-attention update shared across sequences. Apply the supplied dropout only to sequence weights.

Function signature

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

Worked example

Sequence dropout changes contributions, not the key axis

At one residue and head, suppose the sequence-weight logits for two sequences are [0,log(3)]. Their weights are [1/4,3/4]. With sequence_keep=[True,False] and sequence_dropout=0.5, the retained weights become [1/2,0]; do not renormalize them to sum to one. For head width D=1, if the two sequence contributions q·k to one query-key score are [2,6] and its pair bias is zero, that score is (1/2)·2 + 0·6 = 1. If the other key has score 0, key attention is approximately [0.731059,0.268941]. This key attention is then shared across MSA sequences, each mixing its own 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`	Numerical stabilizer used at the locations specified in the task.
`sequence_dropout`floatDefault `0.0`	Probability p in [0,1). Multiply sequence weights by the supplied keep mask and divide by 1-p; this mask is dropout, not padding.

Return

torch.Tensor · $[B,S,L,C]$

Return Attention update, without adding an input residual.

Required behavior

All batch, sequence, residue, head and channel dimensions are positive. There is no padding mask: every supplied residue/sequence participates; supplied keep masks represent dropout retention only.
Normalize features and bias_features independently over their channel axes using affine population LayerNorm with eps (default 1e-5) inside the square root. Form bias-free Q/K/V and a sigmoid value gate from normalized features.
Pair bias is the normalized bias feature projected to heads. Normalize the final attention logits with softmax over key residue $j$ .
Return a projected attention update only, including output bias. Do not add the input residual.
Use CPU float32/float64 tensors with one shared floating dtype/device. Preserve strided layouts and do not mutate supplied tensors.
Use raw tensor operations, not a ready-made implementation of the requested primitive. Supplied parameters and random variates must not be reinitialized.
Compute target-row sequence queries $Q_{ihd}$ and per-sequence keys $K_{sihd}$ with the supplied affine projections of normalized MSA features. Set $a_{sih}=\operatorname{softmax}_s[\sum_d Q_{ihd}K_{sihd}/\sqrt D]$ .
Use $r=a\odot\mathrm{sequence\_keep}/(1-p)$ , with $0\le p<1$ , default 0. Then $A_{ijh}=\operatorname{softmax}_j[\sum_{s,d}r_{sih}q_{sihd}k_{sjhd}/\sqrt D+\beta_{ijh}]$ .
Aggregate $U_{sihd}=\sum_jA_{ijh}v_{sjhd}$ , multiply the matching sigmoid gate, and output-project. All-false sequence_keep leaves bias-only attention, not necessarily zero output.

Why this matters

At each residue, the target MSA row decides which sequences contribute most strongly to the shared residue-attention map. Each sequence then uses that map to aggregate its own value features.

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

Related exercises

Sources and further reading

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