Conditioned pair-bias attention and gated transition block

Your task

Function signature

conditioned_transformer_block(
    activations: dict,
    parameters: dict,
    masks: dict,
    eps: float = 1e-05,
    style: str = 'af3',
) -> dict[str, torch.Tensor]

Worked example

Worked example: attention changes the state consumed by a nonzero transition

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Implement one complete conditioned token block with B,N,C,S,P,H,D,V,M\ge 1. style is 'af3' (default) or 'esmfold2'; they differ here in whether normalized pairs have an affine offset. All shapes below are exact, all active data and learned parameters finite and of one dtype/device. \varepsilon=\mathrm{eps}>0 defaults 10^{-5}.
• Define LN as population last-channel centering divided by \sqrt{\mathrm{variance}+\varepsilon}, without affine.
• For stage j, define: A_{j}(z,c)=\operatorname{LN}(z)\cdot \operatorname{sigmoid}((\operatorname{LN}(c)\cdot \mathrm{cond\_scale}[j])\mathbin{@}w_{\mathrm{scale}}[j]+b_{\mathrm{scale}}[j])+(\operatorname{LN}(c)\cdot \mathrm{cond\_scale}[j])\mathbin{@}w_{\mathrm{shift}}[j]. Here j=0 denotes attention and j=1 denotes the transition.
• Packed stage axes contain DISTINCT learned parameters. Before arithmetic zero x/condition at invalid tokens and pair at disallowed edges. An edge i→k is allowed iff valid_i & valid_k & pair_mask_ik.
• Let z=A_0(x,c). Project Q_{bnhd}=\sum_c z_{bnc}w^q_{chd}+b^q_{hd}, and project K/V similarly without bias. The query projection bias is added after contraction.
• For af3, \mathrm{Bias\_ihk}=(\operatorname{LN}(\mathrm{pair\_ik})\cdot \mathrm{pair\_scale})\mathbin{@}w_{\mathrm{pair}}. The esmfold2 source also has an affine pair_bias [P] before w_pair; its projected contribution is the same scalar for every key in a head.
• Cancel that softmax-invariant term ALGEBRAICALLY before evaluating pair logits, using the same centered formula as af3. Adding a huge offset numerically and then subtracting a softmax maximum can already have erased meaningful pair differences. Keep pair_bias as a declared input with exactly zero VJP. The result uses [B,H,N,N] head order.
• Softmax over key k of \operatorname{dot}(Q_{i},K_{k})/\sqrt{D}+\mathrm{Bias} uses allowed edges; no keys gives exact-zero attention, with a safe denominator and no softmax \epsilon. Average values, multiply by \operatorname{sigmoid}(z\mathbin{@}w_{\mathrm{value\_gate}}), flatten head-major, and multiply w_attention_out.
• Multiply this update by \operatorname{sigmoid}(c\mathbin{@}w_{\mathrm{residual\_gate}}[0]+b_{\mathrm{residual\_gate}}[0]) using RAW c, then add x to get after_attention. Apply A_1 to this FIRST residual result.
• Transition \mathrm{hidden}=\operatorname{SiLU}(A_{1}\mathbin{@}w_{\mathrm{ff\_gate}})\cdot (A_{1}\mathbin{@}w_{\mathrm{ff\_up}}), followed by w_ff_out and \operatorname{sigmoid}(c\mathbin{@}w_{\mathrm{residual\_gate}}[1]+b_{\mathrm{residual\_gate}}[1]); add to after_attention for output.
• Return both complete residual states plus attention. Only b_q, b_scale and b_residual_gate are projection biases; the esmfold2 variant also has pair_bias in its normalization. All other projections are bias-free. Mask both returned states after residuals; invalid payloads may be NaN/Inf.
• Valid rows with no keys still undergo the transition, while their attention update is zero. No causality, dropout, hidden weights, precomputed network update or implicit multi-block stack. Key/value widths may differ. Numerical domain: \varepsilon must be finite and remain strictly positive when represented in the input dtype.
• Active inputs and parameters must be bounded so the required arithmetic intermediates are finite and within that dtype's representable range: normalization sums/means, centered differences, squares, variance-plus-\varepsilon, projections, residual/statistic updates, and intermediates for the graded first-order derivatives.
• Finite inputs alone do not guarantee this; overflow-sized reductions or squares are outside the contract. Representable subnormal \operatorname{sigmoid}/\operatorname{SiLU} tails are in scope and must not be prematurely discarded.
• Use sigmoid(z)=1/(1+e^{-z}) and SiLU(z)=z\,sigmoid(z).

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 conditioned_transformer_block(
    activations: dict,
    parameters: dict,
    masks: dict,
    eps: float = 1e-05,
    style: str = 'af3',
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    output: (B, N, C).
    after_attention: (B, N, C). (B,N,C), after the first residual only.
    attention: (B, H, N, N)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement conditioned_transformer_block")

Sources and further reading

• www.nature.com
• google-deepmind/alphafold3 · diffusion_transformer.py
• google-deepmind/alphafold3 · haiku_modules.py
• www.biorxiv.org
• huggingface/transformers · modeling_esmfold2.py