MSA and pair representations

Tied Pair Row Attention

HardRaw PyTorchAbout 30 min2 public examples plus private tests

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

Your task

Update each residue-pair feature by mixing key-residue values with an attention map shared across value anchors. For row attention, first exchange both pair axes, apply the tied core below, then exchange the update axes back. Return the gated, projected update.

Function signature

tied_pair_row_attention(
    activations: dict,
    parameters: dict,
    eps: float = 1e-05,
) -> torch.Tensor

Worked example

Trace tied attention through a nonsymmetric pair grid

Use one batch, two residues, two feature channels and one head of width 1. The original pair layout is features[0] = [[[-1,1],[1,-1]], [[-1,1],[-1,1]]]. All normalization scales are one and shifts are zero; set eps=3. Each vector has mean 0 and variance 1, so normalization divides it by 2. Query, key and value projections all use the channel weights [-1,1]; pair-bias and gate logits are zero. The final projection uses weights [2,4] and bias [1,-1].

Row attention first transposes both pair axes. The resulting scalar query, key and value matrices are

Q=K=V=M=\begin{bmatrix}1&1\\-1&1\end{bmatrix}.

Rows of $M$ index the tied axis (and value anchor); columns index residue positions. Sum query–key products over that tied axis:

S=\frac{M^{\mathsf T}M}{\sqrt2}=\begin{bmatrix}\sqrt2&0\\0&\sqrt2\end{bmatrix},\qquad p=\sigma(\sqrt2)=\frac{1}{1+e^{-\sqrt2}}\approx0.804430.

Softmax over key columns gives

A=\begin{bmatrix}p&1-p\\1-p&p\end{bmatrix}.

For query 0, anchor 0 has values $[1,1]$ , so its mixture is 1; anchor 1 has $[-1,1]$ , so its mixture is $1-2p$ . The full mixture, with query before anchor, is

U=AV^{\mathsf T}=\begin{bmatrix}1&1-2p\\1&2p-1\end{bmatrix}.

The zero gate logits give a gate of $1/2$ . Each scalar $u$ becomes the two-channel update $(u/2)[2,4]+[1,-1]=[u+1,2u-1]$ . Transpose the query/anchor axes back to the original row-attention layout:

\begin{aligned} \operatorname{output}_{00}=\operatorname{output}_{01}&=[2,1],\\ \operatorname{output}_{10}&=[2-2p,1-4p]\approx[0.391141,-2.217719],\\ \operatorname{output}_{11}&=[2p,4p-3]\approx[1.608859,0.217719]. \end{aligned}

The batch axis is omitted in these pair indices. Return this update tensor with shape (1,2,2,2); do not add the input features.

Parameters

Parameter	Shape and meaning
`activations`dict	Supplied intermediate tensors.
`parameters`dict	Supplied model parameters, keyed as shown below.
`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,L,L,C]$

Return Attention update, without adding an input residual.

Required behavior

Normalize features and bias_features independently over their channel axes using affine population LayerNorm with eps (default 1e-5) inside the square root. Form bias-free Q/K/V and a sigmoid value gate from normalized features.
Pair bias is the normalized bias feature projected to heads. Normalize the final attention logits with softmax over key residue $j$ .
Return a projected attention update only, including output bias. Do not add the input residual.
Use CPU float32/float64 tensors with one shared floating dtype/device. Preserve strided layouts and do not mutate supplied tensors.
Use raw tensor operations, not a ready-made implementation of the requested primitive. Supplied parameters and random variates must not be reinitialized.
Transpose both pair axes of features and bias_features before the calculation, and transpose the final update back.
With batch implicit, $A_{ijh}=\operatorname{softmax}_j[\sum_{s,d}q_{sihd}k_{sjhd}/\sqrt{SD}+\beta_{ijh}]$ . The sequence-axis reduction is a sum scaled by $\sqrt S$ , not a mean.
Use $U_{ikhd}=\sum_j A_{ijh}v_{kjhd}$ , then gate at $(i,k)$ and project over $(h,d)$ . Query $i$ precedes value anchor $k$ ; do not swap those axes.

Why this matters

Tied pair attention forms one residue-to-residue attention map by summing evidence over the other pair axis, then reuses that map for every value row.

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_pair_row_attention(
    activations: dict,
    parameters: dict,
    eps: float = 1e-05,
) -> torch.Tensor:
    """Return a tensor with shape (B,L,L,C)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement tied_pair_row_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.

Related exercises

Sources and further reading

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