MSA and pair representations
Reference Atom Pair Features
Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.
Your task
Build a conditioning vector for each query atom and each supplied key slot. Combine the two atoms’ projected features, then add signed reference-coordinate and distance features only when they share a reference space. Return the local pair vectors before and after the residual MLP, plus the atom/key conditions and geometry-validity mask.
Function signature
af3_reference_atom_pair_features(
activations: dict,
parameters: dict,
masks: dict,
structure: dict,
) -> dict[str, torch.Tensor]Worked example
Distinguish same-space and unrelated atoms
Query atom 0 and key atom 1 share a reference-space ID. Their signed offset is [-2, 0, 0], and inverse squared-distance feature is 1/(1+4)=1/5. The supplied projections give [-2,0]+[2,3]/5+[5,7]=[3.4,7.6]. Query 0 and key 2 belong to different spaces, so their geometry and validity terms are omitted. Their row/column feature projections are zero in this example.
import torch
activations = {
'x': torch.full((1, 4, 3), 0.0, dtype=torch.float64, requires_grad=True),
}
parameters = {
'row_weight': torch.tensor([[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], dtype=torch.float64, requires_grad=True),
'col_weight': torch.tensor([[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], dtype=torch.float64, requires_grad=True),
'offset_weight': torch.tensor([[1.0, 0.0], [0.0, 1.0], [0.0, 0.0]], dtype=torch.float64, requires_grad=True),
'distance_weight': torch.tensor([[2.0, 3.0]], dtype=torch.float64, requires_grad=True),
'valid_weight': torch.tensor([[5.0, 7.0]], dtype=torch.float64, requires_grad=True),
'pair_mlp_weight': torch.full((3, 2, 2), 0.0, dtype=torch.float64, requires_grad=True),
}
masks = {
'mask': torch.tensor([[True, True, True, True]], dtype=torch.bool),
'window_mask': torch.full((1, 4, 3), True, dtype=torch.bool),
}
structure = {
'ref_positions': torch.tensor([[[0.0, 0.0, 0.0], [2.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 4.0, 0.0]]], dtype=torch.float64, requires_grad=True),
'ref_space_uid': torch.tensor([[0, 0, 1, 1]], dtype=torch.int64),
'key_index': torch.tensor([[[0, 2, 1], [1, 0, 0], [3, 2, -1], [2, 1, 3]]], dtype=torch.int64),
}
result = af3_reference_atom_pair_features(activations=activations, parameters=parameters, masks=masks, structure=structure)Expected return
{
'query_condition': torch.full((1, 4, 3), 0.0, dtype=torch.float64),
'key_condition': torch.full((1, 4, 3, 3), 0.0, dtype=torch.float64),
'atom_features': torch.full((1, 4, 3), 0.0, dtype=torch.float64),
'offset_mask': torch.tensor([[[True, False, True], [True, True, True], [True, True, False], [True, False, True]]], dtype=torch.bool),
'pair_features': torch.tensor([[[[7.0, 10.0], [0.0, 0.0], [3.4, 7.6]], [[7.0, 10.0], [7.4, 7.6], [7.4, 7.6]], [[5.2, 4.3], [7.0, 10.0], [0.0, 0.0]], [[5.2, 10.3], [0.0, 0.0], [7.0, 10.0]]]], dtype=torch.float64),
'output': torch.tensor([[[[7.0, 10.0], [0.0, 0.0], [3.4, 7.6]], [[7.0, 10.0], [7.4, 7.6], [7.4, 7.6]], [[5.2, 4.3], [7.0, 10.0], [0.0, 0.0]], [[5.2, 10.3], [0.0, 0.0], [7.0, 10.0]]]], dtype=torch.float64),
}Parameters
| Parameter | Shape and meaning |
|---|---|
activationsdict | Supplied intermediate tensors. |
parametersdict | Supplied model parameters, keyed as shown below. |
masksdict | Supplied masks. Each field below defines what True includes or permits. |
structuredict | Fields supplied together in this dictionary. |
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.
- Gather key_condition at each slot. Initial pair features are ReLU(query_condition) @ row_weight plus ReLU(key_condition) @ col_weight.
- Offsets point from key to query:
r_i-r_j. Add offset @ offset_weight plus1/(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 equals query_condition; this exercise does not use noisy coordinates or trunk features.
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 af3_reference_atom_pair_features(
activations: dict,
parameters: dict,
masks: dict,
structure: dict,
) -> 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_reference_atom_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.
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.