All-atom folding

ESMFold 2 Atom Attention Encoder

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

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 ESMFold2 RMS modulation, rotary query/key features, and raw residual gates.

Function signature

esmfold2_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

In ESMFold2 encoder mode with illustrative eps=1, the valid atom vectors are $[1,3]$ , $[3,1]$ . Write the first as $[l,h]$ and let $r=\sqrt{(l^2+h^2)/2+1}$ . RMS-normalized value vectors are $[l,h]/r$ and $[h,l]/r$ . After q/k RMS normalization, rotate the second atom’s q/k channels by $\pi/2$ , making its key proportional to $[-l,h]$ . Query 0’s self-key and second-key scores therefore differ by $2l^2/((l^2+h^2+1)\sqrt2)$ ; its self weight is the sigmoid of that difference. Condition 1 and residual-modulation weight $1+e^{-1}$ make the raw attention residual multiplier 1, while the value gate is 1/2, so its first hidden channel is $l+[pl+(1-p)h]/(2r)$ . Query 1 uses only its own key. Both atoms map to the same token; ReLU keeps their positive projected channels and pooling averages the two hidden vectors. The supplied transition update 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.
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 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 and rotary values are supplied. There is no sampled dropout, neighbor search, or additional coordinate transformation.
Use $RMSNorm(v)=v/\sqrt{mean(v^2)+eps}$ , $sigmoid(z)=1/(1+e^{-z})$ , and $SiLU(z)=z\,sigmoid(z)$ .
For a head vector split into equal halves $q=[q_1,q_2]$ , define $rotateHalf(q)=[-q_2,q_1]$ and apply $q\odot rope\_cos+rotateHalf(q)\odot rope\_sin$ . Apply the same rule to K.

Why this matters

Atoms attend only to supplied key slots. Conditioning controls how attention and transition updates enter the residual state. The encoder pools atom features into tokens; the decoder combines token context with atom state to predict coordinate changes.

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_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 esmfold2_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.