Transfer Learning From Protein Language Models

The characteristic protein machine-learning problem is 400 labelled sequences and a property you want to predict. Pretrained embeddings make that tractable, and the parameter arithmetic explains exactly why.

The setting: a few hundred labels

Labels in this field come from experiments, so they are expensive and few. A binding assay across a designed library might return 400 measurements. A stability dataset might have 2,000. Deep mutational scanning is the exception at tens of thousands, and it covers one protein.

Training a sequence model from scratch on 400 examples is not possible. The transfer-learning move is to let a model that has already learned the statistics of protein sequences do the representation, and fit only a small function on top of it. What that pretrained model has absorbed, and from what corpus, is the subject of what a protein language model is trained on; here we take it as given and use it.

From residues to one vector

A forward pass over a sequence of length L returns an L-by-d matrix of per-residue representations. For the 650-million-parameter, 33-layer ESM-2 checkpoint, d is 1,280. Your label is a property of the whole protein, so you need one vector, and pooling is where several silent errors live.

• Exclude the special tokens. The tokenizer prepends a beginning-of-sequence token and may append an end-of-sequence and padding tokens. Their representations are in the matrix and they are not residues. Averaging over the raw tensor including padding makes the pooled vector a function of batch composition, which is a bug that produces plausible numbers and irreproducible results.
• Mean over residues is the default and it is length-biased in a subtle way. Longer proteins average over more vectors, so the pooled representation of a long protein is closer to the corpus mean than that of a short one. If your labels correlate with length — and in most curated datasets they do — a model can exploit this. Check by fitting on length alone first.
• Which layer you take matters. The final layer is specialised toward the pretraining objective. Intermediate layers often transfer better on downstream tasks, and treating the layer index as a hyperparameter to sweep costs one forward pass per layer because you can extract them all in a single run.
• If the property is local, do not pool. Binding-site or post-translational-modification prediction is a per-residue label, and pooling destroys exactly the signal. Keep the L-by-d matrix and fit a per-residue head.

Frozen features and a small head

Run every sequence through the model once, save the pooled vectors, and never touch the model again. Your dataset is now a 400-by-1,280 matrix and a 400-vector of labels, which is an ordinary tabular regression.

embeddings   400 x 1,280   (computed once, cached)
labels       400

ridge regression head:
  parameters  1,280 weights + 1 bias  = 1,281
  examples    400

  -> p > n: regularisation is not optional, it is the
     only thing making the fit well-posed

Compare the two parameter counts and the whole argument becomes arithmetic. The backbone has 650,000,000 parameters, which 400 examples could never constrain. The head has 1,281, which 400 examples constrain badly but not hopelessly — provided you regularise, because with more parameters than examples the unregularised least-squares solution is not unique. Ridge regression with the penalty strength chosen by cross-validation is the right default, and for classification, logistic regression with an L2 penalty.

Two baselines must be reported alongside it, and the second is the one people skip. The first is a constant predictor, which tells you the variance you are working against. The second is the same head fitted on a cheap sequence representation — one-hot encoding, amino-acid composition, or a BLOSUM-based encoding. On small datasets over closely related sequences, the cheap representation sometimes wins, and a paper or a pipeline that never ran it cannot know whether the pretrained model contributed anything.

Dimension reduction is worth a moment too. With 1,280 features and 400 examples, projecting to the first 50 principal components of the embeddings before fitting often improves generalisation and always makes the model easier to reason about; see what PCA is computing.

When to unfreeze anything

Full fine-tuning updates the backbone and needs enough data to justify moving 650 million parameters. The intermediate options, in increasing order of data appetite:

1. Frozen embeddings, linear head. Hundreds of examples. Trains in seconds on a CPU. Start here always, because it is the baseline every other option has to beat.
2. Frozen embeddings, small non-linear head. A one- or two-hidden-layer network. Worth trying at low thousands of examples; below that it usually overfits without beating the linear model.
3. Low-rank adaptation of the backbone. Insert small trainable low-rank matrices alongside frozen weights, training a tiny fraction of the parameter count. This adapts the representation without the memory cost of full fine-tuning and is the first option that can actually change what the features are.
4. Unfreeze the last few layers. Or the whole model, with a learning rate around 1e-5 — far lower than you would use training from scratch, because large updates destroy the pretrained representation faster than the small dataset can rebuild it. Use early stopping on a validation set and expect to need thousands of labels.

Cost matters in this ordering. Step 1 requires one forward pass over your dataset, ever. Step 4 requires forward and backward passes over the backbone every epoch, which is several orders of magnitude more compute and needs an accelerator with enough memory for the activations. Going straight to step 4 is common and is usually not justified by the gain.

The split that decides whether it is real

This is the part that makes or breaks the number you report, and it has nothing to do with the model.

Protein datasets are full of homologues. A random train/test split puts near-identical sequences on both sides, so the test set is answerable by recalling a neighbour from training. Reported accuracy under a random split can be very high and mean nothing about a new protein family, which is the case you actually care about.

The correction is to cluster the sequences by identity and split by cluster, so no cluster spans the boundary:

1. Cluster all 400 sequences at a stated identity threshold with a sequence-clustering tool. Thirty percent is the conventional strict setting; 50 percent is a common compromise.
2. Assign whole clusters to train, validation and test folds, never individual sequences.
3. Report the threshold with the result. “Spearman 0.61 at a 30 percent identity split” is a claim. “Spearman 0.61” is not.
4. Report the random-split number too, so the gap between them is visible. A large gap is not a failure — it is the honest measurement of how much of the apparent performance was homology.

Two further evaluation details. With 400 examples, a single held-out fold gives an estimate with a wide confidence interval, so use repeated grouped cross-validation and report the spread rather than a point. And choose the metric to match the use: if the model is going to rank candidates for synthesis, as in the design funnel, a rank correlation or the mean label among the top k is the number that matters, and R-squared over the whole range can look poor while the ranking at the top is excellent, or the reverse.

Finally, a limit that no split fixes: the model predicts the assay it was trained on, under those conditions. A stability predictor fitted to melting temperatures measured in one buffer is a predictor of melting temperature in that buffer. Extrapolating it to behaviour in a cell, or to a therapeutic property, is an assumption that needs its own evidence.