MSA and pair representations

Trunk-Conditioned Atom Pair Features

HardRaw PyTorchAbout 30 min2 public examples plus private tests

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

Your task

Route token-level trunk features back to atoms and local atom pairs. Add each atom’s mapped trunk-single feature to its query condition, and add the ordered trunk-pair feature to each mapped query/key slot before the reference-geometry terms and residual MLP. A separate noisy-coordinate projection changes only atom_features.

Function signature

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

Worked example

Both atoms map to the same token. Its trunk-single projection adds 1, giving query_condition = [2, 3]. Trunk-pair feedback adds 1 to each pair: pair_features = [[10, 10], [11, 13]]. Three half-scale MLP stages multiply positive pairs by $1/8$ before adding the residual. Noisy projections separately give atom_features = [2.5, 4].

import torch

activations = {
    'x': torch.tensor([[[1.0], [2.0]]], dtype=torch.float64, requires_grad=True),
    'trunk_single': torch.tensor([[[-1.0, 1.0]]], dtype=torch.float64, requires_grad=True),
    'trunk_pair': torch.tensor([[[[-1.0, 1.0]]]], dtype=torch.float64, requires_grad=True),
}
parameters = {
    'row_weight': torch.tensor([[1.0]], dtype=torch.float64, requires_grad=True),
    'col_weight': torch.tensor([[2.0]], dtype=torch.float64, requires_grad=True),
    'offset_weight': torch.tensor([[1.0], [0.0], [0.0]], dtype=torch.float64, requires_grad=True),
    'distance_weight': torch.tensor([[2.0]], dtype=torch.float64, requires_grad=True),
    'valid_weight': torch.tensor([[1.0]], dtype=torch.float64, requires_grad=True),
    'pair_mlp_weight': torch.tensor([[[0.5]], [[0.5]], [[0.5]]], dtype=torch.float64, requires_grad=True),
    'trunk_single_scale': torch.tensor([1.0, 1.0], dtype=torch.float64, requires_grad=True),
    'trunk_single_weight': torch.tensor([[-1.0], [1.0]], dtype=torch.float64, requires_grad=True),
    'trunk_pair_scale': torch.tensor([1.0, 1.0], dtype=torch.float64, requires_grad=True),
    'trunk_pair_weight': torch.tensor([[-1.0], [1.0]], dtype=torch.float64, requires_grad=True),
    'noisy_weight': torch.tensor([[0.5], [0.0], [0.0]], dtype=torch.float64, requires_grad=True),
}
masks = {
    'mask': torch.tensor([[True, True]], dtype=torch.bool),
    'window_mask': torch.tensor([[[True, True], [True, True]]], dtype=torch.bool),
    'token_mask': torch.tensor([[True]], dtype=torch.bool),
}
structure = {
    'ref_positions': torch.tensor([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]], dtype=torch.float64, requires_grad=True),
    'ref_space_uid': torch.tensor([[0, 0]], dtype=torch.int64),
    'key_index': torch.tensor([[[0, 1], [0, 1]]], dtype=torch.int64),
    'atom_to_token': torch.tensor([[0, 0]], dtype=torch.int64),
    'noisy_positions': torch.tensor([[[1.0, 0.0, 0.0], [2.0, 0.0, 0.0]]], dtype=torch.float64, requires_grad=True),
}
eps = 3.0

result = af3_trunk_atom_pair_conditioning(activations=activations, parameters=parameters, masks=masks, structure=structure, eps=eps)

Expected return

{
    'query_condition': torch.tensor([[[2.0], [3.0]]], dtype=torch.float64),
    'key_condition': torch.tensor([[[[2.0], [3.0]], [[2.0], [3.0]]]], dtype=torch.float64),
    'atom_features': torch.tensor([[[2.5], [4.0]]], dtype=torch.float64),
    'offset_mask': torch.tensor([[[True, True], [True, True]]], dtype=torch.bool),
    'pair_features': torch.tensor([[[[10.0], [10.0]], [[11.0], [13.0]]]], dtype=torch.float64),
    'output': torch.tensor([[[[11.25], [11.25]], [[12.375], [14.625]]]], 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.
`eps`floatDefault `1e-05`	Positive stabilizer inside population-variance square roots for the scale-only trunk-single and trunk-pair LayerNorm operations.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

x is $[B,A,C]$ ; mask marks valid atoms. key_index $[B,A,K]$ chooses atom slots per query, and window_mask marks allowed slots. A slot requires valid query/key atoms, an in-range index, and window_mask=True; repeated keys stay repeated.
Begin with cleaned x as query_condition. Layer-normalize valid trunk_single with its supplied scale (no bias), project it, and add the feature gathered by a valid atom_to_token/token_mask mapping.
Gather key_condition at each slot. Initial pair features are ReLU(query_condition) @ row_weight plus ReLU(key_condition) @ col_weight.
Normalize valid trunk_pair with its supplied scale (no bias), project it, and gather the ordered query-token/key-token pair into each atom slot. Add it only when both atom-to-token mappings are valid. Normalization uses eps (default 1e-5).
Offsets point from key to query: $r_i-r_j$ . Add offset @ offset_weight plus $1/(1+\|offset\|^2)$ @ distance_weight only when the allowed atoms share ref_space_uid. Also add offset_mask (as a scalar feature) @ valid_weight.
Pass pair_features through three consecutive ReLU-then-matrix-product stages using pair_mlp_weight[0], [1], [2], then add this result to pair_features as a residual. No stage has a bias.
Return zero key_condition/pair_features/output for disallowed slots, zero atom_features for invalid atoms, and boolean offset_mask. query_condition is the cleaned atom condition.
atom_features adds noisy_positions @ noisy_weight to query_condition after cleaning invalid noisy coordinates. This addition does not change query_condition, key_condition, or pair features.
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.

Why this matters

AlphaFold3 combines atom-local reference geometry with token-level trunk information. Atom-to-token maps route the trunk features into each local atom pair; noisy coordinates change atom features separately from these pair 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 af3_trunk_atom_pair_conditioning(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    eps: float = 1e-05,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    pair_features: (B,A,K,P).
    output: (B,A,K,P).
    query_condition: (B, A, C).
    key_condition: (B, A, K, C).
    offset_mask: (B, A, K).
    atom_features: (B, A, C)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement af3_trunk_atom_pair_conditioning")

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.