Structure module, templates and recycling

AlphaFold 2 Recycling Update

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

Prepare the next AlphaFold 2 pass by adding feedback from the previous prediction to fresh base features. Reuse target-row single features, pair features and representative-atom distance feedback only when a cache exists and passes remain. Return the prepared features and explicit next-iteration/continuation state.

Function signature

af2_recycling_update(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    max_passes: int = 3,
    eps: float = 1e-05,
    glycine_index: int = 7,
) -> dict[str, torch.Tensor]

Worked example

Only batch item 0 has cached feedback. The learned updates are [-1, 1] for single and [-2, 2] for pair. The diagonal distance is zero, so geometry adds [1.5, 2.5]. Its single and pair values become [9, 21] and [1.5, 8.5]. Item 1 keeps its base states. next_iteration is [2, 1].

import torch

activations = {
    'single_base': torch.tensor([[[10.0, 20.0]], [[30.0, 40.0]]], dtype=torch.float64, requires_grad=True),
    'pair_base': torch.tensor([[[[2.0, 4.0]]], [[[6.0, 8.0]]]], dtype=torch.float64, requires_grad=True),
    'previous_single': torch.tensor([[[1.0, 3.0]], [[5.0, 7.0]]], dtype=torch.float64, requires_grad=True),
    'previous_pair': torch.tensor([[[[1.0, 3.0]]], [[[5.0, 7.0]]]], dtype=torch.float64, requires_grad=True),
    'distance_edges': torch.tensor([2.0], dtype=torch.float64, requires_grad=True),
}
parameters = {
    'single_scale': torch.tensor([2.0, 2.0], dtype=torch.float64, requires_grad=True),
    'single_bias': torch.tensor([0.0, 0.0], dtype=torch.float64, requires_grad=True),
    'pair_scale': torch.tensor([2.0, 2.0], dtype=torch.float64, requires_grad=True),
    'pair_bias': torch.tensor([0.0, 0.0], dtype=torch.float64, requires_grad=True),
    'w_single': torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float64, requires_grad=True),
    'w_pair': torch.tensor([[2.0, 0.0], [0.0, 2.0]], dtype=torch.float64, requires_grad=True),
    'w_distance': torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float64, requires_grad=True),
    'b_distance': torch.tensor([0.5, 0.5], dtype=torch.float64, requires_grad=True),
}
masks = {
    'has_previous': torch.tensor([True, False], dtype=torch.bool),
    'residue_mask': torch.tensor([[True], [True]], dtype=torch.bool),
    'atom_mask': torch.full((2, 1, 4), True, dtype=torch.bool),
}
structure = {
    'iteration': torch.tensor([1, 0], dtype=torch.int64),
    'atom_positions': torch.tensor([[[[0.0, 1.0, 0.0], [0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]], [[[0.0, 1.0, 0.0], [0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]]], dtype=torch.float64, requires_grad=True),
    'residue_type': torch.tensor([[0], [0]], dtype=torch.int64),
}
max_passes = 3
eps = 3.0
glycine_index = 7

result = af2_recycling_update(activations=activations, parameters=parameters, masks=masks, structure=structure, max_passes=max_passes, eps=eps, glycine_index=glycine_index)

Expected return

{
    'pseudo_beta': torch.tensor([[[0.0, 0.0, 1.0]], [[0.0, 0.0, 0.0]]], dtype=torch.float64),
    'distance_bins': torch.tensor([[[0]], [[-1]]], dtype=torch.int64),
    'single': torch.tensor([[[9.0, 21.0]], [[30.0, 40.0]]], dtype=torch.float64),
    'pair': torch.tensor([[[[1.5, 8.5]]], [[[6.0, 8.0]]]], dtype=torch.float64),
    'next_iteration': torch.tensor([2, 1], dtype=torch.int64),
    'continue': torch.tensor([True, True], dtype=torch.bool),
}

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.
`max_passes`intDefault `3`	Maximum recycling passes.
`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.
`glycine_index`intDefault `7`	Vocabulary index for glycine.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Feedback is active where has_previous is True and iteration < max_passes (default 3). Increment iteration wherever iteration < max_passes, even when no previous prediction is present; continue is next_iteration < max_passes.
Sanitize previous single/pair states with residue_mask, apply supplied affine population LayerNorm (eps default 1e-5 inside square root), and project using w_single/w_pair. Add feedback only for active batch entries to the freshly supplied base states, then mask absent residues/pairs.
Use CA as pseudo-beta for glycine (residue_type == glycine_index, default 7), and supplied CB for other residues. The selected atom mask determines pseudo-beta validity.
Distance bin index is the number of squared distance_edges strictly below the squared pseudo-beta separation. Embed the resulting one-hot bins with w_distance plus b_distance, mask invalid residue pairs, and add this geometry contribution only as active feedback.
Return zero pseudo_beta unless its selected atoms, residue, and feedback are valid. Return distance_bins = −1 unless both endpoints have valid pseudo-beta coordinates; otherwise return the strict-threshold bin index.
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

Recycling feeds information from an earlier prediction into fresh base representations; it does not itself rerun the network.

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 af2_recycling_update(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    max_passes: int = 3,
    eps: float = 1e-05,
    glycine_index: int = 7,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    single: (B, N, Cs).
    pair: (B, N, N, Cz).
    pseudo_beta: (B, N, 3).
    distance_bins: (B, N, N).
    next_iteration: (B,).
    continue: (B,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement af2_recycling_update")

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.