How the Function Token Decoder Maps InterPro Annotations to Protein Sequences

The Function Token Decoder converts quantized function tokens into per-position InterPro classifications using a compact Transformer stack, then collapses adjacent predictions into contiguous annotation ranges on the protein sequence.

The Function Token Decoder in the Biohub/esm repository serves as the interpretability layer for the ESM3 model's functional predictions. It transforms discrete token representations produced by the InterProQuantizedTokenizer back into human-readable InterPro annotations mapped to specific residue ranges on a protein sequence. Understanding this mapping pipeline is essential for extracting biological function from the model's latent space.

Configuration and InterPro Vocabulary Initialization

During instantiation in esm/models/function_decoder.py, the decoder loads the complete InterPro entry table to establish its prediction vocabulary. The configuration is defined by the FunctionTokenDecoderConfig dataclass:

@dataclass(frozen=True)
class FunctionTokenDecoderConfig:
    d_model: int = 1024
    n_heads: int = 8
    n_layers: int = 3
    function_token_vocab_size: int = 260
    function_token_depth: int = 8
    num_interpro_classes: int = 29026   # ← total InterPro classes available

    interpro_entry_list: str = field(default_factory=lambda: str(C.INTERPRO_ENTRY))

Source: esm/models/function_decoder.py (lines 21-38)

Inside __init__, the decoder reads the TSV file specified by interpro_entry_list and builds efficient lookup structures:

with AnyPath(config.interpro_entry_list).open("r") as f:
    df = pd.read_csv(f, sep="\t")
self.interpro_ids = sorted(df.ENTRY_AC)            # list of all InterPro IDs

self.interpro2index = {interpro_id: i for i, interpro_id in enumerate(self.interpro_ids)}
assert len(self.interpro_ids) == config.num_interpro_classes

Source: esm/models/function_decoder.py (lines 62-70)

This creates self.interpro_ids (a sorted list of 29,026 entries) and self.interpro2index for O(1) class-to-index mapping.

Transformer Forward Pass for Per-Position Classification

The decoder embeds input tokens and processes them through a lightweight Transformer (default: 3 layers, 8 heads, 1024 dimensions). The architecture includes a RegressionHead that projects hidden states to logits for every InterPro class:

"interpro_logits": RegressionHead(
    d_model=config.d_model,
    output_dim=config.num_interpro_classes,
    hidden_dim=4 * config.d_model,
)

Source: esm/models/function_decoder.py (lines 43-48)

Calling self(token_ids) (the forward method) returns a dictionary containing "interpro_logits"—a tensor of shape [batch_size, seq_len, 29026] representing the presence probability for every InterPro class at each sequence position.

Decoding Logits into Protein Sequence Annotations

The decode method orchestrates the conversion from raw logits to FunctionAnnotation objects anchored to specific protein coordinates.

Thresholding and Binarization

Raw logits are first converted to probabilities and thresholded to create a boolean mask:

interpro_preds = F.sigmoid(outputs["interpro_logits"])
interpro_preds = interpro_preds >= annotation_threshold   # binary mask

interpro_preds[~where_decode, :] = False                 # ignore padded positions

Source: esm/models/function_decoder.py (lines 133-137)

Constructing FunctionAnnotation Objects

For every (position, class) pair marked True in the mask, the decoder instantiates a FunctionAnnotation:

for position_index, class_index in zip(*preds.nonzero()):
    interpro_id = self.interpro_ids[class_index]
    annotation = FunctionAnnotation(
        label=interpro_id,
        start=position_index,          # 1-indexed, inclusive (BOS shifts +1)

        end=position_index,            # initially single-residue span

    )
    annotations.append(annotation)

Source: esm/models/function_decoder.py (lines 140-148)

Positions are 1-indexed to account for the BOS token, with start and end initially identical for single-residue hits.

Merging Contiguous Predictions

Single-residue hits are grouped into functional domains using merge_annotations:

annotations = merge_annotations(
    annotations, merge_gap_max=annotation_gap_merge_max
)

Source: esm/models/function_decoder.py (lines 150-152)

The merge_annotations utility (from esm/utils/misc.py) groups predictions by label, builds range objects, and collapses adjacent or nearby ranges (within annotation_gap_merge_max residues) into non-overlapping intervals with inclusive coordinates.

Minimum Length Filtering

Spurious short predictions are removed via length filtering:

if annotation_min_length is not None:
    annotations = [
        a for a in annotations
        if a.end - a.start + 1 >= annotation_min_length
    ]

Source: esm/models/function_decoder.py (lines 154-158)

Complete Usage Example

The following example demonstrates the end-to-end pipeline from quantized tokens to InterPro annotations:

import torch
from esm.models.function_decoder import FunctionTokenDecoder, FunctionTokenDecoderConfig
from esm.tokenization.function_tokenizer import InterProQuantizedTokenizer

# 1️⃣ Load the tokenizer (defines depth and special tokens)

tokenizer = InterProQuantizedTokenizer(depth=8)

# 2️⃣ Create dummy token matrix [seq_len, depth]

# In practice, this comes from model inference

dummy_token_ids = torch.randint(
    low=0,
    high=tokenizer.vocab_to_index["<lsh:0>"] + 1,
    size=(100, tokenizer.depth),  # example sequence length of 100

    dtype=torch.int32,
)

# 3️⃣ Initialize the decoder

decoder = FunctionTokenDecoder(FunctionTokenDecoderConfig())

# 4️⃣ Decode tokens → InterPro annotations

decoded = decoder.decode(
    function_token_ids=dummy_token_ids,
    tokenizer=tokenizer,
    annotation_threshold=0.1,      # confidence cutoff

    annotation_min_length=5,      # drop fragments shorter than 5 residues

)

# 5️⃣ Access mapped annotations

for ann in decoded["interpro_annotations"]:
    print(f"{ann.label}: residues {ann.start}-{ann.end}")

Key implementation details:

• The tokenizer provides the depth dimension and special token indices (<pad>, <none>, <unk>)
• decode returns three keys: "interpro_logits" (raw scores), "interpro_preds" (binary mask), and "interpro_annotations" (final range objects)

Summary

• The Function TokenDecoder loads 29,026 InterPro classes from INTERPRO_ENTRY during initialization
• A Transformer regression head generates per-position logits across all classes, producing a tensor of shape [batch, seq_len, 29026]
• Sigmoid thresholding converts logits to binary predictions, which are materialized as FunctionAnnotation objects with 1-indexed coordinates
• The merge_annotations utility (in esm/utils/misc.py) collapses adjacent hits into contiguous functional domains
• Final output maps specific InterPro IDs to inclusive residue ranges on the protein sequence, filtered by minimum length requirements

Frequently Asked Questions

What input format does the Function Token Decoder require?

The decoder accepts quantized function token IDs produced by the InterProQuantizedTokenizer, typically as an int32 tensor of shape [sequence_length, depth] where depth defaults to 8. These tokens represent compressed functional signatures extracted from the protein sequence.

How does the decoder handle multiple InterPro predictions at the same position?

Each position independently predicts across all 29,026 classes via binary classification. If multiple classes exceed the confidence threshold, the decoder generates separate FunctionAnnotation objects for each, allowing biologically accurate overlapping functional assignments on the same residue.

Why are annotation positions 1-indexed in the output?

The decoder accounts for the BOS (beginning-of-sequence) token prepended during tokenization, shifting coordinates by +1 to align with standard protein numbering conventions where the first residue is position 1. This ensures compatibility with biological databases and visualization tools.

Where does the decoder load InterPro entry definitions from?

It reads the TSV file specified by config.interpro_entry_list, which defaults to the path C.INTERPRO_ENTRY defined in esm/utils/constants/esm3.py. This file contains the official InterPro accession numbers (e.g., IPR000123) used to label the predictions.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}