ESMFold 2 Atom Attention Decoder

Your task

Function signature

esmfold2_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

Parameters

Return

dict[str, torch.Tensor]

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.
• Compute SiLU(condition) @ w_mod as six vectors in order: attention shift, scale, gate; transition shift, scale, gate. Each modulated feature is RMSNorm(x) * (1 + scale) + shift. Residual gates are RAW modulation values, not sigmoids.
• Project Q/K/V without biases; apply RMS normalization separately over each Q/K head width. Apply the supplied half-split rotary embedding to Q/K before gathering slots. The head width is even. There is no pair-logit bias.
• Compute Q\cdot K/\sqrt D; 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 and add readout_bias, 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.
• For even head width D, define \operatorname{rotate\_half}(q)=\operatorname{concat}(-q[...,D/2:],q[...,:D/2]). Replace Q and K by q\odot rope\_cos+\operatorname{rotate\_half}(q)\odot rope\_sin before gathering key slots.
• RMSNorm is x/\sqrt{\operatorname{mean}_{last}(x^2)+eps}. LayerNorm is (x-\mu)/\sqrt{\operatorname{mean}_{last}((x-\mu)^2)+eps}, where \mu=\operatorname{mean}_{last}(x).

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 esmfold2_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 esmfold2_atom_attention_decoder")

Related exercises

• Builds onESMFold 2 Atom Attention Encoder
• Combined exerciseLocal atom attention, encoder pooling and decoder skip/readout

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