QuantaFold finetunes a model for protein function identification from a 5-second prediction, achieving 97.9% accuracy on the 1,000 most common protein families while democratizing access to computational biology tools. Inspired by AlphaFold
Hugging Face Model Card: https://huggingface.co/Tarive/esm2_t12_35M_UR50D-finetuned-pfam-1k
Hugging Face Dataset - 100k - https://huggingface.co/datasets/Tarive/within_family_test_set
Hugging Face Dataset - 400k - https://huggingface.co/datasets/Tarive/quantafold
Weights and Biases Tracker - 400k model - https://wandb.ai/tarive22-shivoham/huggingface/runs/n0g2n290?nw=nwusertarive22
Weights and Biases Tracker - 100k model - https://wandb.ai/tarive22-shivoham/huggingface/runs/vklqmoh2?nw=nwusertarive22
The Challenge: Protein function identification is a critical bottleneck in drug discovery and biological research. Making AlphaFold grade models availabe to researchers with less amount of resources availabe.
The Solution: QuantaFold is a complete end-to-end system that fine-tunes the lightweight ESM-2 model to classify proteins into 5,000 functional families based solely on amino acid sequences, running efficiently on a single GPU while maintaining research-grade accuracy.
| Model | Dataset Size | Families | Training Time | Status | Accuracy |
|---|---|---|---|---|---|
| Specialist Model | 1K balanced samples | 1,000 top families | 45 minutes | ✅ Completed | 97.9% |
| Optimized Generalist | 70K stratified samples | 5,000 families | ~3 hours | ✅ Completed | 98.4% |
| Full-Scale Generalist | 400K balanced samples | 5,000 families | ~4 hours | ✅ Completed | 98.1% |
- Training Time Reduction: From impossible 19+ hours → manageable 3-4 hours (80%+ reduction)
- Memory Efficiency: 50% reduction through FP16 mixed-precision training
- Dataset Optimization: Intelligent stratified sampling (400K → 70K) while preserving all 5,000 families
- Parallel Training Strategy: Running multiple model variants to compare optimization impact
- Workflow Acceleration: 10,000x speedup from days of manual analysis to seconds of automated prediction
- Source: Google AI Pfam Dataset on Kaggle
- Scale: ~1.34 million protein sequences across 17,929 families
- Structure: Curated protein domains with family annotations (family_accession, sequence, aligned_sequence)
- Quality: Gold standard benchmarking dataset used in leading computational biology publications
- Base Model: ESM-2 (Evolutionary Scale Modeling v2) - 35M parameters
- Framework: PyTorch + Hugging Face Transformers
- Fine-tuning Strategy: Classification head adaptation for multi-class protein family prediction
- Optimization: Custom WeightedTrainer to handle severe class imbalance (5,000 families)
Balanced Dataset Creation (400K sequences from 1.34M original):
- Original dataset: 1,339,083 sequences across 17,929 families
- Optimized dataset: 400,000 sequences across 5,000 families
- Size reduction: 3.3x smaller dataset, 3.3x faster training
- Method: Intelligent stratified sampling
- Top 1,000 families: 200 sequences each (200,000 total)
- Next 4,000 families: 50 sequences each (200,000 total)
- Perfect balance maintaining statistical representation
Figure 1: Comprehensive dataset balancing strategy showing family distribution and sequence count optimization
Key Findings from Length Distribution Analysis:
- Mean sequence length: 154 amino acids
- Median sequence length: 119 amino acids
- Long sequence outliers: >381 amino acids identified as computational bottlenecks
- Optimization strategy: Strategic truncation balancing information retention vs. efficiency
Figure 2: Complete sequence length analysis revealing optimization opportunities and truncation strategy
# Performance optimizations implemented:
- FP16 Mixed Precision Training (50% memory reduction)
- 8-bit AdamW Optimizer (bitsandbytes)
- Gradient Accumulation for effective large batch training
- GPU memory optimization techniques# Custom WeightedTrainer implementation
class WeightedTrainer(Trainer):
def compute_loss(self, model, inputs, return_outputs=False):
# Implemented inverse frequency weighting
# Forces model attention to rare protein families
# Critical for building robust generalist model- Deep Learning: PyTorch, Hugging Face Transformers, Accelerate
- Data Processing: pandas, scikit-learn, Hugging Face Datasets
- Optimization: bitsandbytes (8-bit optimization), FP16 mixed precision
- Monitoring: Weights & Biases for experiment tracking
- Web App: Gradio interactive interface
- Hosting: Hugging Face Spaces & Hub
- Accessibility: Zero-code-required user experience
- Hardware: NVIDIA A100 GPU (primary), T4 GPU (testing)
- Tutorials: Hugging Face Protein Language Modeling Guide
- Notebook Reference: Protein Language Modeling Colab
Objective: Validate core approach with balanced dataset
Implementation:
- Dataset: 1,000 most common protein families (balanced sampling)
- Training time: 45 minutes on A100
- Result: ✅ 97.9% accuracy - validating core methodology
Key Learning: High accuracy achieved, but model completely failed on proteins outside training families - revealing critical limitation for real-world deployment.
Objective: Build robust models handling real-world data distribution
Critical Challenges Identified:
- Severe Class Imbalance: Real protein data follows power law distribution
- Computational Feasibility: Original 19+ hour training time estimate
- Memory Constraints: GPU memory limitations with large batches
Engineering Solutions & Parallel Training Strategy:
- Smart Data Curation: Stratified sampling preserving all 5,000 families
- Training Status: 🔄 Currently training (~3 hours estimated)
- Optimization: WeightedTrainer + FP16 + 8-bit optimizers
- Balanced Dataset: Strategic 400K sample curation from 1.34M original
- Training Status: 🔄 Currently training (~4 hours estimated)
- Purpose: Maximum balanced data utilization with optimized class distribution
Advanced Training Pipeline (Both Models):
- Implemented WeightedTrainer for class imbalance
- FP16 mixed precision for memory efficiency
- 8-bit optimizers for speed optimization
- Comprehensive W&B logging for performance comparison
Objective: Deploy completed model and analyze training results
Current Implementation:
- Gradio web interface (deployed with 97.9% model)
- Hugging Face Spaces hosting
- Real-time monitoring of parallel training runs
- Comparative analysis preparation for final results
[Include Weights & Biases screenshots showing:]
- Training/Validation Loss Curves: Demonstrating stable convergence over 45 minutes
- Final Accuracy Metrics: 97.9% validation accuracy achievement
- Class Distribution: Balanced 1K family performance analysis
- Real-time Training Progress: Live loss curves and accuracy tracking
- Memory Optimization Impact: GPU utilization efficiency gains
- Class Imbalance Handling: WeightedTrainer performance on rare families
- Comparative Training Metrics: Side-by-side with 70K model
- Resource Utilization: Full dataset computational requirements
- Convergence Analysis: Training stability with maximum data
- Training Speed vs. Data Size: Performance scaling relationships
- Optimization Impact: FP16 and 8-bit optimizer effectiveness
- Accuracy vs. Efficiency Trade-offs: Comprehensive performance matrix
[Include W&B visualizations showing:]
- Confusion Matrix: Per-family classification performance
- Accuracy by Family Size: Performance correlation with training data availability
- Inference Speed Benchmarks: Latency analysis across different sequence lengths
- Resource Utilization: GPU memory and compute efficiency metrics
- Accessibility: Enables small labs and startups to perform advanced protein analysis without supercomputing infrastructure
- Cost Reduction: Eliminates need for expensive computational resources
- Speed: Accelerates research timelines from weeks to minutes
- Target Identification: Rapid hypothesis generation for new protein functions
- Pipeline Optimization: Reduces R&D bottlenecks in pharmaceutical development
- Academic Research: Enables broader participation in computational biology research
- Open Source: Freely available tools and methodologies
- Reproducible Research: Documented approach enabling further research
- Educational Resource: Demonstrates practical ML engineering for biology
- Efficient Fine-tuning Pipeline: Optimized ESM-2 adaptation for large-scale classification
- Class Imbalance Solution: Custom weighted training approach for biological data
- Computational Optimization: Advanced techniques reducing training time by 84%
- Deployment Strategy: User-friendly interface bridging research and application
- Professional-Grade Workflow: Complete ML lifecycle from data analysis to deployment
- Performance Optimization: Multiple levels of computational efficiency improvements
- Scalable Architecture: Design supporting future expansion to larger protein databases
- Quality Assurance: Rigorous validation using established benchmarks
- Dataset: Pfam Seed Random Split - Google AI
- Model: ESM-2 by Meta AI
- Tutorial: Deep Learning with Proteins - Hugging Face
- Implementation Guide: Protein Language Modeling Notebook
- ESM-2 Paper: "Language models enable zero-shot prediction of the effects of mutations on protein function" (Meta AI)
- Pfam Database: "The Pfam protein families database" (Nucleic Acids Research)
- Benchmark Reference: "Can Deep Learning Classify the Protein Universe?" (Bileschi et al.)
- Ensemble Methods: Combining multiple model architectures for improved accuracy
- Active Learning: Intelligent selection of proteins for manual annotation
- Multi-task Learning: Simultaneous prediction of function, structure, and interactions
- Real-time Analysis: Integration with laboratory sequencing workflows
- Collaborative Platform: Community-driven protein annotation system
- Commercial Applications: Licensed solutions for pharmaceutical R&D
@misc{adhsaksham,
author = {Adhikari, Saksham and Sharma, Kusum},
year = {2026},
month = {09},
pages = {4},
title = {QuantaFold: Scaling Protein Language Model Fine-tuning to 5,000 Families Through Systematic Optimization}
}- Google TPU Research Cloud for compute access