GitHub repo: Embed GitHub

In this project, I fine-tuned ESM2 on the CAFA3 labeled proteins and annotate them with the Gene Ontology labels. The goal is to predict labeled protein functions from an input amino acid sequence. There is a similar paper that executes a similar methodology and comes to similar conclusions. The hardware used in this includes an NVIDIA 4070ti GPU with 12GB VRAM, Intel Core i7-13700KF CPU with 24 threads. The operating system is Ubuntu 24.04. The pipeline is as follows:

input any arbitrary protein sequence → ESM-2 tokenizer → classifier model → classification profile of protein with GO annotation

General outline

1. prepare the data
1. load protein sequences from CAFA3 (CAFA3 provides labeled examples: the sequence + experimentally verified GO functions)
2. load GO term annotations
3. create synthetic data for improved training
4. split data into train/test/validation sets
2. extract ESM2 features
1. load pretrained ESM2 into RAM
2. tokenize protein sequences from CAFA3 using the ESM2 tokenizer
3. extract mean embeddings from ESM2’s last hidden layer
4. store embeddings as features for each protein
3. create classifier model architecture
1. MLP on top of ESM2 embeddings
2. input: ESM2 1280 dimension embeddings
3. hidden layers: 512 → 256 → num_targets
4. output: binary predictions for each GO term (multi-label classification)
4. train
1. loss function: sigmoid binary cross-entropy
2. optimizer: Adam with learning rate 0.001
3. training loop: 2000 steps with an eval every 100 steps
4. batch size: 32
5. run evaluations
1. metrics: accuracy, precision, auPRC, auROC
2. per-function evaluation: calculate metrics fro each GO term
3. final test: evaluate on held-out set
6. iterate, if needed
7. extrapolate to new datasets
1. tbd, but arguably the most important part to see if it is actually valuable

Data preprocessing

We start with the protein sequence dataset (in this case, CAFA3), then pass those inputs to the ESM2 tokenizer.

After running the protein sequence datasets through the ESM2 tokenizer, we then need to convert those token embeddings into sequence embeddings to pass into our hand-built classifier.

After embedding our protein sequences into a data format that our classifier can work with (a concise, fixed-size embedding for each protein sequence), we then need to use these embeddings to predict the protein's functions.

*notes on data amount

number of CAFA3 protein sequences: ~1,800-2,000

• train split: ~1,440-1,600 proteins (80%)
• validation split: ~180-200 proteins (10%)
• test split: ~360-400 proteins (20%)

total number of GO terms: 303 

but what really matters is the meaning we can get out of the dataset rather than raw bytes size

Training steps

1. prepare data
1. take CAFA3 protein sequences, map each to its GO function terms
2. convert terms into multi-hot vectors
3. split data into train/validation/test
2. tokenize
1. convert CAFA3 sequences into ESM2 token IDs (padding/truncation handled by Triton kernel)
3. embed
1. feed tokenized sequences into ESM2 and get per-sequence embeddings
4. train head:
1. pass embeddings into a linear head
2. compare outputs against GO vectors using BCEWithLogits
3. update head weights
5. validate and save
1. evaluate on CAFA3 validation split (F-max metric, auRPC, auROC)
2. store trained head, tokenizer, label map

Inference steps

1. tokenization
1. the raw user-provided string (“MKT…”) is converted into tokens the ESM2 model understands
2. embedding
1. ESM2 turns those tokens into a vector representation (hidden states learned during that model’s training)
3. head projection
1. the trained linear head maps that embedding to scores (logits) over GO terms
4. prediction
1. apply sigmoid to get probabilities for each GO term
5. interpretation
1. the top terms are the protein’s predicted functions

Results

In summary, we fine tuned ESM2-150M on CAFA3 and GO term datasets, built an MLP classifier on those CAFA3 protein sequences for function prediction. We implemented a Triton kernel for optimized sequence embedding pooling and transfer learning from the frozen ESM2-150M embeddings, then a trained MLP head for the 303 GO function terms.

We used 2,000 training steps and achieved decent performance: validation AUPRC 0.583, test AUPRC 0.591, and AUROC ~0.94 across both splits. The model excels at specific functions like deubiquitinase activity (AUPRC 0.976) and G protein-coupled receptor activity (AUPRC 0.915). We had 98.2% overall accuracy with consistent generalization between validation and test sets.

This obviously needs way more data to validate, specifically in a business/research context.

Cloud deployment

The model size ended up being relatively small at 2.05 MB, which is efficient for deployment. This means it can be deployed in AWS Lambda, which is extremely cost effective.

If we want to deploy both the ESM2-150M model (~600MB) for embedding generation and the trained MLP classifier (2.05MB) for function prediction, we could use ECS Fargate, specifically with 2 vCPUs and 8GB RAM, which would cost ~$50-80/month. The customer would send a request to POST /predict with {"sequences": ["MKLLILTCLVAV..."]}.

It would be better to host both, since we could do raw protein sequences → function predictions in one API call (on unseen proteins). No preprocessing would be required since the users just send protein sequences, we could optimize caching to get much faster for every subsequent call, and we can retrain the MLP without changing the API. But as an MVP, we could just host the Lambda function.

Figures