MSA and pair representations

MSA cluster profiles and deletion features

HardRaw PyTorchAbout 30 min1 public example plus private tests

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

Your task

Assign extra MSA rows to supplied cluster centers and build masked cluster profiles and deletion features. Return assignments, profiles, and the requested center/extra-row feature tensors.

Function signature

msa_profile_features(
    msa: torch.Tensor,
    extra_msa: torch.Tensor,
    msa_mask: torch.Tensor,
    extra_mask: torch.Tensor,
    deletions: torch.Tensor,
    extra_deletions: torch.Tensor,
    agreement_weights: torch.Tensor,
) -> dict[str, torch.Tensor]

Worked example

Count observations within each cluster

Extra row 0 exactly matches representative 0, with weighted agreement $1+1+1+0.2=3.2$ ; extra row 1 matches representative 1. Extra row 2 matches representative 0 at its third position, giving agreement $1$ . Thus assignments are [0, 1, 0]. In cluster 0 at position 0, the representative and first extra row contain category 0, while the third extra row contains category 3: frequencies are $2/3$ and $1/3$ .

import torch

msa = torch.tensor([[[0, 1, 2, 3], [1, 2, 0, 3]]], dtype=torch.int64)
extra_msa = torch.tensor([[[0, 1, 2, 3], [1, 2, 0, 3], [3, 3, 2, 1]]], dtype=torch.int64)
msa_mask = torch.full((1, 2, 4), True, dtype=torch.bool)
extra_mask = torch.full((1, 3, 4), True, dtype=torch.bool)
deletions = torch.tensor([[[0.33440397994552407, 0.7618063009752785, 0.3232192763322387, 0.6017582635498979], [0.6090226915412557, 0.21896612648967853, 0.42954764715125526, 0.41624465502273716]]], dtype=torch.float64, requires_grad=True)
extra_deletions = torch.tensor([[[0.311435034228908, 0.5125377709315249, 0.3756854033911365, 0.6103182835033594], [0.3925585032328245, 0.8529914727477061, 0.32150636483641926, 0.423268971317183], [0.21148476464956434, 0.4895142803803796, 0.4337284046223584, 0.5169061875189971]]], dtype=torch.float64, requires_grad=True)
agreement_weights = torch.tensor([1.0, 1.0, 1.0, 0.2, 0.0], dtype=torch.float64)

result = msa_profile_features(msa=msa, extra_msa=extra_msa, msa_mask=msa_mask, extra_mask=extra_mask, deletions=deletions, extra_deletions=extra_deletions, agreement_weights=agreement_weights)

Expected return

{
    'assignment': torch.tensor([[0, 1, 0]], dtype=torch.int64),
    'profile': torch.tensor([[[[0.6666666666666666, 0.0, 0.0, 0.3333333333333333, 0.0], [0.0, 0.6666666666666666, 0.0, 0.3333333333333333, 0.0], [0.0, 0.0, 1.0, 0.0, 0.0], [0.0, 0.3333333333333333, 0.0, 0.6666666666666666, 0.0]], [[0.0, 1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0, 0.0]]]], dtype=torch.float64),
    'mean_deletion': torch.tensor([[[0.2857745929413322, 0.5879527840957276, 0.3775443614485779, 0.5763275781907514], [0.5007905973870401, 0.5359787996186923, 0.37552700599383726, 0.41975681316996005]]], dtype=torch.float64),
    'features': torch.tensor([[[[1.0, 0.0, 0.0, 0.0, 0.0, 0.33440397994552407, 0.07067099371170948, 0.6666666666666666, 0.0, 0.0, 0.3333333333333333, 0.0, 0.06046081621029981], [0.0, 1.0, 0.0, 0.0, 0.0, 0.7618063009752785, 0.15831406843724105, 0.0, 0.6666666666666666, 0.0, 0.3333333333333333, 0.0, 0.12320585706872142], [0.0, 0.0, 1.0, 0.0, 0.0, 0.3232192763322387, 0.06832570248268699, 0.0, 0.0, 1.0, 0.0, 0.0, 0.07969841635957112], [0.0, 0.0, 0.0, 1.0, 0.0, 0.6017582635498979, 0.126024640440868, 0.0, 0.3333333333333333, 0.0, 0.6666666666666666, 0.0, 0.12082843409856038]], [[0.0, 1.0, 0.0, 0.0, 0.0, 0.6090226915412557, 0.12750588033819898, 0.0, 1.0, 0.0, 0.0, 0.0, 0.10530014209586326], [0.0, 0.0, 1.0, 0.0, 0.0, 0.21896612648967853, 0.046383804277378204, 0.0, 0.0, 1.0, 0.0, 0.0, 0.11255074789230625], [1.0, 0.0, 0.0, 0.0, 0.0, 0.42954764715125526, 0.09053747816281954, 1.0, 0.0, 0.0, 0.0, 0.0, 0.07927695955931023], [0.0, 0.0, 0.0, 1.0, 0.0, 0.41624465502273716, 0.08776950283838884, 0.0, 0.0, 0.0, 1.0, 0.0, 0.08850061222300255]]]], dtype=torch.float64),
    'extra_features': torch.tensor([[[[1.0, 0.0, 0.0, 0.0, 0.0, 0.311435034228908, 0.0658526810863654], [0.0, 1.0, 0.0, 0.0, 0.0, 0.5125377709315249, 0.10772383555365878], [0.0, 0.0, 1.0, 0.0, 0.0, 0.3756854033911365, 0.07931005376090094], [0.0, 0.0, 0.0, 1.0, 0.0, 0.6103182835033594, 0.12776990915364153]], [[0.0, 1.0, 0.0, 0.0, 0.0, 0.3925585032328245, 0.08283287352129885], [0.0, 0.0, 1.0, 0.0, 0.0, 0.8529914727477061, 0.17635630393825594], [1.0, 0.0, 0.0, 0.0, 0.0, 0.32150636483641926, 0.06796636049696612], [0.0, 0.0, 0.0, 1.0, 0.0, 0.423268971317183, 0.08923148672474324]], [[0.0, 0.0, 0.0, 1.0, 0.0, 0.21148476464956434, 0.04480434011383874], [0.0, 0.0, 0.0, 1.0, 0.0, 0.4895142803803796, 0.10297069404882583], [0.0, 0.0, 1.0, 0.0, 0.0, 0.4337284046223584, 0.09140666897406385], [0.0, 1.0, 0.0, 0.0, 0.0, 0.5169061875189971, 0.10862433366899243]]]], dtype=torch.float64),
}

