rustambekurokov/esm2-650m-protein-localization-lora

ESM-2 650M + LoRA for Protein Subcellular Localization

Fine-tuned ESM-2 650M with LoRA for predicting protein subcellular localization (10 classes).

Results

Model	Params Trained	Accuracy	F1 (macro)	MCC
ESM-2 8M linear probe	100% (head only)	69.6%	0.581	0.614
ESM-2 35M full fine-tune	100%	74.3%	0.647	0.677
ESM-2 150M full fine-tune	100%	76.6%	0.696	0.706
ESM-2 650M LoRA	2.4%	76.5%	0.668	0.704

Usage

import torch
from transformers import AutoTokenizer, EsmModel
from peft import LoraConfig, get_peft_model

# Load base model + LoRA weights
tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
# See full inference code in the repository

Labels

ID	Location
0	Cytoplasm
1	Nucleus
2	Extracellular
3	Cell membrane
4	Mitochondrion
5	Endoplasmic reticulum
6	Membrane
7	Golgi apparatus
8	Lysosome/Vacuole
9	Peroxisome

Training

• Dataset: DeepLoc 2.0 (17,266 train / 3,700 val / 3,701 test)
• LoRA config: r=16, alpha=32, target_modules=[query, key, value]
• Training: 10 epochs, lr=2e-4, batch_size=32, cosine schedule
• Hardware: NVIDIA DGX Spark (128GB unified memory)

Citation

@article{lin2023evolutionary,
  title={Evolutionary-scale prediction of atomic-level protein structure with a language model},
  author={Lin, Zeming and Akin, Halil and Rao, Roshan and others},
  journal={Science},
  year={2023}
}