Learned irrep query-key attention and tensor values

Your task

Function signature

se3_equivariant_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    parity: str = 'so3',
) -> dict[str, torch.Tensor]

Worked example

Two heads prefer different incoming vector messages

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Use row-stored real irreducible representations. features [N,T,C,D] holds T irrep types with degrees node_degrees [T] and inversion parities node_parities [T] in {−1,+1}; type t uses only m=0,...,2\cdot l_{t}. D is a common padded component width at least every irrep dimension.
• All counts except E must be positive; E\ge 0, N\ge 1. node_valid [N] and edge_valid [E] use True for present items. edge_index [2,E] is (\mathrm{receiver},\mathrm{sender}), not (\mathrm{sender},\mathrm{receiver}). Duplicate edges are separate attention items; self edges are allowed.
• Valid indices are in [0,N); disabled edges may have arbitrary integer indices. An edge is active iff edge_valid and both endpoint node masks are True. All masked node/edge payloads and padded components may contain NaN/Inf and must be replaced by zero before arithmetic; observed payloads and all parameters are finite.
• edge_sh [E,S,D] contains geometric edge irreps with edge_degrees [S] and edge_parities [S]. These are supplied harmonic evaluations, not supplied keys, values, or scores. edge_attr [E,R] is a supplied invariant radial/chemical feature vector.
• The caller builds these with the radius/harmonic exercises; coordinates are not an additional input here. Under a rotation each active component block of features and edge_sh transforms by its specified orthogonal real irrep matrix. Under inversion it is multiplied by its parity.
• paths [P,3] has (\mathrm{input\_node\_type},\mathrm{edge\_harmonic\_type},\mathrm{output\_node\_type}). Every index is valid and every degree path satisfies |l_{\mathrm{in}}-l_{\mathrm{edge}}|\le l_{\mathrm{out}} \le l_{\mathrm{in}}+l_{\mathrm{edge}}.
• Fixed cg [P,D,D,D] has axes (\mathrm{output\_component},\mathrm{input\_component},\mathrm{harmonic\_component}), zero padding, and is an intertwiner in the supplied common basis. Fixed path_scale [P] gives any desired path-normalization convention. It is not a learned radial weight.
• parity defaults to 'so3': any angularly valid path is allowed, but reflections are not promised. With 'o3', require p_{\mathrm{out}}=p_{\mathrm{in}}\cdot p_{\mathrm{edge}} for every path; reject a violation with ValueError. Unsupported parity strings, negative degrees, insufficient D, and triangle violations also raise ValueError.
• Other malformed shapes/tables are outside the contract. General supplied irreps are supported; graded rotations exercise l=0,1,2 in the real Cartesian/STF basis.
• Let A_{e}=\operatorname{SiLU}(\mathrm{edge\_attr\_e} \mathbin{@} \mathrm{edge\_weight} + \mathrm{edge\_bias}), where \operatorname{SiLU}(z)=z/(1+\operatorname{exp}(-z)), edge_weight [R,U], edge_bias [U].
• Form \mathrm{Kcoef}[e,p,c,h,k]=\sum_{u} A[e,u]\cdot \mathrm{key\_weight}[u,p,c,h,k]+\mathrm{key\_bias}[p,c,h,k], and Vcoef analogously using value_weight/value_bias. The key/value weight shapes are [U,P,C,H,K] and bias shapes [P,C,H,K]. These scalar coefficients can mix channels and paths, never irrep components independently.
• Q[i,t,h,k,m]=\sum_{c} \mathrm{features}[i,t,c,m]\cdot \mathrm{query\_weight}[t,c,h,k], with query_weight [T,C,H,K], no bias.
• For path p=(a,b,t), compute: Z[e,p,c,m]=\mathrm{path\_scale}[p]\cdot \sum_{r,s} \mathrm{cg}[p,m,r,s]\cdot \mathrm{features}[\mathrm{sender\_e},a,c,r]\cdot \mathrm{edge\_sh}[e,b,s]. Then K[e,t,h,k,m]=\sum_{p:\mathrm{out}=t,c} Z[e,p,c,m]\cdot \mathrm{Kcoef}[e,p,c,h,k], and V likewise.
• Thus an edge key/value is a learned tensor-product convolution. Its scalar logit is L[e,h]=\sum_{t,k,m} Q[\mathrm{receiver\_e},t,h,k,m]\cdot K[e,t,h,k,m] / \sqrt{H\cdot K\cdot \sum_{t}(2\cdot l_{t}+1)}. The H factor deliberately matches the audited RFdiffusion backend's full-key-fiber scaling, not a per-head-only denominator.
• Return zero logits on inactive edges. Normalize \operatorname{exp}(L) stably over ACTIVE incoming edges of each receiver, separately per head. Return attention [E,H], zero for inactive edges. Receivers with no active edges get an exactly zero message.
• Set M[i,t,h,k,m]=\sum_{e:\mathrm{receiver}=i} \mathrm{attention}[e,h]\cdot V[e,t,h,k,m].
• Merge heads and apply a learned self interaction: \mathrm{output}[i,t,o,m]=\sum_{h,k} M[i,t,h,k,m]\cdot \mathrm{merge\_weight}[t,h,k,o]+\sum_{c} \mathrm{features}[i,t,c,m]\cdot \mathrm{self\_weight}[t,c,o], with merge_weight [T,H,K,\mathrm{Cout}] and self_weight [T,C,\mathrm{Cout}].
• This is the bias-free projection of the concatenated message and residual, not an additional nonlinear block. Return features_out [N,T,\mathrm{Cout},D], zero for missing nodes and padded components, together with attention and logits. An isolated valid node retains its self interaction.
• Empty E and all-masked graphs must have finite connected zero derivatives for unused inputs. Do not modify input tensors. No other keys, dropout, feature normalization, or geometry updates are used.

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 se3_equivariant_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    parity: str = 'so3',
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    features_out: (N, T, Cout, D).
    attention: (E, H).
    logits: (E, H)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement se3_equivariant_attention")

Sources and further reading

• arxiv.org
• www.nature.com
• RosettaCommons/RFdiffusion · SE3_network.py
• RosettaCommons/RFdiffusion · attention.py
• RosettaCommons/RFdiffusion · convolution.py