MTL-Peptide-Classifier

MTL Peptide Classifier (22 Tasks)

Multi-Task Learning peptide classifier covering 22 binary peptide-activity tasks. Built on a frozen ESM-2 (650M) backbone with a parallel Transformer + CNN feature extractor and per-task heads, following a PDeepPP-inspired design.

Held-out Test Set Performance (Averaged across 22 tasks)

Best Val Avg F1 (used for checkpoint selection): 85.56%

Per-Task Test Metrics

Architecture

• Shared encoder: frozen ESM-2 (facebook/esm2_t33_650M_UR50D, 650M params) + learnable base embedding, mixed at esm_ratio=0.9
• Feature extraction (parallel): 4-layer Transformer + CNN (kernel=7, padding=3) → concatenated to 2560-dim features
• Heads: 22 binary classifiers (2560 → 256 → 128 → 2) with masked average pooling
• Loss: TIM (Threshold-Independent Multi-task) loss + label smoothing 0.1

Tasks

Usage

import os
from huggingface_hub import hf_hub_download
import torch
from transformers import EsmTokenizer

from mtl_peptide_classifier import MTLPeptideClassifier, get_all_peptide_tasks

REPO = "tuankg1028/MTL-Peptide-Classifier"
checkpoint_dir = "MTL-Peptide-Classifier"
os.makedirs(checkpoint_dir, exist_ok=True)

for fname in ["heads.pt", "shared_backbone.pt", "ablation_config.json"]:
    hf_hub_download(repo_id=REPO, filename=fname, local_dir=checkpoint_dir)

tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
task_configs = get_all_peptide_tasks("datasets")  # needs local datasets/ dir for task names

model = MTLPeptideClassifier(
    task_configs=task_configs,
    hidden_dim=1280,
    esm_ratio=0.9,
    num_transformer_layers=4,
    dropout=0.3,
    use_transformer=True,
    use_cnn=True,
    unfreeze_esm=False,
)

device = "cuda" if torch.cuda.is_available() else "cpu"
backbone = torch.load(f"{checkpoint_dir}/shared_backbone.pt", map_location=device)
heads = torch.load(f"{checkpoint_dir}/heads.pt", map_location=device)

model.base_embed.load_state_dict(backbone["base_embed"])
if "transformer" in backbone:
    model.transformer.load_state_dict(backbone["transformer"])
if "cnn" in backbone:
    model.cnn.load_state_dict(backbone["cnn"])
    model.layer_norm.load_state_dict(backbone["layer_norm"])
for name, head in model.heads.items():
    if name in heads:
        head.load_state_dict(heads[name])

model = model.to(device).eval()

sequence = "MKWVTFISLLFLFSSAYSRGVFRR"
tokens = " ".join(list(sequence))
inputs = tokenizer(tokens, return_tensors="pt", max_length=128, padding="max_length", truncation=True)
with torch.no_grad():
    logits = model(inputs["input_ids"].to(device), inputs["attention_mask"].to(device), task_name="Antimicrobial")
    probs = torch.softmax(logits, dim=-1)

Training

• Base model: facebook/esm2_t33_650M_UR50D (frozen)
• Batch size: 16, learning rate: 1e-4, 50 epochs, dropout: 0.3
• 3-way split per task: 80% train / 20% val (checkpoint selection) / held-out test CSV evaluated once
• Mixed precision, gradient clipping 1.0, cosine LR with 5 warmup epochs
• TIM loss + label smoothing 0.1

Files

• heads.pt — per-task classifier heads
• shared_backbone.pt — base embedding, Transformer, CNN, LayerNorm
• ablation_config.json — architecture configuration for reproducibility
• test_results.json — held-out test metrics (per task + averages)
• mtl_peptide_classifier.py — model code

Requirements

torch>=2.0.0
transformers>=4.30.0
huggingface_hub
numpy
pandas
scikit-learn