Prott5

ProtT5: Protein Language Model Embeddings

Overview

ProtT5 is a family of T5-based protein language models from Rostlab, trained on millions of UniRef protein sequences. The recommended model — ProtT5-XL-U50 — is an encoder-decoder transformer with 3B parameters whose encoder produces 1024-dimensional embeddings that outperform earlier BERT-based protein models and match or exceed ESM2 on most tasks.

Primary use cases:

• Per-residue and per-protein embeddings for downstream ML
• Protein classification (localization, membrane, function)
• Regression (stability, fitness, thermostability)
• Sequence similarity search

Installation

pip install transformers torch sentencepiece

Model Selection

Model	Type	Params	HF ID	Best for
ProtT5-XL-U50	T5 encoder	3B	Rostlab/prot_t5_xl_half_uniref50-enc	Best default — half-precision encoder only
ProtT5-XL-U50 (full)	T5 enc-dec	3B	Rostlab/prot_t5_xl_uniref50	When you also need the decoder
ProtT5-XXL-U50	T5 encoder	11B	Rostlab/prot_t5_xxl_uniref50	Maximum accuracy, multi-GPU
ProtBERT-BFD	BERT	~420M	Rostlab/prot_bert_bfd	Faster/lighter, lower quality
ProtBERT	BERT	~420M	Rostlab/prot_bert	UniRef100 trained BERT

Use prot_t5_xl_half_uniref50-enc as the default: it's encoder-only (no decoder weights), loads in half-precision, and is what all published benchmarks use.

Critical Preprocessing

This is the most important thing to get right. ProtT5 requires two preprocessing steps that ESM2 does not:

import re

def preprocess(sequence: str) -> str:
    # 1. Map rare/ambiguous amino acids to X
    sequence = re.sub(r"[UZOB]", "X", sequence)
    # 2. Space-separate every amino acid (ProtT5 is character-level)
    return " ".join(list(sequence))

sequence = "MKTAYIAKQRQISFVK"
processed = preprocess(sequence)
# → "M K T A Y I A K Q R Q I S F V K"

Skip either step and you will get garbage embeddings — the model was trained on this exact format.

Core Usage

Single Sequence

import re
import torch
from transformers import T5Tokenizer, T5EncoderModel

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = T5Tokenizer.from_pretrained("Rostlab/prot_t5_xl_half_uniref50-enc", do_lower_case=False)
model = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_half_uniref50-enc").to(device)
if device.type == "cpu":
    model = model.float()   # half-precision not supported on CPU
model.eval()

sequence = "MKTAYIAKQRQISFVK"
processed = " ".join(list(re.sub(r"[UZOB]", "X", sequence)))

ids = tokenizer(processed, return_tensors="pt", add_special_tokens=True).to(device)

with torch.no_grad():
    output = model(**ids)

# output.last_hidden_state: (1, seq_len + 1, 1024)
# T5 encoder adds an EOS token at the end — slice to remove it
per_residue = output.last_hidden_state[0, :len(sequence)]   # (L, 1024)
per_sequence = per_residue.mean(dim=0)                       # (1024,)

print(per_residue.shape)   # (16, 1024)
print(per_sequence.shape)  # (1024,)

Note: T5 encoder output has shape (seq_len + 1, 1024) — the +1 is the trailing EOS token. Slice [:len(sequence)] to get only residue embeddings. Unlike BERT/ESM2, there is no leading [CLS] token in the T5 encoder output.

Batch Processing

def get_embeddings(model, tokenizer, sequences, device):
    """Mean-pooled per-sequence embeddings for a list of sequences."""
    # Preprocess all sequences
    processed = [" ".join(list(re.sub(r"[UZOB]", "X", s))) for s in sequences]
    lengths = [len(s) for s in sequences]

    ids = tokenizer(
        processed,
        return_tensors="pt",
        padding="longest",
        add_special_tokens=True,
    ).to(device)

    with torch.no_grad():
        output = model(**ids)

    embeddings = []
    for i, seq_len in enumerate(lengths):
        # Slice per-residue embeddings (no leading special token, EOS at end)
        per_residue = output.last_hidden_state[i, :seq_len]  # (L, 1024)
        per_sequence = per_residue.mean(dim=0)               # (1024,)
        embeddings.append(per_sequence.cpu().float())

    return torch.stack(embeddings)  # (N, 1024)

See scripts/embed.py for a CLI that reads FASTA and saves embeddings.

Scripts

# Batch embed a FASTA file → .npy
python scripts/embed.py sequences.fasta --output embeddings.npy

# Save as HDF5, L2-normalized
python scripts/embed.py sequences.fasta --output embeddings.h5 --normalize

# Use ProtBERT-BFD for faster/lighter inference
python scripts/embed.py sequences.fasta --model Rostlab/prot_bert_bfd --output embeddings.npy

Fine-Tuned Checkpoints

Rostlab publishes ready-to-use fine-tuned models:

Task	HF ID
Secondary structure (3-class)	Rostlab/prot_bert_bfd_ss3
Membrane prediction	Rostlab/prot_bert_bfd_membrane
Subcellular localization	Rostlab/prot_bert_bfd_localization

Load with the appropriate HuggingFace pipeline for quick prediction without training.

Comparison with ESM2 and ESM-C

	ProtT5-XL-U50	ESM2-650M	ESM-C 600M
Architecture	T5 encoder	BERT encoder	Custom encoder
Package	transformers	transformers	esm SDK
Hidden dim	1024	1280	1152
Preprocessing	space-sep + X mapping	none	none
Special token handling	slice [:seq_len]	slice [1:-1]	none
Subcell. Loc. (DeepLoc)	86	83	—
Variant effect (DMS)	0.53	—	—

ProtT5 is a strong choice for localization, function prediction, and variant effect tasks. ESM-C is faster for pure embedding throughput.

Best Practices

• Always use half-precision on GPU (prot_t5_xl_half_uniref50-enc is already half-precision by default)
• Max sequence length: 1022 residues (1024 minus 2 for special tokens) — truncate or chunk longer sequences
• L2-normalize before cosine similarity or FAISS indexing
• For comparison with published benchmarks, always use prot_t5_xl_half_uniref50-enc

Resources

• GitHub: https://github.com/agemagician/ProtTrans
• HuggingFace: https://huggingface.co/Rostlab
• Paper: Elnaggar et al., IEEE TPAMI 2021 — https://doi.org/10.1109/TPAMI.2021.3095381
• Pre-computed embeddings: available via UniProt for all reviewed sequences
• LambdaPP webservice: https://embed.predictprotein.org/

References

• references/prott5-api.md — batch processing, layer extraction, ProtBERT usage, downstream tasks, FAISS, fine-tuning patterns