MSA and pair representations

Triangle attention: starting and ending node

HardRaw PyTorchAbout 30 min3 public examples plus private tests

AlphaFold2OpenFoldAlphaFold3Boltz-1Boltz-2ProtenixESMFold

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Your task

Update each directed residue pair by attending along pairs sharing its starting or ending residue. Return the projected attention update using the selected orientation and masks.

Function signature

triangle_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
    orientation: str = 'starting',
    chunk_size: int = 0,
    eps: float = 1e-05,
) -> torch.Tensor

Worked example

A closing edge changes the attention weight

Use the starting orientation and fix target pair (i,j). Inspect a one-head, D=1 stage after normalization/projection: the two allowed key edges (i,0),(i,1) have scaled QK scores [0,0], values [2,6], and closing-edge biases β[j,0]=0, β[j,1]=log(3). Softmax over k gives [1/4,3/4], so the mixed value is 5. With query gate 1/2, w_out=[[2,4]] and b_out=[1,-1], this valid target update is [6,9]. Ending orientation instead fixes j and reads key/value edges (k,j) with closing bias β[k,i].

Parameters

Parameter	Shape and meaning
`activations`dict	Supplied intermediate tensors.
`parameters`dict	Supplied model parameters, keyed as shown below.
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`orientation`strDefault `'starting'`	Choose "starting" for pairs sharing the query pair’s first residue, or "ending" for pairs sharing its second residue.
`chunk_size`intDefault `0`	Optional chunk width; zero uses the unchunked operation.
`eps`floatDefault `1e-05`	Positive value added to final-axis population variance inside each LayerNorm square root. It is not a mask/count denominator or an output offset.

Return

torch.Tensor · $[B, N, N, C]$

Return [B,N,N,C] gated pair update; invalid query edges exactly zero.

Required behavior

Implement a gated multihead pair UPDATE, not a residual sum. $B,N,C,H,D$ are positive; all floating inputs are finite and share dtype/device. True means a valid directed edge in pair_mask.
Normalize each $z$ $[b,i,j,:]$ independently: $x=(z-\operatorname{mean}_{C}(z))/\sqrt{\operatorname{mean}_{C}((z-\operatorname{mean}_{C}(z))^2)+\varepsilon}\cdot \mathrm{norm\_scale}+\mathrm{norm\_bias}$ . The variance is population variance and eps is inside the square root.
Let $q_{ijhd}=\sum_{c}\,x_{ijc}\cdot \mathrm{w\_q}[\mathrm{chd}]$ , with analogous bias-free $k$ and $v$ . Closing-edge bias is $\beta_{ijh}=\sum_{c}\,x_{ijc}\cdot \mathrm{w\_bias}[\mathrm{ch}]$ . Gate $g_{ijhd}=\sigma(\sum_{c}\,x_{ijc}\cdot \mathrm{w\_gate}[\mathrm{chd}]+\mathrm{b\_gate}[\mathrm{hd}])$ .
orientation='starting': for fixed $i$ , $\mathrm{logits}[h,j,k]=\operatorname{dot}(q_{ijh},k_{ikh})/\sqrt{D}+\mathrm{beta\_jkh}$ ; permitted keys are those with $\mathrm{pair\_mask}[i,k]=\mathrm{True}$ . Average $v_{ikh}$ over $k$ with a stable masked softmax.
orientation='ending': for fixed $j$ , $\mathrm{logits}[h,i,k]=\operatorname{dot}(q_{ijh},k_{kjh})/\sqrt{D}+\beta_{kih}$ ; permitted keys have $\mathrm{pair\_mask}[k,j]=\mathrm{True}$ and values are $v_{kjh}$ .
Thus $\mathrm{ending}(z,m)=\mathrm{transpose}_{ij}(\mathrm{starting}(\mathrm{transpose}_{ij}(z),\mathrm{transpose}_{ij}(m)))$ with the SAME parameters. The closing-edge bias is not separately masked: pair_mask controls key edges and returned query updates, not a three-edge clique filter.
Multiply each averaged $[H,D]$ result by $g_{ij}$ , flatten in head-major/channel-minor order, then $\mathrm{update}_{ij}=\operatorname{flatten}(g\cdot \mathrm{average})\,\mathrm{w\_out}+\mathrm{b\_out}$ . Return exact zero when pair_mask $[i,j]$ is False or the fixed row/column has no valid keys, including the output bias.
$A$ fully masked softmax is zero, not uniform; no epsilon is added to its denominator. Softmax is over $k$ , never heads. All $q$ / $k$ / $v$ and closing-bias projections have no biases; gate and output biases are supplied.
Config: orientation is 'starting' (default) or 'ending'; $\varepsilon>0$ defaults to $1\times10^{-5}$ ; chunk_size is a nonnegative integer default $0$ (dense).
Positive chunk_size partitions the fixed-node axis after orientation transposition, retaining all $N$ keys for every query; it must produce the same answer, including a partial final chunk. No dropout, residual addition, or hidden parameters.

Why this matters

A query edge compares neighboring edges around one fixed endpoint. A third, closing edge contributes a learned bias to each comparison; that bias is not itself an extra validity mask.

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 triangle_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
    orientation: str = 'starting',
    chunk_size: int = 0,
    eps: float = 1e-05,
) -> torch.Tensor:
    """Return a tensor with shape (B, N, N, C)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement triangle_attention")

Ready to implement it?

3 progressive hints, a solution walkthrough, and grading that checks values, shapes and gradients are waiting in the workspace.

Write and run your solution Sign in

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Sources and further reading

aqlaboratory/openfold · triangular_attention.py

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.