MSA and pair representations

Tied Pair Column 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. Column attention uses the supplied pair-axis orientation. In the core output U[i,k], i is the attention query and k selects the value-anchor row; then apply its gate and output projection.

Function signature

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

Worked example

Compute the shared map before mixing value anchors

Consider the score-and-readout stage after normalization and projection, with two residues, one head and width 1. Let Q=K=V=M=[[1,-1],[1,1]]; rows are tied-axis/value anchors and columns are residue positions. Set the projected pair bias to zero. The tied scores are M.T @ M / sqrt(2) = [[sqrt(2),0],[0,sqrt(2)]]. With p=1/(1+exp(-sqrt(2)))≈0.804430, key softmax gives A=[[p,1-p],[1-p,p]]. The mixture is U=A @ V.T=[[2p-1,1],[1-2p,1]], with query before anchor. Zero gate logits give 1/2; output weights [2,4] and bias [1,-1] map each scalar u to [u+1,2u-1]. Thus output pair (0,0) is [2p,4p-3]≈[1.608859,0.217719], and (0,1) is [2,1]. Column mode transposes neither input nor output.

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

All dimensions are positive; both pair axes have length L. There is no padding or dropout mask: every supplied residue participates.
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.
Do not transpose features, bias_features, or the returned update.
With batch implicit, $A_{ijh}=\operatorname{softmax}_j[\sum_{s,d}q_{sihd}k_{sjhd}/\sqrt{LD}+\beta_{ijh}]$ . The sequence-axis reduction is a sum scaled by $\sqrt L$ , 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 attention builds one residue-to-residue map by combining evidence across an additional axis. The pair variants have a deliberately unusual value layout; the full block then adds these updates through supplied dropout masks and a transition.

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_column_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_column_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.