Local atom attention, encoder pooling and decoder skip/readout

Your task

Function signature

windowed_atom_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    eps: float = 1e-05,
    style: str = 'af3',
    stage: str = 'decoder',
    radius: int = -1,
) -> dict[str, torch.Tensor]

Worked example

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

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Implement one learned local atom block and its encoder or decoder boundary. B,A,T,K,C,S,H,D,V,M,U \ge 1; all active data/weights are finite and share dtype/device, indices int64, masks bool.
• stage is 'decoder' (default) or 'encoder'; style is 'af3' (default) or 'esmfold2'; \varepsilon>0 defaults 10^{-5} and radius is an integer \ge -1, default -1. Atom a is usable iff valid [a], 0\le \mathrm{atom\_to\_token}[a]<T and its token_valid is True.
• Invalid payloads may be NaN/Inf and must be sanitized before normalization/projection. Encoder starts x=\mathrm{atom}. Decoder starts x=\mathrm{atom}(\mathrm{skip})+\operatorname{gather}(\mathrm{token}\mathbin{@}w_{\mathrm{token}},\mathrm{atom\_to\_token}); projection is bias-free and invalid atoms are zero.
• key_index[B,A,K] supplies global within-batch atom IDs; no batch mixing. A slot is permitted iff its key ID is in [0,A), query/key atoms are usable, and window_mask is True.
• If \mathrm{radius}\ge 0 also require \left|\mathrm{key\_id}-\mathrm{query\_id}\right|\le \mathrm{radius}, a bidirectional INDEX window, not a spatial neighbor search. -1/\mathrm{out}-\mathrm{of}-\mathrm{range} keys are missing. Repeated slots count separately; repeated selection across queries accumulates gradients.
• All-missing windows have zero attention, but valid atoms retain residual/transition computation.
• AF3 style uses A_{j}(z,c)=\operatorname{LN}(z)\cdot \operatorname{sigmoid}((\operatorname{LN}(c)\cdot \mathrm{cond\_scale}[j])\mathbin{@}w_{\mathrm{scale}}[j]+b_{\mathrm{scale}}[j])+(\operatorname{LN}(c)\cdot \mathrm{cond\_scale}[j])\mathbin{@}w_{\mathrm{shift}}[j], population LN with \varepsilon inside root; j=0 query,1 key,2 transition.
• Query/key normalized atom states project to Q/K; only Q has b_q. Values come from key-normalized atoms. Per-slot logit bias is (\operatorname{LN}(\mathrm{pair})\cdot \mathrm{pair\_scale})\mathbin{@}w_{\mathrm{pair}}.
• Residual gates are \operatorname{sigmoid}(\mathrm{raw} \mathrm{condition}\mathbin{@}w_{\mathrm{residual\_gate}}[j]+b_{\mathrm{residual\_gate}}[j]), j=0 attention,1 transition.
• ESMFold2 style uses six vectors from \operatorname{SiLU}(\mathrm{condition}) contracted with w_mod [S,6,C], ordered attention_shift, attention_scale, attention_gate, transition_shift, transition_scale, transition_gate.
• Attention \mathrm{input}=\operatorname{RMS}(x)\cdot (1+\mathrm{attention\_scale})+\mathrm{attention\_shift}, shared by Q/K/V and value gating.
• \operatorname{RMS}(z)=z/\sqrt{\operatorname{mean_last}(z^{2})+\varepsilon}, no affine; additionally \operatorname{RMS}-\operatorname{normalize} Q and K independently over each head's D coordinates. D is even.
• Apply supplied full-width rope_cos/rope_sin by half-split rotation \operatorname{rotate_half}(q)=\operatorname{concat}(-q[D/2:],q[:D/2]); rotate Q/K before gathering keys. The cosine/sine construction from 3 D coordinates and reference UIDs belongs to spatial_rotary_embedding; here those tensors are supplied differentiable inputs.
• There is no b_q or pair bias in this style. Residual gates are the RAW attention_gate and transition_gate vectors, not sigmoid gates.
• In both styles, attention logits are \operatorname{dot}(Q_{\mathrm{query}},K_{\mathrm{key}})/\sqrt{D} plus any pair bias, stabilized masked softmax over K slots, with exact zero empty rows.
• Weighted values are multiplied by \operatorname{sigmoid}(\mathrm{query}-\mathrm{normalized}-\mathrm{input} \mathrm{contracted} \mathrm{with} w_{\mathrm{value\_gate}}), flattened head-major, and projected with w_attention_out, without bias. Add the appropriate attention-gated update to x.
• Normalize/modulate this NEW x for transition (AF3 A_2; ESMFold2 \operatorname{RMS}(x)\cdot (1+\mathrm{transition\_scale})+\mathrm{transition\_shift}). \mathrm{Transition}=\operatorname{SiLU}(z\mathbin{@}w_{\mathrm{ff\_gate}})\cdot (z\mathbin{@}w_{\mathrm{ff\_up}}), then w_ff_out, then the style-specific transition residual gate.
• Add it and mask to obtain hidden. Decoder: af3 \mathrm{position\_update}=(\operatorname{LN}(\mathrm{hidden})\cdot \mathrm{readout\_scale})\mathbin{@}w_{\mathrm{readout}}; esmfold2 adds learned readout_bias [C] after the LN scale and before w_readout.
• Mask the result to zero; it is a coordinate UPDATE, not final coordinates and not a residual with atom positions. Encoder: token_output is the exact per-token mean of \operatorname{ReLU}(\mathrm{hidden}\mathbin{@}w_{\mathrm{pool}}) over usable assigned atoms; empty/masked tokens give zero. All boundary linear projections are bias-free.
• Final LayerNorm has scale only for AF3 and both scale and bias for ESMFold2. Return hidden/attention plus only the stage-specific position_update or token_output. No weights, geometry descriptors, UID encodings, schedules or window maps are inferred; no dropout, resampling, frame updates or input mutation.
• Source ESMFold2 atom RMS \varepsilon=10^{-6} can be selected explicitly; one \varepsilon governs this exercise's normalizers. Numerical domain: \varepsilon must be finite and remain strictly positive when represented in the input dtype.
• Active inputs and parameters must be bounded so the required arithmetic intermediates are finite and within that dtype's representable range: normalization sums/means, centered differences, squares, variance-plus-\varepsilon, projections, residual/statistic updates, and intermediates for the graded first-order derivatives.
• Finite inputs alone do not guarantee this; overflow-sized reductions or squares are outside the contract. Representable subnormal \operatorname{sigmoid}/\operatorname{SiLU} tails are in scope and must not be prematurely discarded.

Why this matters

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 windowed_atom_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    eps: float = 1e-05,
    style: str = 'af3',
    stage: str = 'decoder',
    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). Returned only for stage="decoder".
    token_output: (B, T, U). Returned only for stage="encoder"."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement windowed_atom_attention")

Related exercises

• Builds onAlphaFold 3 Atom Attention Encoder
• Builds onAlphaFold 3 Atom Attention Decoder
• Builds onESMFold 2 Atom Attention Encoder
• Builds onESMFold 2 Atom Attention Decoder

Sources and further reading

• www.nature.com
• google-deepmind/alphafold3 · diffusion_transformer.py
• google-deepmind/alphafold3 · atom_cross_attention.py
• www.biorxiv.org
• huggingface/transformers · modeling_esmfold2.py
• huggingface/transformers · modeling_esmfold2.py