Follow Blog posts from the web or conferences

Protein language models: from sequence to ecology

A walk through based on a lecture by Antrea Pavlou at the 11th EBAME workshop

Protein language models (PLMs) read proteins the way an LLM reads text. They are trained on hundreds of millions of sequences, and they learn which amino acids tend to appear together. ESM2 is the best known.

We can’t read what a PLM has learnt. But we can turn each protein into a vector of numbers, and train a small model on those numbers. In this tutorial we ask one question: can we predict where a protein was found (deep or surface water) from its sequence alone?

We expect basic Python (pandas, running scripts) and a working conda. New to conda? Install Miniconda first.

The idea

We use two models. They do very different jobs.

protein sequence → [ frozen PLM ] → embedding → [ small predictor ] → habitat
                    (not trained)                (we train this)

Three words we’ll use:

Token: one amino acid. The model reads a protein as a string of tokens.
Embedding: a vector of numbers for each token. It holds what the model learnt about that residue in context.
Frozen model: a pre-trained PLM, used as it is. Training one from scratch needs far too much compute.

PLMs learn by masked prediction: hide some residues, guess them, get penalised when wrong. We don’t repeat the training here. The optional section at the end shows it in action.

Setup

We create one environment for the whole tutorial.

mamba create -n plm -c conda-forge -c bioconda \
    python=3.11 pandas scikit-learn biopython matplotlib mmseqs2
mamba activate plm
pip install torch transformers

We need two input files:

proteins.faa: our proteins, in FASTA format
metadata.tsv: one row per protein, with a habitat column

protein_id	habitat
prot_0001	deep
prot_0002	surface
prot_0003	deep

Quick check that the two files agree:

grep -c '>' proteins.faa
tail -n +2 metadata.tsv | wc -l

Both numbers must match. The FASTA ids must be identical to protein_id.

The default CPU install of PyTorch is fine for the small model below. For thousands of proteins with a bigger model, we want a GPU.

One protein first

We start with a single sequence. Save this as embed_one.py:

import torch
from transformers import AutoTokenizer, EsmModel

MODEL = "facebook/esm2_t6_8M_UR50D"
SEQ = "MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVKALPDAQFEVV"

tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = EsmModel.from_pretrained(MODEL)
model.eval()

inputs = tokenizer(SEQ, return_tensors="pt")
with torch.no_grad():
    out = model(**inputs)

# One vector per token
per_residue = out.last_hidden_state[0]
print(per_residue.shape)

# One vector per protein: average over the real residues
mean_embedding = per_residue[1:-1].mean(dim=0)
print(mean_embedding.shape)

python embed_one.py

The first run downloads the model (a few tens of MB). Then:

torch.Size([72, 320])
torch.Size([320])

where:

MODEL — a small ESM2 model (8 million parameters), hosted on Hugging Face
model.eval() — we use the model as it is, no training
torch.no_grad() — no gradients needed, so it’s faster and uses less memory
[1:-1] — drops the two special tokens the model adds, <cls> at the start and <eos> at the end

Our protein has 70 residues. We get 72 vectors, each with 320 numbers. The mean gives a single vector of 320 numbers for the whole protein.

Averaging is simple and it works for many tasks. But it throws information away: the order of residues is gone. Use per-residue embeddings if the task depends on specific positions.

Change MODEL to facebook/esm2_t12_35M_UR50D and run again. What happens to the second number? Bigger models have longer vectors, but bigger isn’t always better. For many tasks the small ones do just as well.

Many proteins

For the records, we write here a prototype of script that embeds a whole FASTA file. Adapt the variables at the top.

Save this as embed_all.py:

#!/usr/bin/env python3
import numpy as np
import torch
from Bio import SeqIO
from transformers import AutoTokenizer, EsmModel

MODEL = "facebook/esm2_t6_8M_UR50D"
FASTA = "proteins.faa"
OUT = "embeddings"
BATCH = 16
MAXLEN = 1022   # ESM2 accepts at most 1024 tokens, including <cls> and <eos>

