MSA and pair representations · INTEGRATION

Single-to-pair features and atom reference conditioning

HardRaw PyTorchAbout 30 min5 public examples plus private tests

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

Your task

Construct one feature vector describing each ordered token pair, or each mapped query-atom/key slot in the AF3 atom variant. Token variants combine projected single features through the stated products and differences, then apply their pair head. The AF3 atom variant combines atom conditions, reference geometry, and optional trunk/noise routes. Return the selected branch’s pair features, readout, and declared diagnostics.

Function signature

single_to_pair_features(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    options: dict,
) -> dict[str, torch.Tensor]

Worked example

An ordered ESM3 pair

In this ESM3 example, raw token projection gives q=[1,3] and k=[2,5]. Pair (0,1) is [k₀q₁,q₁−k₀]=[6,1]; pair (1,0) is [k₁q₀,q₀−k₁]=[5,−4]. The hidden projection, exact GELU, hidden LayerNorm, and final readout then operate on each pair independently.

import torch

activations = {
    'x': torch.tensor([[[1.0, 2.0], [3.0, 5.0]]], dtype=torch.float64, requires_grad=True),
}
parameters = {
    'weight': torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float64, requires_grad=True),
    'bias': torch.tensor([0.0, 0.0], dtype=torch.float64, requires_grad=True),
    'hidden_weight': torch.tensor([[0.21216724624478506, 0.06723158926988397, -0.22075659103655332, -0.26883329892386415], [0.2893647720587629, 0.050376568994800024, -0.27911394562654585, 0.4091670321049685]], dtype=torch.float64, requires_grad=True),
    'hidden_bias': torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64, requires_grad=True),
    'hidden_scale': torch.tensor([0.795208817925842, 0.020389689607759776, -0.0021728448198021777, 0.6797717713123919], dtype=torch.float64, requires_grad=True),
    'hidden_norm_bias': torch.tensor([0.4990308466741449, 0.412780913199289, -0.09772624102845953, -0.13636775593871922], dtype=torch.float64, requires_grad=True),
    'out_weight': torch.tensor([[0.11925094884659707, -0.2225214435612696, 0.19144161103199026], [-0.7885706318288417, -0.3458584542086859, 0.4039829309365077], [0.4198020222911164, -0.10846524434043135, 0.20913154379117166], [-0.13320172563627275, -0.9911054716945893, 0.17832603147645124]], dtype=torch.float64, requires_grad=True),
    'out_bias': torch.tensor([0.0, 0.0, 0.0], dtype=torch.float64, requires_grad=True),
}
masks = {
    'mask': torch.tensor([[True, True]], dtype=torch.bool),
}
structure = {
}
options = {
    'variant': 'esm3',
}

result = single_to_pair_features(activations=activations, parameters=parameters, masks=masks, structure=structure, options=options)

Expected return

