ESM3 multitrack embedding and confidence fusion

Your task

Function signature

multitrack_embedding_fusion(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
) -> torch.Tensor

Worked example

Add categorical tracks and a nontrivial confidence response

Parameters

Return

torch.Tensor · [B, N, C]

Required behavior

• Fuse the complete supplied ESM3-style multitrack input representation. All token axes are [B,N], B,N\ge 1, and mask [B,N] is True means valid.
• The four main tracks (\mathrm{sequence},\mathrm{structure},\mathrm{ss8},\mathrm{sasa}) each have int64 tokens [B,N] and a floating table [V_{\mathrm{track}},C]; add their row lookups. Token 0 is an ordinary learned category in these four tracks.
• function_tokens [B,N,H] address function_table [H,\mathrm{Vf},\mathrm{Dh}], H,\mathrm{Dh}\ge 1 with C=H\cdot \mathrm{Dh}. Look up each head independently and CONCATENATE head 0..H-1 embeddings, not sum/average. Function ID 0 is padding and contributes exact zero even if the supplied table's zero row is nonzero.
• annotation_tokens [B,N,A] address annotation_table [\mathrm{Va},C]. SUM annotations within a token, excluding ID 0. A may be 0 and duplicate IDs count repeatedly. Add this bag to the main/function embeddings. confidence [B,N,2] contains average-structure and per-residue confidence in that order.
• For supplied centers [R] and positive scalar width, compute \phi_{\mathrm{bntr}}=\operatorname{exp}(-((\mathrm{confidence\_bnt}-\mathrm{centers\_r})/\mathrm{width})^{2}), with R\ge 1. Each t has its own confidence_weight [2,R,C] and confidence_bias [2,C].
• Add \sum_{t}(\phi_{t}\mathbin{@}\mathrm{weight\_t}+\mathrm{bias\_t}). The original open encoder uses R=16, centers evenly spaced from 0 to 1 inclusive, width 1/16; these resources are supplied so smaller examples and derivatives can be studied. There is no one-half in the exponent and no normalization over centers.
• Return embedding [B,N,C], the sum of these eight contributions (four main tracks, function concatenation, annotation sum and two confidence projections). All floating tensors share \mathrm{float32}/64. Valid token IDs must index their table; valid confidence, tables and parameters are finite and width> 0.
• Scaled confidence differences, their squares, and projection/reduction intermediates must remain finite in the common dtype. Masked token IDs may be arbitrary and masked confidence may be NaN: sanitize before lookup or arithmetic; final output is exact zero. has no keys.
• Tokenizers, structure quantization, annotation retrieval, confidence prediction and pretrained parameters are outside this computation.

Why this matters

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 multitrack_embedding_fusion(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
) -> torch.Tensor:
    """Return a tensor with shape (B, N, C)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement multitrack_embedding_fusion")

Sources and further reading

• www.science.org
• Biohub/esm · esm3.py
• Biohub/esm · misc.py