# Sorting by length keeps the padding small inside each batch
records = list(SeqIO.parse(FASTA, "fasta"))
records.sort(key=lambda r: len(r.seq))
ids = [r.id for r in records]
seqs = [str(r.seq).replace("*", "")[:MAXLEN] for r in records]

device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = EsmModel.from_pretrained(MODEL).to(device).eval()

chunks = []
for i in range(0, len(seqs), BATCH):
    batch = seqs[i:i + BATCH]
    inputs = tokenizer(batch, return_tensors="pt", padding=True).to(device)
    with torch.no_grad():
        hidden = model(**inputs).last_hidden_state

    # Mean over real residues only: remove padding, <cls> and <eos>
    mask = inputs["attention_mask"].clone()
    lengths = mask.sum(dim=1)
    mask[:, 0] = 0
    mask[torch.arange(mask.size(0)), lengths - 1] = 0
    mask = mask.unsqueeze(-1)

    mean = (hidden * mask).sum(dim=1) / mask.sum(dim=1)
    chunks.append(mean.cpu().numpy())
    print(f"{min(i + BATCH, len(seqs))}/{len(seqs)}", end="\r")

np.save(f"{OUT}.npy", np.vstack(chunks))
with open(f"{OUT}.ids.txt", "w") as f:
    f.write("\n".join(ids) + "\n")

python embed_all.py

We get two files: embeddings.npy (one row per protein) and embeddings.ids.txt (the protein ids, in the same row order).

Proteins longer than MAXLEN are cut. Check how many we lose before trusting the results.

Training the predictor

Now the small model. First, a rule we must follow.

Proteins from the same gene cluster are similar. If one is in the training set and its sibling is in the validation set, the model can cheat. So we cluster the proteins first, and keep each cluster on one side of the split.

mmseqs easy-cluster proteins.faa clusters tmp --min-seq-id 0.5 -c 0.8

where:

--min-seq-id 0.5 — proteins with at least 50% identity go in the same cluster
-c 0.8 — the alignment must cover at least 80% of the sequences

The file clusters_cluster.tsv has two columns: the cluster representative and the member. We use the representative as the group label.

Save this as train.py:

import numpy as np
import pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import confusion_matrix
from sklearn.model_selection import GroupKFold, cross_val_predict, cross_val_score
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

EMB = "embeddings.npy"
IDS = "embeddings.ids.txt"
META = "metadata.tsv"
CLUSTERS = "clusters_cluster.tsv"
FOLDS = 5

X = np.load(EMB)
ids = [line.strip() for line in open(IDS)]

meta = pd.read_csv(META, sep="\t").set_index("protein_id")
clusters = pd.read_csv(CLUSTERS, sep="\t", header=None,
                       names=["rep", "member"]).set_index("member")

y = meta.loc[ids, "habitat"].values
groups = clusters.loc[ids, "rep"].values

model = make_pipeline(StandardScaler(), LogisticRegression(max_iter=1000))
cv = GroupKFold(n_splits=FOLDS)

# Real labels
scores = cross_val_score(model, X, y, groups=groups, cv=cv,
                         scoring="balanced_accuracy")
print(f"Balanced accuracy: {scores.mean():.2f} (+/- {scores.std():.2f})")

# Control: the same test with shuffled labels
y_shuffled = np.random.default_rng(1).permutation(y)
control = cross_val_score(model, X, y_shuffled, groups=groups, cv=cv,
                          scoring="balanced_accuracy")
print(f"Shuffled labels:   {control.mean():.2f} (+/- {control.std():.2f})")

# Where do the mistakes go?
pred = cross_val_predict(model, X, y, groups=groups, cv=cv)
labels = sorted(set(y))
print(labels)
print(confusion_matrix(y, pred, labels=labels))

python train.py

where:

GroupKFold — splits the data into folds without breaking clusters. Each fold takes a turn as the validation set, so the result doesn’t depend on one lucky split
StandardScaler — puts all 320 dimensions on the same scale
LogisticRegression — a deliberately simple predictor. If the signal is in the embeddings, a simple model finds it
balanced_accuracy — fair when one habitat has many more proteins than the other

