ESM-2 Enzyme Family Classification System: A Production-Grade Fine-Tuning Scheme Based on Protein Language Models

This project introduces a production-grade fine-tuning system for enzyme family classification based on the ESM-2 protein language model, aiming to address the limitations of traditional sequence alignment methods in classifying distantly homologous proteins. The system integrates key technologies such as LoRA parameter-efficient fine-tuning, homology-aware data splitting, temperature scaling calibration, and integrated gradient interpretability to realize a complete workflow from data processing and model training to production deployment. It provides a reliable solution for enzyme function annotation and can be applied in fields like drug discovery and synthetic biology.

Enzyme family classification is a fundamental task in computational biology and drug discovery, crucial for target identification and selectivity analysis. Traditional methods (e.g., BLAST, HMM) rely on sequence similarity and perform poorly in classifying distantly homologous proteins (low sequence similarity but similar function and structure). Protein language models (e.g., ESM-2) can capture functional and structural information by learning evolutionary covariation, providing a theoretical basis for classification beyond shallow sequence similarity.

The project focuses on four key scientific questions:

1. Representation Capability Evaluation: The extent to which ESM-2 embedding vectors encode enzyme family identity, and the performance difference between fine-tuned and zero-shot classification;
2. Data Splitting Strategy: How to avoid homology leakage (which leads to inflated performance);
3. Uncertainty Quantification: Whether the model's confidence is well-calibrated and when to trust predictions;
4. Interpretability Analysis: Which sequence positions drive classification decisions and whether they correspond to functionally important residues.

Core Technologies

1. LoRA Parameter-Efficient Fine-Tuning: Freeze ESM-2 pre-trained weights, inject low-rank matrices into q/v projection layers, update only 0.36% of parameters, avoid catastrophic forgetting, and achieve performance close to full fine-tuning;
2. Homology-Aware Splitting: Use MMseqs2 to cluster at a 30% similarity threshold, ensuring sequences from the same cluster are only in one dataset (training/validation/test) to avoid homology leakage;
3. Temperature Scaling Calibration: Calibrate confidence using temperature scaling after training, report ECE and reliability plots, and mark low-confidence sequences;
4. Integrated Gradient Interpretability: Replace unreliable attention weights, generate residue importance attributions that can be cross-validated with MSA conservation.

Dataset Preprocessing

• Source: UniProtKB/Swiss-Prot (about 570,000 proteins with EC annotations);
• Filtering: Retain sequences with EC numbers, length between 50-1024 residues, and deduplicate at 100% similarity;
• Splitting: After clustering with MMseqs2 at a 30% similarity threshold, split into training/validation/test sets in a 70%/15%/15% ratio.

Training Configuration

• Tools: HuggingFace Transformers, PEFT library;
• Optimization: AdamW, cosine learning rate, gradient clipping, early stopping (based on validation set macro F1);
• Tracking: MLflow records hyperparameters and experimental results.

Production Deployment

The system provides an inference endpoint via FastAPI, supports Docker packaging, and can be deployed to production environments. The endpoint receives protein sequences and returns classification results, calibrated confidence, and attribution visualizations.

Application Prospects

Applicable to:

• Target identification and selectivity analysis in new drug discovery;
• Enzyme engineering in synthetic biology;
• Enzyme function annotation in metagenomics;
• Function prediction in protein design.

This system combines the strong representation of PLMs with rigorous machine learning practices to realize the transformation of cutting-edge research into a reliable production system.