MSA and pair representations

AlphaFold 3 Relative Position Buckets

MediumRaw PyTorchAbout 20 min1 public example plus private tests

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

Your task

Describe each ordered token pair with residue-offset and token-offset categories, a same-entity indicator, and a relative symmetry-copy category. Project their concatenation to obtain one learned pair embedding. Token offsets are meaningful only within the same chain and residue; AF3 symmetry offsets are meaningful when the entities match.

Function signature

af3_relative_position_buckets(
    residue_index: torch.Tensor,
    mask: torch.Tensor,
    asym_id: torch.Tensor,
    token_index: torch.Tensor,
    entity_id: torch.Tensor,
    sym_id: torch.Tensor,
    weight: torch.Tensor,
    residue_radius: int = 32,
    chain_radius: int = 2,
) -> dict[str, torch.Tensor]

Worked example

Two tokens from one residue

Take two valid tokens with residue_index [10,10], token_index [0,2], asym_id [0,0], entity_id [7,7], and sym_id [0,0]. Set residue_radius=1 and chain_radius=1. Pair (0,1) gets residue/token/symmetry indices [1,0,1]: zero residue difference, token difference −2 clipped to the lowest bucket, and zero symmetry difference. Its features are [0,1,0,0 | 1,0,0,0 | 1 | 0,1,0,0], where bars separate the four blocks. With scalar projection weight[k]=k for k=0…12, the embedding is 1+4+8+10=23. The reverse pair has indices [1,2,1] and embedding 25.

Parameters

Parameter	Shape and meaning
`residue_index`torch.Tensor	$[B, N]$ Supplied residue numbering, not required to start at zero.
`mask`torch.Tensor	$[B, N]$ True for valid tokens.
`asym_id`torch.Tensor	$[B, N]$ Chain identity per token; equal IDs mean the same chain.
`token_index`torch.Tensor	$[B, N]$ Token position within the supplied residue, used when the two tokens share chain and residue identity.
`entity_id`torch.Tensor	$[B, N]$ Chemical/entity identity; equal values permit a relative symmetry-copy bucket.
`sym_id`torch.Tensor	$[B, N]$ Index of the symmetry-equivalent chain copy.
`weight`torch.Tensor	$[4r+2c+7,C]$ Learned projection from the ordered one-hot feature blocks to output channels.
`residue_radius`intDefault `32`	Clipping radius for residue-index differences.
`chain_radius`intDefault `2`	Clipping radius for symmetry-copy differences.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

residue_index and mask have shape $[B,N]$ , with True marking a valid token. residue_radius r (default 32) and chain_radius c (default 2, when used) are nonnegative integers. Weights project the concatenated bucket features into embedding channels.
Same-chain pairs have equal asym_id. Their residue bucket is clamp(residue_index[i]−residue_index[j]+r,0,2r); other chains use sentinel 2r+1.
Invalid endpoint pairs have indices -1, all-zero features, and all-zero embeddings.
The token bucket uses token_index[i]−token_index[j] clipped/shifted by r only for same-chain pairs at the same residue_index. Every other pair uses sentinel 2r+1.
Append same_entity as one scalar indicator. The symmetry bucket uses sym_id[i]−sym_id[j] clipped/shifted by c when entity_id matches; otherwise use sentinel 2c+1.
Feature order is residue one-hot(2r+2), token one-hot(2r+2), same_entity(1), symmetry one-hot(2c+2). Return the three corresponding integer indices in residue/token/symmetry order.

Why this matters

Relative offsets become discrete feature channels, but source models differ in signs, special buckets and which identity comparisons enable each feature.

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_relative_position_buckets(
    residue_index: torch.Tensor,  # (B, N)
    mask: torch.Tensor,  # (B, N)
    asym_id: torch.Tensor,  # (B, N)
    token_index: torch.Tensor,  # (B, N)
    entity_id: torch.Tensor,  # (B, N)
    sym_id: torch.Tensor,  # (B, N)
    weight: torch.Tensor,  # (4r+2c+7,C)
    residue_radius: int = 32,
    chain_radius: int = 2,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    indices: (B,N,N,3).
    features: (B,N,N,4r+2c+7).
    embedding: (B, N, N, C)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement af3_relative_position_buckets")

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.