Want a small neural network instead? Swap in MLPClassifier(hidden_layer_sizes=(128,)) from sklearn.neural_network. Start simple, and move up only if it helps.

Checking the results

Balanced accuracy: 0.xx (+/- 0.xx)
Shuffled labels:   0.50 (+/- 0.xx)
['deep', 'surface']
[[... ...]
 [... ...]]

Success looks like this:

The real score is clearly above the shuffled one. With two balanced classes, shuffled labels should sit near 0.5.
The confusion matrix shows errors in both directions, not just one.

A good score is not the end. We still need to ask what the embeddings are really telling us.

Taxonomy can fake an ecological signal. If one clade lives only in deep water and another only at the surface, the model may just be recognising clades. Test it: put whole taxa on one side of the split, not just whole clusters.

Public training data is biased. PLMs have seen mostly well-studied organisms. Proteins from poorly studied environments may be embedded less well, and we can’t easily measure by how much.

Habitat is many factors at once. Depth also means temperature, pressure, light and nutrients. A hit doesn’t tell us which one the model has learnt.

Build a baseline without any PLM: use amino acid composition (20 numbers per protein) and run the same train.py. If composition scores almost as well, the embeddings add little. This happened in published work on optimal growth temperature: the embeddings improved on composition only modestly.

from Bio import SeqIO
import numpy as np

AA = "ACDEFGHIKLMNPQRSTVWY"

def composition(seq):
    seq = str(seq)
    return [seq.count(a) / len(seq) for a in AA]

records = {r.id: r.seq for r in SeqIO.parse("proteins.faa", "fasta")}
ids = [line.strip() for line in open("embeddings.ids.txt")]
X_comp = np.array([composition(records[i]) for i in ids])
np.save("composition.npy", X_comp)

Then set EMB = "composition.npy" in train.py and compare the two scores.

Optional extras

Everything below is optional.

See the embeddings

A PCA plot is a quick look at the structure in the data. Two axes rarely tell the whole story. A predictor can find signal in 320 dimensions even when the plot looks like a blob.

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA

X = np.load("embeddings.npy")
ids = [line.strip() for line in open("embeddings.ids.txt")]
meta = pd.read_csv("metadata.tsv", sep="\t").set_index("protein_id")
y = meta.loc[ids, "habitat"].values

pcs = PCA(n_components=2).fit_transform(X)
for habitat in sorted(set(y)):
    sel = y == habitat
    plt.scatter(pcs[sel, 0], pcs[sel, 1], s=5, label=habitat)
plt.legend()
plt.xlabel("PC1")
plt.ylabel("PC2")
plt.savefig("embeddings_pca.png", dpi=150)

Watch masked prediction

This is how the PLM was trained. We hide one residue and ask the model to guess it.

import torch
from transformers import AutoTokenizer, EsmForMaskedLM

MODEL = "facebook/esm2_t6_8M_UR50D"
SEQ = "MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVKALPDAQFEVV"
POS = 10   # 0-based position to hide

tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = EsmForMaskedLM.from_pretrained(MODEL).eval()

masked = SEQ[:POS] + tokenizer.mask_token + SEQ[POS + 1:]
inputs = tokenizer(masked, return_tensors="pt")
with torch.no_grad():
    logits = model(**inputs).logits

# +1 because of the <cls> token at the start
probs = logits[0, POS + 1].softmax(dim=-1)
top = probs.topk(5)
for p, idx in zip(top.values, top.indices):
    print(tokenizer.decode(idx), f"{p:.2f}")

print("Real residue:", SEQ[POS])

If the top guess is the real residue, the model has learnt something about the context. If it’s not, look at which amino acids it prefers. They are often similar to the real one.

The probabilities are the model’s loss function at work. During training, the model is penalised when the real residue gets a low probability.

01 May 2026

« Host removal with Deacon

Microbiome binfies

Protein language models: from sequence to ecology

The idea

Setup

One protein first

Many proteins

Training the predictor

Checking the results

Optional extras

See the embeddings

Watch masked prediction

Explore →