Invariant point attention (IPA)

Your task

Function signature

invariant_point_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    eps: float = 1e-08,
) -> dict[str, torch.Tensor]

Worked example

Nearby frames receive more weight

All local q_points and k_points: [0,0,0]
All local v_points: [1,0,0]
All scalar q, k and pair_bias: 0
Scalar values v: 2 and 6
pair_values (query, key): [[1,3],[4,2]]
eps=1; b_out=[0]
w_out = [[1],[0],[0],[0],[0],[1]]

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Implement the complete scalar/point/pair AGGREGATION CORE of monomer-style IPA followed by its supplied affine output projection.
• This explicitly starts after input projections: q,k,v are projected single-feature tensors; q_points,k_points,v_points are projected LOCAL points (including any learned point-projection biases); pair_bias is the projected per-head pair logit; pair_values is the pair representation to aggregate.
• Do not invent input projections, layer norms or additional weights. All dimensions B,N,H,D,\mathrm{Dv},P,Q,\mathrm{Cz},\mathrm{Cout} are positive, and all floating tensors are finite with common dtype/device.
• Frames use column-vector convention T_{i}(x)=\mathrm{rotation\_i}\mathbin{@}x+\mathrm{translation\_i} with rotation_i a proper \operatorname{SO}(3) matrix. Transform every q/k/v point to global coordinates.
• For each batch, query i, key j and head h, define L_{ijh}=\frac{q_{ih}^{\top}k_{jh}}{\sqrt{3D}}+\frac{\texttt{pair\_bias}_{ijh}}{\sqrt3}-\frac12\operatorname{softplus}(\texttt{point\_weight\_raw}_h)\sqrt{\frac{2}{27P}}\sum_p\left\|T_i(\texttt{q\_points}_{ihp})-T_j(\texttt{k\_points}_{jhp})\right\|^2 .
• \operatorname{softplus}(x)=\operatorname{log}(1+\operatorname{exp}(x)), implemented without overflow; distance has a NEGATIVE sign, sums over all P query/key points and all three coordinates, and uses squared distance without \epsilon. The nonnegative distance weight is learned separately for every head. True is valid.
• Permit key j exactly when mask_i AND mask_j AND pair_mask_ij. Stable softmax over j produces attention [B,H,N,N]. An empty neighborhood gives all-zero attention; no \epsilon is added to its denominator.
• Aggregate scalar o_{\mathrm{scalar\_ih}}=\sum_{j} a_{\mathrm{hij}}\cdot v_{\mathrm{jh}} and pair o_{\mathrm{pair\_ih}}=\sum_{j} a_{\mathrm{hij}}\cdot \mathrm{pair\_values\_ij}.
• For every value point p, aggregate its GLOBAL position, then inverse-transform into the query frame: \mathrm{local\_points\_ihp}=\mathrm{rotation\_i}^{\top}\mathbin{@}(\sum_{j} a_{\mathrm{hij}}\cdot T_{j}(v_{\mathrm{points\_jhp}})-\mathrm{translation\_i}). Explicitly set local_points to zero for an empty neighborhood.
• Build features by concatenating, in THIS order: all scalar heads flattened (H,\mathrm{Dv}); all local x coordinates flattened (H,Q); all local y; all local z; all point lengths \sqrt{\sum_{\mathrm{xyz}}(\mathrm{local\_points}^{2})+\varepsilon} flattened (H,Q); all pair heads flattened (H,\mathrm{Cz}).
• Each flatten is head-major. Feature length F=H\cdot (\mathrm{Dv}+4\cdot Q+\mathrm{Cz}). Empty-neighborhood length features are also zero, not \sqrt{\varepsilon}. \mathrm{update}=\mathrm{features}\mathbin{@}w_{\mathrm{out}}+b_{\mathrm{out}}, then zero the entire update for an empty neighborhood.
• Return update, attention and local_points; there is no residual, final gate, dropout or frame update. has only \varepsilon>0 (default 10^{-8}), used inside the OUTPUT point-length square root, not the attention distance.
• All three outputs are invariant when a common proper rotation G and translation u change frames to rotation'i=G\mathbin{@}\mathrm{rotation\_i} and translation'i=G\mathbin{@}\mathrm{translation\_i}+u; local inputs and pair/scalar features stay fixed.
• This is a full projected-input core, not a checkpoint-compatible full IPA layer.

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 invariant_point_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    eps: float = 1e-08,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    update: (B, N, Cout).
    attention: (B, H, N, N).
    local_points: (B, N, H, Q, 3)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement invariant_point_attention")

Sources and further reading

• aqlaboratory/openfold · structure_module.py