Parameters

Parameter	Shape and meaning
`msa`torch.Tensor	$[B, S, N]$ Supplied MSA cluster-center token sequences.
`extra_msa`torch.Tensor	$[B, E, N]$ Extra MSA sequences to assign to those centers.
`msa_mask`torch.Tensor	$[B, S, N]$ True marks an observed sequence/residue entry in the MSA.
`extra_mask`torch.Tensor	$[B, E, N]$ True marks an observed residue in an extra MSA sequence.
`deletions`torch.Tensor	$[B, S, N]$ Deletion counts/continuous relaxation
`extra_deletions`torch.Tensor	$[B, E, N]$ Deletion counts/continuous relaxation
`agreement_weights`torch.Tensor	$[A]$ Fixed category weights used only when scoring token agreement.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Build nearest-center MSA cluster profiles and deletion features. $B,S,N,A$ positive; $E\ge 0$ extra rows is allowed. msa $[B,S,N]$ and extra_msa $[B,E,N]$ are int64 category IDs in $[0,A)$ . agreement_weights $[A]$ is a fixed nonnegative vector: for example,
ordinary amino acids $1$ , gaps downweighted, masked-token agreement $0$ . msa_mask/extra_mask are Boolean True means an observed position. Finite deletions $[B,S,N]$ and extra_deletions $[B,E,N]$ are nonnegative at observed positions, and can be treated as continuous noninteger values for differentiation.
Missing payloads are ignored. Score each extra row $e$ against center $s$ with $\mathrm{score}_{b,e,s}=\sum_n\mathbf{1}[\text{both masks and equal tokens}]\,\mathrm{agreement\_weights}[\mathrm{token}]$ . assignment $[b,e]$ is the maximum-score center; exact ties choose the lowest center index, even an entirely unobserved extra row.
$E=0$ returns an empty assignment.
For each center $s$ and residue $n$ , include the center observation if masked valid and all valid extra observations assigned to $s$ . profile $[b,s,n,a]$ is $\operatorname{count}(\mathrm{category}\,a)/\mathrm{observed\_count}$ , and mean_deletion is the same observed-position mean of deletion values.
Empty clusters/positions return exact zeros; do not divide by the number of sequences irrespective of per-position masks. Let onehot_center and onehot_extra be zero at missing positions.
$\mathrm{has\_deletion}=\operatorname{clamp}(d,0,1)$ , $\mathrm{bounded\_deletion}=(2/\pi)\cdot \arctan(d/3)$ , and $\mathrm{bounded\_mean}=(2/\pi)\cdot \arctan(\mathrm{mean\_deletion}/3)$ .
Concatenate features $[B,S,N,2A+3]$ in order $[\underbrace{\mathrm{center\_onehot}}_{A},\underbrace{\mathrm{has\_deletion}}_{1},\underbrace{\mathrm{bounded\_deletion}}_{1},\underbrace{\mathrm{profile}}_{A},\underbrace{\mathrm{bounded\_mean}}_{1}]$ , then zero all channels where the CENTER position is masked out.
Returned profile/mean_deletion themselves can still reflect valid extra rows at that position. extra_features $[B,E,N,A+2]$ = $[\underbrace{\mathrm{extra\_onehot}}_{A},\underbrace{\mathrm{extra\_has\_deletion}}_{1},\underbrace{\mathrm{extra\_bounded\_deletion}}_{1}]$ , zero on missing extra positions.
Return assignment, profile, mean_deletion, features and extra_features. Database search, sequence alignment, stochastic MSA subsampling and masked-token corruption occur outside this deterministic tensor operation.
The strict missing-center/count-zero behavior differs from preprocessing paths that assume unpadded valid cluster centers.

Why this matters

Each extra sequence joins the center with the strongest weighted agreement. Within a cluster, each residue position gets its own observed count: a missing character must not dilute a profile or deletion average.

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 msa_profile_features(
    msa: torch.Tensor,  # (B, S, N)
    extra_msa: torch.Tensor,  # (B, E, N)
    msa_mask: torch.Tensor,  # (B, S, N)
    extra_mask: torch.Tensor,  # (B, E, N)
    deletions: torch.Tensor,  # (B, S, N)
    extra_deletions: torch.Tensor,  # (B, E, N)
    agreement_weights: torch.Tensor,  # (A,)
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    assignment: (B, E).
    profile: (B, S, N, A).
    mean_deletion: (B, S, N).
    features: (B, S, N, 2A+3).
    extra_features: (B, E, N, A+2)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement msa_profile_features")

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.