All-atom folding

AlphaFold 3 Atom Attention Decoder

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

Add each atom’s mapped token context, update the atom using its allowed key slots and conditioned transition, then read out a 3D coordinate update from its hidden feature vector. Each position_update is a displacement output, not a final atom position. Use AF3 adaptive normalization and sigmoid residual gates.

Function signature

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

Worked example

Worked example: competing neighbor slots, an excluded slot, and an active residual update

Use AF3 decoder mode with two valid atoms [1,3], [4,2] mapped to one valid token. After adding token context $[1,1]$ , then adaptive normalization with eps=3, the query/key inputs are $[-0.25,0.25]$ and $[0.25,-0.25]$ . Query 0 can use both key slots; scores are $\pm0.125/\sqrt2$ , giving first-key weight $p=\sigma(0.25/\sqrt2)\approx0.544079$ . The weighted value is $[-0.25(2p-1),0.25(2p-1)]$ ; half value and residual gates add $[-d,d]$ , where $d=(2p-1)/16$ . Query 1 can use only key 1, so its update is $[1/16,-1/16]$ . The first hidden vector is $[2-d,4+d]$ . Its centered entries are $[-(1+d),1+d]$ ; final normalization, scale 2, and readout give $[-a,a,-2a]$ , $a=2(1+d)/\sqrt{(1+d)^2+3}$ . Here $\sigma(z)=1/(1+e^{-z})$ . The supplied transition is zero.

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 every LayerNorm variance or RMS mean-square root used by this block.
`radius`intDefault `-1`	Maximum absolute difference between query and key atom indices; -1 disables this additional index-window restriction.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

An atom is usable only when valid=True, atom_to_token lies in range, and its mapped token_valid=True. Sanitize unusable atom/condition payloads. A key slot additionally requires an in-range usable key and window_mask=True.
radius defaults to −1 (unlimited). A nonnegative radius further requires abs(query_atom_index−key_atom_index) ≤ radius. This is an INDEX window, not a geometric radius. Repeated slots remain independent softmax entries.
Begin with the cleaned atom features plus the projected token feature token @ w_token gathered to each usable atom.
For normalization stage j, let $c=LN(condition)\odot cond\_scale[j]$ . Adaptive features are $LN(x)\odot sigmoid(cW_{scale}[j]+b_{scale}[j])+cW_{shift}[j]$ . Stages0,1,2 are query, key/value, and transition respectively.
Project Q from query-normalized features with b_q; project K/V from key-normalized features without biases. Add pair @ w_pair to attention logits AFTER normalizing pair with LN and pair_scale. Pair bias has one value per head.
Attention and transition residual gates are sigmoid(raw condition @ w_residual_gate + b_residual_gate), using their respective stage indices.
Compute $Q\cdot K/\sqrt D$ plus any declared pair bias; softmax over allowed key slots only. Empty rows have exact zero attention weights. Weight V, multiply each head by sigmoid(query_features @ w_value_gate), flatten head-major, and apply w_attention_out before the gated residual.
Normalize the updated residual state for the transition. Compute SiLU(z @ w_ff_gate) * (z @ w_ff_up), then w_ff_out, then its gated residual. All normalization epsilons use eps (default 1e-5) inside the square root.
Return hidden, attention, and position_update. Apply LayerNorm to hidden, multiply readout_scale, then w_readout. This is a coordinate UPDATE, not coordinates; zero unusable atom outputs.
All parameters and rotary values are supplied. There is no sampled dropout, neighbor search, or additional coordinate transformation.
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.
An empty attention window gives zero attention update, but a usable atom retains its skip state and still runs the transition and coordinate readout.

Why this matters

The AF3 decoder broadcasts token context to cached atom features, applies local conditioned attention and a transition, then predicts a coordinate update for each usable atom.

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_atom_attention_decoder(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    eps: float = 1e-05,
    radius: int = -1,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    hidden: (B, A, C).
    attention: (B, H, A, K).
    position_update: (B, A, 3)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement af3_atom_attention_decoder")

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.