Fine-tuning ESM-2 for protein subcellular localization: LoRA matches full fine-tuning with 2.4% of the parameters

The full article is on Medium. This is the short version.

Where a protein localises in the cell is fundamental to understanding what it does, and determining it experimentally is slow and expensive. I fine-tuned ESM-2 for 10-class subcellular localization on the DeepLoc 2.0 data to see how much of that can be read off the sequence alone, and how cheaply.

The head is the same throughout — ESM-2 backbone, mean pooling, LayerNorm, linear, GELU, linear, 10 classes — so that four training strategies could be compared under one harness: a linear probe on a frozen 8M model, full fine-tuning of the 35M and 150M models, and rank-16 LoRA on the query, key and value projections of the 650M model.

Results

Rank-16 LoRA on the 650M model reached 76.5% accuracy (MCC 0.704) while training 2.4% of the parameters. Full fine-tuning of the 150M model reached 76.6% (MCC 0.706). Four times the backbone bought nothing at this data size.

Two things beside the headline are worth keeping:

1. The linear probe on the frozen 8M model already gives 69.6%. The pretrained representation encodes localization-relevant features before any adaptation at all.
2. At roughly 17,000 training examples, the ceiling looks like the dataset rather than the model. More data or multi-task training is the lever here, not more parameters.

One practical detail

ESM-2 names its attention modules query, key and value, not q_proj, k_proj and v_proj the way LLaMA-style models do. LoRA target modules have to be configured accordingly. That cost me an afternoon.

Training ran on an NVIDIA DGX Spark. The four checkpoints are on the Hugging Face Hub, the code is on GitHub, and there is a Gradio demo you can paste a sequence into.