MSA and pair representations

ESMFold Single-to-Pair Features

MediumRaw PyTorchAbout 20 min2 public examples plus private tests

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

Your task

Convert single-token vectors into an ordered feature vector for every token pair. After input normalization and projection into q and k, pair (i,j) joins the componentwise product k_i*q_j with the difference q_j−k_i. Return these pair features and their supplied linear readout; reversing a pair can change its features.

Function signature

esmfold_single_to_pair(
    activations: dict,
    parameters: dict,
    masks: dict,
    eps: float = 1e-05,
    output_norm: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

Check both pair directions

LayerNorm gives [-0.5, 0.5] and [0.5, -0.5], producing $q=[1,-1]$ and $k=[0.5,-0.5]$ . Pair (0, 1) uses $q_1k_0=-0.5$ and $q_1-k_0=-1.5$ . Pair (1, 0) uses $q_0k_1=-0.5$ and $q_0-k_1=1.5$ . Here the product agrees in both directions while the difference changes sign.

import torch

activations = {
    'x': torch.tensor([[[-1.0, 1.0], [1.0, -1.0]]], dtype=torch.float64, requires_grad=True),
}
parameters = {
    'weight': torch.tensor([[-1.0, 0.0], [1.0, 1.0]], dtype=torch.float64, requires_grad=True),
    'bias': torch.tensor([0.0, 0.0], dtype=torch.float64, requires_grad=True),
    'norm_scale': torch.tensor([1.0, 1.0], dtype=torch.float64, requires_grad=True),
    'norm_bias': torch.tensor([0.0, 0.0], dtype=torch.float64, requires_grad=True),
    'out_weight': torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float64, requires_grad=True),
    'out_bias': torch.tensor([0.0, 0.0], dtype=torch.float64, requires_grad=True),
}
masks = {
    'mask': torch.tensor([[True, True]], dtype=torch.bool),
}
eps = 3.0
output_norm = False

result = esmfold_single_to_pair(activations=activations, parameters=parameters, masks=masks, eps=eps, output_norm=output_norm)

Expected return

(torch.tensor([[[[0.5, 0.5], [-0.5, -1.5]], [[-0.5, 1.5], [0.5, -0.5]]]], dtype=torch.float64), torch.tensor([[[[0.5, 0.5], [-0.5, -1.5]], [[-0.5, 1.5], [0.5, -0.5]]]], dtype=torch.float64))

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 stabilizer inside each population LayerNorm variance square root, including optional output normalization.
`output_norm`boolDefault `False`	Apply the supplied final affine normalization when True.

Return

tuple[torch.Tensor, torch.Tensor]

Return these tensors in the stated order.

Required behavior

x is $[B,N,C]$ ; mask is $[B,N]$ , with True meaning present. Sanitize invalid token features before normalization/projection, and zero both returned pair tensors unless both residue endpoints are present.
Apply affine population LayerNorm to x before the initial weight/bias projection.
Split the projected feature axis into q then k. Pair(i,j) is the concatenation $[k_i\odot q_j,\;q_j-k_i]$ ; this ordered difference is directional.
Project pair_features directly with out_weight and out_bias.
output_norm defaults to False. When True, apply affine population LayerNorm to the final output using output_scale/output_bias.
Every LayerNorm uses eps (default 1e-5) inside the population-variance square root. Return pair_features before the learned pair head and output after it.

Why this matters

Pair representations describe relationships between two positions. Product channels capture shared feature strength, while difference channels preserve direction. The AF3 atom variant instead combines atom conditioning, reference geometry and optional trunk context.

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 esmfold_single_to_pair(
    activations: dict,
    parameters: dict,
    masks: dict,
    eps: float = 1e-05,
    output_norm: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    pair_features: (B,N,N,2H).
    output: (B,N,N,D)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement esmfold_single_to_pair")

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.