Template-axis attention and mean aggregation

Your task

Function signature

template_pair_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
    mode: str = 'attention',
    mean_eps: float = 1e-07,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

A query favors one template

Parameters

Return

tuple[torch.Tensor, torch.Tensor]

Required behavior

• Return a template-derived pair UPDATE, without residual addition. All dimensions are positive (T\ge 1; missing templates use False masks), all floats finite and same dtype/device. template[B,T,N,N,C_t] contains already pair-processed template features.
• template_mask[B,T,N,N] and pair_mask[B,N,N] are Boolean True means usable; effective mask is their conjunction. Invalid template payloads are ignored and invalid/empty query pairs return exact zero, including output bias.
• mode='attention' (default): query[B,N,N,C_q], w_q[C_q,H,D], w_k[C_t,H,D], w_v[C_t,H,V] are bias-free projections.
• For each b,i,j,h, compute \mathrm{logits}_{t}=\operatorname{dot}(\mathrm{query}_{ij}\,\mathrm{w\_q}_{h},\mathrm{template}_{tij}\,\mathrm{w\_k}_{h})/\sqrt{D}. Apply numerically stable masked softmax over T ONLY. Expose weights[B,N,N,H,T]; all-masked weights are zero.
• Average projected template values, flatten H,V in head-major order, apply w_out[H\cdot V,C_q] and b_out[C_q], then zero empty/invalid pairs. There is no gate or LayerNorm inside this attention variant; callers supply the desired preprocessed template/query representations.
• mode='mean': query and attention weights are absent. Average usable template vectors with denominator \operatorname{count}+\mathrm{mean\_eps} (\mathrm{mean\_eps}\ge 0 default 1\times10^{-7}); if \operatorname{count}=0 use an exact zero average. Apply elementwise ReLU, then w_mean[C_t,C_q] and b_mean[C_q].
• Return weights[B,T,N,N] equal to \mathrm{effective\_mask}/(\operatorname{count}+\mathrm{mean\_eps}), zero for empty pairs. This is AF3's mean/ReLU/readout algebra; with all templates present the denominator is T+1\times10^{-7}, as in the pinned code.
• A pair-specific mask is an educational extension, not AF3's total-template denominator for a partially invalid pair. ReLU inputs exactly zero are excluded from smooth-gradient fixtures unless a structural all-masked zero. No dropout, pair-trunk iteration, or hidden parameters.

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 template_pair_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
    mode: str = 'attention',
    mean_eps: float = 1e-07,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    update: (B, N, N, Cq).
    weights: Attention mode: (B, N, N, H, T). Mean mode: (B, T, N, N)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement template_pair_attention")

Related exercises

• Builds onTemplate Pair Cross-Attention
• Builds onTemplate Pair Mean

Sources and further reading

• www.nature.com
• www.nature.com
• www.nature.com
• google-deepmind/alphafold · modules.py
• google-deepmind/alphafold3 · template_modules.py
• RosettaCommons/RFdiffusion · Embeddings.py