MSA and pair representations

ESM3 Single-to-Pair Features

MediumRaw PyTorchAbout 20 min1 public example plus private tests

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

Your task

Build ESM3 pair features from projected single-token vectors, then apply the supplied GELU, normalization and output head.

Function signature

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

Worked example

Check the raw pair direction

The identity projection gives q=[1,3] and k=[2,5]. Pair (0,1) is [k₀q₁,q₁−k₀]=[6,1], whereas pair (1,0) is [k₁q₀,q₀−k₁]=[5,−4]. Apply the supplied hidden projection, exact GELU, hidden LayerNorm, and final projection to each pair independently.

import torch

activations = {
    'x': torch.tensor([[[1.0, 2.0], [3.0, 5.0]]], dtype=torch.float64, requires_grad=True),
}
parameters = {
    'weight': torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float64, requires_grad=True),
    'bias': torch.tensor([0.0, 0.0], dtype=torch.float64, requires_grad=True),
    'hidden_weight': torch.tensor([[0.21216724624478506, 0.06723158926988397, -0.22075659103655332, -0.26883329892386415], [0.2893647720587629, 0.050376568994800024, -0.27911394562654585, 0.4091670321049685]], dtype=torch.float64, requires_grad=True),
    'hidden_bias': torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64, requires_grad=True),
    'hidden_scale': torch.tensor([0.795208817925842, 0.020389689607759776, -0.0021728448198021777, 0.6797717713123919], dtype=torch.float64, requires_grad=True),
    'hidden_norm_bias': torch.tensor([0.4990308466741449, 0.412780913199289, -0.09772624102845953, -0.13636775593871922], dtype=torch.float64, requires_grad=True),
    'out_weight': torch.tensor([[0.11925094884659707, -0.2225214435612696, 0.19144161103199026], [-0.7885706318288417, -0.3458584542086859, 0.4039829309365077], [0.4198020222911164, -0.10846524434043135, 0.20913154379117166], [-0.13320172563627275, -0.9911054716945893, 0.17832603147645124]], dtype=torch.float64, requires_grad=True),
    'out_bias': torch.tensor([0.0, 0.0, 0.0], dtype=torch.float64, requires_grad=True),
}
masks = {
    'mask': torch.tensor([[True, True]], dtype=torch.bool),
}
eps = 1e-05
include_pairwise = False

result = esm3_single_to_pair(activations=activations, parameters=parameters, masks=masks, eps=eps, include_pairwise=include_pairwise)

Expected return

(torch.tensor([[[[2.0, -1.0], [6.0, 1.0]], [[5.0, -4.0], [15.0, -2.0]]]], dtype=torch.float64), torch.tensor([[[[-0.06920852641660866, 0.6486721557558881, 0.2181486014425758], [-0.0543721550180028, 0.14350769645932748, 0.3727625527042009]], [[-0.40388905713661477, 0.2904779634860536, 0.009239601425780924], [-0.05849757277360139, 0.12680340136716683, 0.37476460055425737]]]], 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 the hidden LayerNorm variance square root, after exact GELU.
`include_pairwise`boolDefault `False`	Append supplied pairwise-state channels 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.
Do not normalize 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 with hidden_weight/hidden_bias and apply exact erf GELU. Apply affine LayerNorm to that hidden vector before out_weight/out_bias.
include_pairwise defaults to False. When True, append the supplied masked pairwise_state after the product/difference features and before the hidden projection.
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.
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.
Exact GELU is $x[1+\operatorname{erf}(x/\sqrt2)]/2$ .

Why this matters

ESM3 forms directional pair features from projected single-token channels. Products capture shared strength; the ordered difference distinguishes the two directions.

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