RFdiffusion tied axial and sequence-weighted attention

Your task

Function signature

tied_axial_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    variant: str = 'pair_row',
    eps: float = 1e-05,
    sequence_dropout: float = 0.0,
    pair_dropout: float = 0.0,
    transition_dropout: float = 0.0,
) -> torch.Tensor

Worked example

One attention row, two value anchors

Parameters

Return

torch.Tensor · Exactly features.shape: [B, S, L, C].

Required behavior

• Implement variant='pair_row' (default), 'pair_column', 'msa_row', or the complete 'pair_block'. All floating inputs/parameters are supplied, finite float32/float64 CPU tensors of one dtype; dimensions B,S,L,C,C_b,H,D are positive and arithmetic must remain representable. No input mutation or random draws.
• There is no padding/residue mask: every supplied residue/sequence participates. Boolean keep masks denote RETAINED supplied dropout events, not biological validity.
• Define \operatorname{LN}(v;a,b)=a\cdot (v-\operatorname{mean}_{\mathrm{channels}}(v))/\sqrt{\operatorname{mean}_{\mathrm{channels}}((v-\operatorname{mean}_{\mathrm{channels}}(v))^2)+\varepsilon}+b, with population variance, \varepsilon>0 default 1\times10^{-5} INSIDE sqrt.
• Define \sigma(z)=1/(1+\exp(-z)), evaluated stably; \operatorname{ReLU}(z)=\max(z,0). CORE INPUTS: features X[B,S,L,C], bias_features Y[B,L,L,C_b], feature_scale/bias[C], bias_scale/bias[C_b], w_q,w_k,w_v,w_gate[C,H,D], b_gate[H,D], w_bias[C_b,H], w_out[H,D,C], b_out[C].
• Pair cores require S=L.
• For pair_row ONLY, transpose BOTH X and Y over their two pair axes before the following equations, and transpose the resulting update back. pair_column and msa_row do not transpose.
• Let x=\operatorname{LN}(X;\mathrm{feature\_scale},\mathrm{feature\_bias}), y=\operatorname{LN}(Y;\mathrm{bias\_scale},\mathrm{bias\_bias}); q_{sihd}=\sum_{c}\,x_{sic}\cdot \mathrm{w\_q}_{chd} (similarly k,v), \beta_{ijh}=\sum_{c}\,y_{ijc}\cdot \mathrm{w\_bias}_{ch}, g_{sihd}=\sigma(\sum_{c}\,x_{sic}\cdot \mathrm{w\_gate}_{chd}+\mathrm{b\_gate}_{hd}).
• Batch b is implicit.
• For pair cores, A_{ijh}=\operatorname{softmax}_{j}[\sum_{s,d}\,q_{sihd}\cdot k_{sjhd}/\sqrt{S\cdot D}+\beta_{ijh}]. This SUM ties one attention map across the first pair axis, with \sqrt{S}, not mean-by-S, normalization.
• The exact pinned value contraction is U_{ikhd}=\sum_{j}\,A_{ijh}\cdot v_{kjhd}: query i precedes value-anchor k. Gate at x_{ik}, then return \mathrm{update}_{ikc}=\sum_{h,d}\,U_{ikhd}\cdot g_{ikhd}\cdot \mathrm{w\_out}_{hdc}+\mathrm{b\_out}_{c}. Do not silently replace U_{ik} by U_{ki}.
• These are not per-edge triangle logits or a triangle closing-edge bias. MSA ROW: additionally supply w_seq_q,w_seq_k[C,H,D], b_seq_q,b_seq_k[H,D] and sequence_keep[B,S,L,H].
• Use normalized TARGET ROW ZERO: Q_{ihd}=\sum_{c}\,x_{0ic}\cdot \mathrm{w\_seq\_q}_{chd}+\mathrm{b\_seq\_q}_{hd}; K_{sihd}=\sum_{c}\,x_{sic}\cdot \mathrm{w\_seq\_k}_{chd}+\mathrm{b\_seq\_k}_{hd}.
• Let a_{sih}=\operatorname{softmax}_{s}(\sum_{d}\,Q_{ihd}\cdot K_{sihd}/\sqrt{D}); r_{sih}=a_{sih}\cdot \mathrm{sequence\_keep}_{sih}/(1-\mathrm{sequence\_dropout}).
• Then A_{ijh}=\operatorname{softmax}_{j}[\sum_{s,d}\,r_{sih}\cdot q_{sihd}\cdot k_{sjhd}/\sqrt{D}+\beta_{ijh}], U_{sihd}=\sum_{j}\,A_{ijh}\cdot v_{sjhd}, and \mathrm{update}_{sic}=\sum_{h,d}\,U_{sihd}\cdot g_{sihd}\cdot \mathrm{w\_out}_{hdc}+\mathrm{b\_out}_{c}.
• The dropout is on SequenceWeight before weighting queries, never on main values or final attention. Even all-false sequence_keep leaves bias-only attention and generally nonzero output. The three core variants return the update tensor, without residual addition.
• PAIR BLOCK: only features[B,L,L,C], ca_coordinates[B,L,3], the pair-block tensors listed below, and prefixed attention parameters are supplied; unprefixed core parameters/bias_features are absent. Compute d_{ij}=\Vert \mathrm{ca}_{i}-\mathrm{ca}_{j}\Vert with exact zero self-distances and no epsilon added to d.
• The 36 bare Gaussian features are \rho_{ijm}=\exp(-((d_{ij}-\mu_{m})/(20/36))^2), \mu_{m}=20\cdot m/35, m=0,\ldots,35. There is no cutoff or cosine envelope.
• Geometry bias Y=\operatorname{ReLU}(\rho\,\mathrm{w\_rbf\_in}+\mathrm{b\_rbf\_in})\,\mathrm{w\_rbf\_out}+\mathrm{b\_rbf\_out}, with w_rbf_in[36,D], b_rbf_in[D], w_rbf_out[D,C], b_rbf_out[C].
• Separate row_ and column_ prefixes supply every one of the twelve core attention parameter names, both using C_b=C and the same H,D.
• First Z_1=X+\mathrm{pair\_row\_core}(X,Y;\mathrm{row\_params})\cdot \mathrm{row\_keep}/(1-\mathrm{pair\_dropout}), then Z_2=Z_1+\mathrm{pair\_column\_core}(Z_1,Y;\mathrm{column\_params})\cdot \mathrm{column\_keep}/(1-\mathrm{pair\_dropout}).
• row_keep[B,1,L,C] broadcasts over the FIRST pair axis; column_keep[B,L,1,C] broadcasts over the SECOND. Attention output biases are inside these masks.
• Let h=\operatorname{ReLU}(\operatorname{LN}(Z_2;\mathrm{transition\_scale},\mathrm{transition\_bias})\,\mathrm{w\_transition\_in}+\mathrm{b\_transition\_in}), with input matrix[C,2C] and bias[2C]. Let \mathrm{h\_drop}=h\cdot \mathrm{transition\_keep}/(1-\mathrm{transition\_dropout}), mask[B,L,L,2C].
• Return output Z_2+\mathrm{h\_drop}\,\mathrm{w\_transition\_out}+\mathrm{b\_transition\_out}, with output matrix[2C,C] and bias[C]. Unlike the core modes, pair_block returns the COMPLETE residual-composed pair features. Its final transition output bias is after hidden dropout.
• The only keys are variant, eps, sequence_dropout, pair_dropout, transition_dropout. All dropout probabilities are finite in [0,1), default 0 (evaluation); sequence_dropout is used only by msa_row, and pair_dropout/transition_dropout only by pair_block.
• For evaluation supply all-true keep masks. Incorrect variant, \varepsilon\le 0, dropout outside[0,1), nonpositive feature dimensions or nonsquare pair features raise ValueError. Other malformed shapes and nonfinite values are outside the contract.
• No trained checkpoint, surrounding three-track iteration, \mathrm{SE}(3) transform, or biological design outcome is implied.

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 tied_axial_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    variant: str = 'pair_row',
    eps: float = 1e-05,
    sequence_dropout: float = 0.0,
    pair_dropout: float = 0.0,
    transition_dropout: float = 0.0,
) -> torch.Tensor:
    """Return a tensor with shape features.shape: (B, S, L, C)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement tied_axial_attention")

Related exercises

• Builds onTied Pair Row Attention
• Builds onTied Pair Column Attention
• Builds onSequence-Weighted MSA Attention
• Builds onRoseTTAFold Pair Block

Sources and further reading

• www.nature.com
• RosettaCommons/RFdiffusion · Attention_module.py
• RosettaCommons/RFdiffusion · Attention_module.py
• RosettaCommons/RFdiffusion · Track_module.py
• RosettaCommons/RFdiffusion · util_module.py
• RosettaCommons/RFdiffusion · Attention_module.py