MSA and pair representations

RoseTTAFold Pair Block

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

Use Cα geometry to condition communication between residue pairs, then update the pair state in three stages: a tied row-attention residual, a tied column-attention residual, and a transition residual. Each returned pair vector contains its starting features plus these successive learned updates.

Function signature

rosettafold_pair_block(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    eps: float = 1e-05,
    pair_dropout: float = 0.0,
    transition_dropout: float = 0.0,
) -> torch.Tensor

Worked example

How geometry reaches attention

For the distance-feature stage, take two Cα coordinates (0,0,0) and (20/35,0,0). Their off-diagonal distance is $d=20/35$ , exactly the center of radial channel $m=1$ , so that channel equals 1. Radial channel $m=0$ instead equals

\rho_0=\exp\!\left[-\left(\frac{d}{20/36}\right)^2\right]=\exp\!\left[-\left(\frac{36}{35}\right)^2\right]\approx0.347164.

On a diagonal pair the distance is zero and channel $m=0$ equals 1. Thus different residue pairs supply different radial vectors to the geometry MLP; its output becomes the attention bias in both residual attention stages.

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.
`structure`dict	Fields supplied together in this dictionary.
`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.
`pair_dropout`floatDefault `0.0`	Dropout probability for supplied pair-update keep masks.
`transition_dropout`floatDefault `0.0`	Dropout probability for supplied transition keep masks.

Return

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

Return Updated pair state including both attention residuals and transition.

Required behavior

features has shape $[B,L,L,C]$ ; ca_coordinates has shape $[B,L,3]$ . All axes are positive; no biological padding mask or random draw is used.
Compute exact Euclidean Cα pair distances, with zero self-distance. Use 36 Gaussian channels $\exp(-((d-20m/35)/(20/36))^2)$ , $m=0,\ldots,35$ , with no cutoff or envelope.
Project these channels through the supplied ReLU radial MLP to obtain geometry bias Y. Use the separately supplied row_ and column_ parameters for the tied attention cores.
Compute $Z_1=X+\mathrm{row}(X,Y)\odot\mathrm{row\_keep}/(1-p)$ , then $Z_2=Z_1+\mathrm{column}(Z_1,Y)\odot\mathrm{column\_keep}/(1-p)$ . Row masks broadcast over the first pair axis; column masks over the second. Attention biases are inside dropout.
Layer-normalize Z2, project C→2C with bias, apply ReLU and supplied hidden dropout, then project 2C→C and add the final bias AFTER dropout. Add this transition to Z2.
eps defaults to 1e-5 inside each population-variance square root. pair_dropout/transition_dropout are finite in [0,1), default 0; use supplied all-true masks for evaluation. Return complete residual features, not just an update.
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.
For each tied core, normalize features and geometry bias separately. Form q,k,v by their supplied bias-free projections. With S=L, form A_ijh = softmax_j(sum_s,d(q_sihdk_sjhd)/sqrt(SD) + projected_bias_ijh). Mix values as U_ikhd=sum_j A_ijh*v_kjhd, multiply by sigmoid(x_ik @ w_gate + b_gate), project over H,D and add b_out. The row core transposes both pair axes before this computation and transposes the update back; the column core uses the original orientation.

Why this matters

Tied attention builds one residue-to-residue map by combining evidence across an additional axis. The pair variants have a deliberately unusual value layout; the full block then adds these updates through supplied dropout masks and a transition.

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

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.