MSA and pair representations

Triangle multiplicative update: incoming and outgoing

HardRaw PyTorchAbout 30 min2 public examples plus private tests

AlphaFold2OpenFoldAlphaFold3Boltz-1Boltz-2ProtenixRoseTTAFold2RoseTTAFold All-AtomESMFold ESMFold2

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

Your task

Update each directed residue pair by multiplying projected pair features along two-edge paths through an intermediate residue. Use the selected incoming or outgoing orientation and return the gated update.

Function signature

triangle_multiplicative_update(
    activations: dict,
    parameters: dict,
    masks: dict,
    eps: float = 1e-05,
    orientation: str = 'outgoing',
) -> torch.Tensor

Worked example

Which two edges meet at the intermediate residue?

Inspect one hidden channel after the input projections, gates and leg masks have produced a=[[1,2],[3,4]] and b=[[5,6],[7,8]], with row/column indices ordered by residue 0 then 1. For target pair (0,1), outgoing contracts a[0,k] with b[1,k]: 1×7+2×8=23. Incoming contracts a[k,0] with b[k,1]: 1×6+3×8=30. These are hidden-channel sums before the required hidden LayerNorm, output projection and final gate. Neither is divided by a path count, and neither sums over hidden channels.

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.
`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.
`orientation`strDefault `'outgoing'`	Choose "outgoing" to combine (i,k) with (j,k), or "incoming" to combine (k,i) with (k,j). These are directed pair features.

Return

torch.Tensor · $[B, N, N, \mathrm{Cz}]$

Return [B,N,N,Cz] gated pair update only; masked targets exactly zero.

Required behavior

Implement BOTH incoming and outgoing triangle multiplicative updates with explicit learned parameters. $\operatorname{LN}(x;\,\mathrm{scale},\mathrm{bias})\,=\,\mathrm{scale}\cdot (x-\operatorname{mean}_{\mathrm{last}}(x))/\sqrt{\operatorname{mean}_{\mathrm{last}}((x-\operatorname{mean}_{\mathrm{last}}(x))^2)+\varepsilon}+\mathrm{bias}$ .
Use population variance, over channels only; eps is inside the square root. All dimension sizes are positive, all floating inputs are finite and share a dtype/device, and all learned tensors are supplied (no parameter initialization). First $x=\operatorname{LN}(z;\,\mathrm{z\_scale},\mathrm{z\_bias})$ .
For every directed pair ( $i,j$ ) and hidden channel $h$ , $a_{ijh}=\mathrm{pair\_mask}_{ij}\cdot (\sum_{c}\,x_{ijc}\cdot \mathrm{w\_a}_{ch}+\mathrm{b\_a}_{h})\cdot \sigma(\sum_{c}\,x_{ijc}\cdot \mathrm{w\_a\_gate}_{ch}+\mathrm{b\_a\_gate}_{h})$ ; $\mathrm{b\_ijh}$ is defined analogously using w_b,b_b,w_b_gate,b_b_gate.
Batch indices are independent and omitted.
For orientation='outgoing', $t_{ijh}=\sum_{k}\,a_{ikh}\cdot \mathrm{b\_jkh}$ .
For orientation='incoming', $t_{ijh}=\sum_{k}\,a_{kih}\cdot \mathrm{b\_kjh}$ . These sums run over intermediate residues $k$ ONLY, not hidden channels; do not divide by $N$ or a valid count. Normalize $t$ over hidden channels: $u=\operatorname{LN}(t;\,\mathrm{hidden\_scale},\mathrm{hidden\_bias})$ , using the SAME eps.
Project $p_{ijc}=\sum_{h}\,u_{ijh}\cdot \mathrm{w\_out\_hc}+\mathrm{b\_out}_{c}$ . The final gate is from INPUT-normalized $x$ , $g_{ijc}=\sigma(\sum_{d}\,\mathrm{x\_ijd}\cdot \mathrm{w\_out\_gate\_dc}+\mathrm{b\_out\_gate}_{c})$ . Return $\mathrm{update}_{ijc}=\mathrm{pair\_mask}_{ij}\cdot p_{ijc}\cdot g_{ijc}$ .
Both input legs are masked after their affine projections/gates; apply final target-pair masking after output bias and gate. No residual addition, dropout, softmax, symmetry enforcement, or detached intermediates.
$A$ masked target is exactly zero; a valid target with no contributing triangles may still be nonzero due to learned hidden/output biases. Config: orientation ('outgoing' default, or 'incoming'); eps (positive float, default $1\times10^{-5}$ ).

Why this matters

Two learned views of the pair tensor form the legs of a triangle. For each hidden channel, the operator adds products over intermediate residues, then normalizes and gates the resulting target-pair update.

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

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

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