All-atom folding

AlphaFold 3 Atom Attention Encoder

HardRaw PyTorchAbout 30 min1 public example plus private tests

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

Your task

Update each atom by attending over its supplied allowed key slots, then form each token’s output as the mean of projected, ReLU-activated hidden features from its usable atoms. One attention entry belongs to a query atom, head, and key slot. Use AF3 adaptive normalization and sigmoid residual gates.

Function signature

af3_atom_attention_encoder(
    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 encoder mode with two valid atoms [1,3], [4,2] mapped to one valid token. After 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 transition weights in this example give zero. Pooling the first hidden channel gives $((1-d)+(4+1/16))/2$ for their shared token. Here $\sigma(z)=1/(1+e^{-z})$ .

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.
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 and attention. Token output is the exact per-token mean of ReLU(hidden @ w_pool), using usable atom assignments; tokens with no atoms are zero.
All parameters, pair descriptors and window maps 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.

Why this matters

AF3 atom attention exchanges information within supplied local windows. After attention and transition residuals, a projection and ReLU prepare each atom for mean pooling into its token.

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_encoder(
    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).
    token_output: (B, T, U)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement af3_atom_attention_encoder")

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.