{
    'pair_features': torch.tensor([[[[2.0, -1.0], [6.0, 1.0]], [[5.0, -4.0], [15.0, -2.0]]]], dtype=torch.float64),
    'output': torch.tensor([[[[-0.06920852641660866, 0.6486721557558881, 0.2181486014425758], [-0.0543721550180028, 0.14350769645932748, 0.3727625527042009]], [[-0.40388905713661477, 0.2904779634860536, 0.009239601425780924], [-0.05849757277360139, 0.12680340136716683, 0.37476460055425737]]]], dtype=torch.float64),
}

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.
`options`dict	Scalar settings. Omitted keys use the defaults below.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Construct a directed pair tensor from supplied single-token features, then apply the full variant's readout. $x$ [ $B,N,C$ ], mask[ $B,N$ ] ( $\mathrm{True}=\mathrm{valid}$ ) and floating projection parameters share $\mathrm{float32}/64$ . $B,N,C,H,D \ge 1$ .
Sanitize masked $x$ before arithmetic; invalid $i$ , $j$ pairs have exact-zero pair_features and output, including all biases.
variant='esmfold' (default): population LN over $C$ with supplied norm_scale/norm_bias and $\varepsilon(\mathrm{default} 10^{-5})>0$ , then $[q_{i},k_{i}]=\operatorname{LN}(x_{i})\mathbin{@}\mathrm{weight}+\mathrm{bias}$ , with weight $[C,2H]$ .
For matrix row $i$ ,column $j$ use $f_{\mathrm{ij}}=\operatorname{concat}(q_{j}\cdot k_{i}, q_{j}-k_{i})$ , width $2 H$ . $\mathrm{output\_ij}=f_{\mathrm{ij}}\mathbin{@}\mathrm{out\_weight}+\mathrm{out\_bias}$ , out_weight $[2H,D]$ . This is the original ESMFold orientation, not q_i-k_j and not a full $H-\mathrm{by}-H$ outer product.
variant='esmfold2': $u_{i}=x_{i}\mathbin{@}\mathrm{weight}+\mathrm{bias}$ with weight[ $C,H$ ]; $f_{\mathrm{ij}}=\operatorname{concat}(u_{i}\cdot u_{j},u_{i}-u_{j})$ . Apply $h=\operatorname{GELU}(f\mathbin{@}\mathrm{hidden\_weight}+\mathrm{hidden\_bias})$ , then $\mathrm{output}=h\mathbin{@}\mathrm{out\_weight}+\mathrm{out\_bias}$ .
hidden_weight $[2H,K]$ , out_weight[ $K,D$ ]; GELU is the exact erf definition $x\cdot (1+\operatorname{erf}(x/\sqrt{2}))/2$ , not the tanh approximation. This variant has no input LN inside this primitive; normalization before language-layer mixing is a separate exercise.
Optional output_norm (defaultFalse) applies population LN to the $D$ -channel readout using output_scale $[D]$ ,output_bias $[D]$ , before pair masking. ESMFold2's language-model encoder uses this option; it is optional teaching composition for esmfold.
These two variants return pair_features[ $B,N,N$ , $2 H$ ] and output[ $B,N,N,D$ ], without residual or symmetric averaging.
variant='esm3' implements the structural decoder's PairwisePredictionHead. It uses RAW supplied $x$ , with no input LayerNorm: $[q_{i},k_{i}]=x_{i}\mathbin{@}\mathrm{weight}+\mathrm{bias}$ , weight $[C,2H]$ , and $f_{\mathrm{ij}}=\operatorname{concat}(q_{j}\cdot k_{i},q_{j}-k_{i})$ , the same directed $q/k$ orientation as esmfold.
If include_pairwise=True (defaultFalse), append supplied pairwise_state[ $B,N,N,Z$ ] after the product/difference channels, with $Z\ge 1$ . Sanitize that state at invalid token pairs before concatenating; pair_features returns this complete concatenation, width $2H+Z$ (or $2 H$ without it).
For esm3, $u=\mathrm{pair\_features}\mathbin{@}\mathrm{hidden\_weight}+\mathrm{hidden\_bias}$ , hidden_weight[ $2 H$ + $Z,K$ ]; $h=\operatorname{LN}(\operatorname{GELU}(u);\mathrm{hidden\_scale},\mathrm{hidden\_norm\_bias})$ , with both normalization parameters $[K]$ .
Then $\mathrm{output}=h\mathbin{@}\mathrm{out\_weight}+\mathrm{out\_bias}$ , out_weight[ $K,D$ ]. GELU is exact, followed by affine population LayerNorm over $K$ before the final readout; there is no LN on the input $x$ or on output logits. output_norm does not apply to this variant.
The original head class permits biases; supplying zero projection biases reproduces the released structure decoder's bias-free builder. Upstream final-stack normalization, if present, has already happened in supplied $x$ and must not be applied a second time here. Invalid token pairs are zeroed in both outputs.
variant='af3_atom' has a different atom/window layout: $x$ [ $B,A,C$ ] is supplied categorical/numeric reference-atom conditioning, mask[ $B,A$ ] marks real atoms, key_index[ $B,A,K$ ] gives each query atom's candidate global key indices, and window_mask[ $B,A,K$ ] selects permitted slots. $A,K \ge 1$ .
A slot is legal only when both atoms are valid, its key index lies in $[0,A)$ , and its window mask is True; arbitrary out-of-range indices are ignored safely. This is a supplied layout map, not nearest-neighbor discovery. Option trunk=True (defaultFalse) adds token conditioning. Let LN0 use supplied scale, no bias, and $\varepsilon$ .
With atom_to_token[ $B,A$ ], token_mask[ $B,T$ ], trunk_single[ $B,T,S$ ], define $c_{i}=x_{i}+\operatorname{gather}(\operatorname{LN0}(\mathrm{trunk\_single})\mathbin{@}\mathrm{trunk\_single\_weight})$ . Missing/out-of-range/masked token mappings contribute zero.
Likewise normalize/project trunk_pair[ $B,T,T,Z$ ] and gather token-pair features at $(\mathrm{token\_i},\mathrm{token\_j})$ . With trunk=False, $c_{i}=x_{i}$ . Define $\mathrm{row}=\operatorname{ReLU}(c_{i})\mathbin{@}\mathrm{row\_weight}$ and $\mathrm{col}=\operatorname{ReLU}(c_{j})\mathbin{@}\mathrm{col\_weight}$ .
For a legal pair, $m_{\mathrm{ij}}=(\mathrm{ref\_space\_uid\_i}=\mathrm{ref\_space\_uid\_j})$ , $\delta=\mathrm{ref\_positions\_i}-\mathrm{ref\_positions\_j}$ .
Build $a_{\mathrm{ij}}=\mathrm{row}+\mathrm{col}+\mathrm{gathered\_trunk\_pair} + m_{\mathrm{ij}}\cdot (\delta\mathbin{@}\mathrm{offset\_weight} + [1/(1+\sum(\delta^{2}))]\mathbin{@}\mathrm{distance\_weight} + \mathrm{valid\_weight}[0])$ . The geometry is excluded across unrelated reference spaces, but row/column/trunk conditioning is retained.
Every projection is bias-free. Then $\mathrm{h1}=\operatorname{ReLU}(a)\mathbin{@}\mathrm{pair\_mlp\_weight}[0]$ , $\mathrm{h2}=\operatorname{ReLU}(\mathrm{h1})\mathbin{@}\mathrm{pair\_mlp\_weight}[1]$ , $\mathrm{h3}=\operatorname{ReLU}(\mathrm{h2})\mathbin{@}\mathrm{pair\_mlp\_weight}[2]$ , and $\mathrm{output}=a+\mathrm{h3}$ .
All three matrices are[ $P,P$ ]. ReLU's derivative at zero is defined as $0$ . pair_features returns a; illegal slots return zero in both tensors. The af3_atom result also returns $\mathrm{query\_condition}=c$ [ $B,A,C$ ], gathered legal key_condition[ $B,A,K,C$ ], offset_mask[ $B,A,K$ ], and atom_features[ $B,A,C$ ].
If noisy=True (defaultFalse), $\mathrm{atom\_features}=c+\mathrm{noisy\_positions}\mathbin{@}\mathrm{noisy\_weight}$ ; otherwise it is $c$ . Noisy positions are ALREADY supplied in the encoder's preconditioned coordinate scale, so this is not another EDM scaling or denoising step.
These outputs connect directly to windowed_atom_attention's condition/pair/atom inputs. Masked atom payloads and masked trunk payloads may be NaN and are sanitized before operations. All valid inputs and weights are finite, ref-position magnitudes are below $10^{10}$ , token axes are nonempty, and IDs fit int64.
for this variant is variant, $\varepsilon$ ,trunk,noisy (output_norm does not apply). Invalid variant or nonpositive $\varepsilon$ raises ValueError; inconsistent shapes and nonfinite valid inputs/parameters are outside the contract.
Required arithmetic intermediates, including centered squares, pair products and learned projections, must remain finite in the common dtype; $\varepsilon$ must remain strictly positive when represented in that dtype. No MSA averaging, atom attention, chemical parsing, learned weights, or coordinate generation is hidden in this primitive.

Why this matters

Pair representations describe relationships between two positions. Product channels capture shared feature strength, while difference channels preserve direction. The AF3 atom variant instead combines atom conditioning, reference geometry and optional trunk context.

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 single_to_pair_features(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    options: dict,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    pair_features: ESM variants: (B, N, N, 2H), with an extra Z channels for esm3 plus
        include_pairwise=True. Atom variant: (B, A, K, P).
    output: ESM variants: (B, N, N, D). Atom variant: (B, A, K, P).
    query_condition: (B, A, C). Returned only for variant="af3_atom".
    key_condition: (B, A, K, C). Returned only for variant="af3_atom".
    offset_mask: (B, A, K). Returned only for variant="af3_atom".
    atom_features: (B, A, C). Returned only for variant="af3_atom"."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement single_to_pair_features")

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.