27.5 Building a Protein Property Predictor

"You give me a sequence, I give you a melting temperature. It took evolution four billion years to figure out thermostability. It took me three epochs and a rank-8 adapter."

A LoRA Adapter With Delusions of Grandeur

1. The Meltome Atlas Dataset

Can a neural network, given nothing but a string of amino acids, predict the temperature at which a protein falls apart? The Meltome Atlas (Jarzab et al., 2020) makes that question testable: it contains melting temperatures for 48,000 proteins across 13 organisms, measured by thermal proteome profiling (TPP). In TPP, cells are heated to a range of temperatures, and the fraction of soluble protein remaining at each temperature is quantified by mass spectrometry. The melting temperature \(T_m\) is the temperature at which 50% of the protein has denatured.

Thermal proteome profiling (TPP) measures protein stability across an entire proteome in a single experiment, rather than purifying and testing one protein at a time. TPP provides the large, organism-scale labeled datasets that supervised machine learning requires. Before TPP, melting temperatures existed for only a few hundred well-studied proteins. The technique heats replicate cell lysates to a series of temperatures (typically 37 to 67 degrees Celsius in roughly ten steps). Centrifugation pellets the aggregated proteins, and tandem mass spectrometry quantifies the remaining soluble fraction. Use TPP-derived datasets like the Meltome Atlas when you need broad coverage across many protein families; use differential scanning calorimetry (DSC) or circular dichroism (CD) when you need precise, per-protein thermodynamic parameters such as enthalpy of unfolding.

Melting temperature is a proxy for thermal stability: proteins with higher \(T_m\) resist heat-induced unfolding more effectively. Enzyme engineering depends on this property, because industrial enzymes must survive high process temperatures. Protein therapeutics demand it too, since biologics must remain stable during storage. Fundamental biology also relies on \(T_m\): thermophilic organisms have systematically higher values. In short: a few megabytes of learned adapter weights, grafted onto a pretrained protein language model, turn raw amino acid strings into quantitative stability predictions.

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split

# Load the Meltome dataset
# Available from: https://github.com/J-SNACKKB/FLIP/tree/main/splits/meltome
# We use a preprocessed version with sequences and T_m values
meltome = pd.read_csv("meltome_data.csv")

print(f"Total proteins: {len(meltome):,}")
print(f"Columns: {list(meltome.columns)}")
print(f"\nT_m distribution:")
print(f"  Mean:   {meltome['melting_temp'].mean():.1f} C")
print(f"  Std:    {meltome['melting_temp'].std():.1f} C")
print(f"  Min:    {meltome['melting_temp'].min():.1f} C")
print(f"  Max:    {meltome['melting_temp'].max():.1f} C")
print(f"  Median: {meltome['melting_temp'].median():.1f} C")

# Filter: keep proteins with sequence length <= 1024 (ESM-2 context window)
meltome = meltome[meltome["sequence"].str.len() <= 1024].copy()
print(f"\nAfter length filter (<= 1024): {len(meltome):,} proteins")

# Sequence length distribution
seq_lengths = meltome["sequence"].str.len()
print(f"Sequence length: mean={seq_lengths.mean():.0f}, "
      f"median={seq_lengths.median():.0f}, "
      f"max={seq_lengths.max()}")

# Train/validation/test split (80/10/10)
train_df, temp_df = train_test_split(meltome, test_size=0.2, random_state=42)
val_df, test_df = train_test_split(temp_df, test_size=0.5, random_state=42)

print(f"\nSplit sizes: train={len(train_df)}, "
      f"val={len(val_df)}, test={len(test_df)}")

Total proteins: 28,790
Columns: ['sequence', 'melting_temp', 'organism']

T_m distribution:
  Mean:   52.3 C
  Std:    8.7 C
  Min:    28.1 C
  Max:    86.4 C
  Median: 51.8 C

After length filter (<= 1024): 26,412 proteins
Sequence length: mean=387, median=332, max=1024

Split sizes: train=21129, val=2641, test=2642

2. Building the Dataset and DataLoader

PyTorch's Dataset and DataLoader classes handle batching, shuffling, and tokenization. For protein sequences of varying lengths, we need a custom collation function (a callback that assembles individual samples into a single batch tensor, handling padding and alignment) that pads sequences within each batch to the same length.

import torch
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer

class MeltomeDataset(Dataset):
    """PyTorch dataset for protein thermostability prediction."""

    def __init__(self, sequences: list, temperatures: list,
                 tokenizer, max_length: int = 1026):
        self.sequences = sequences
        self.temperatures = torch.tensor(temperatures, dtype=torch.float32)
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.sequences)

    def __getitem__(self, idx):
        encoding = self.tokenizer(
            self.sequences[idx],
            truncation=True,
            max_length=self.max_length,
            padding=False,            # pad at collation time
            return_tensors="pt"
        )
        return {
            "input_ids": encoding["input_ids"].squeeze(0),
            "attention_mask": encoding["attention_mask"].squeeze(0),
            "labels": self.temperatures[idx]
        }

def collate_fn(batch):
    """Pad sequences to the longest in the batch."""
    max_len = max(item["input_ids"].size(0) for item in batch)

    input_ids = torch.stack([
        torch.nn.functional.pad(
            item["input_ids"],
            (0, max_len - item["input_ids"].size(0)),
            value=1  # pad token id = 1
        )
        for item in batch
    ])
    attention_mask = torch.stack([
        torch.nn.functional.pad(
            item["attention_mask"],
            (0, max_len - item["attention_mask"].size(0)),
            value=0
        )
        for item in batch
    ])
    labels = torch.stack([item["labels"] for item in batch])

    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels
    }

# Create datasets and loaders
tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")

train_dataset = MeltomeDataset(
    train_df["sequence"].tolist(),
    train_df["melting_temp"].tolist(),
    tokenizer
)
val_dataset = MeltomeDataset(
    val_df["sequence"].tolist(),
    val_df["melting_temp"].tolist(),
    tokenizer
)
test_dataset = MeltomeDataset(
    test_df["sequence"].tolist(),
    test_df["melting_temp"].tolist(),
    tokenizer
)

train_loader = DataLoader(
    train_dataset, batch_size=8, shuffle=True, collate_fn=collate_fn
)
val_loader = DataLoader(
    val_dataset, batch_size=16, shuffle=False, collate_fn=collate_fn
)
test_loader = DataLoader(
    test_dataset, batch_size=16, shuffle=False, collate_fn=collate_fn
)

# Verify
batch = next(iter(train_loader))
print(f"Batch input_ids shape: {batch['input_ids'].shape}")
print(f"Batch attention_mask shape: {batch['attention_mask'].shape}")
print(f"Batch labels shape: {batch['labels'].shape}")
print(f"Sample labels: {batch['labels'][:4].tolist()}")

3. Configuring ESM-2 with LoRA

A single laboratory stability assay costs days of technician time and thousands of dollars in reagents; screening a library of ten thousand enzyme variants experimentally is prohibitive. A predictor that converts raw amino acid sequence into a reliable thermostability estimate lets researchers filter candidates computationally, reserving wet-lab validation for only the most promising designs.

Now we assemble the complete model: ESM-2 backbone with LoRA adapters and a regression head, where the regression head is a small feedforward network that maps the model's internal representation to a single numeric output (the predicted T_m). This builds on the LoRA concepts from Section 27.4 and the model pipeline from Section 27.1, using the PEFT library (Parameter-Efficient Fine-Tuning, a Hugging Face toolkit that implements LoRA and related adapter methods) to inject the low-rank adapters into ESM-2. Figure 27.9 illustrates the full architecture from input sequence to predicted melting temperature. Figure 27.5.1 illustrates ESM-2 + LoRA thermostability prediction pipeline.

import torch.nn as nn
from transformers import EsmModel
from peft import LoraConfig, get_peft_model, TaskType

class ThermostabilityPredictor(nn.Module):
    """ESM-2 + LoRA model for protein melting temperature prediction."""

    def __init__(self, model_name: str = "facebook/esm2_t33_650M_UR50D",
                 lora_rank: int = 8, lora_alpha: int = 16,
                 dropout: float = 0.1):
        super().__init__()

        # Load pretrained ESM-2
        self.esm = EsmModel.from_pretrained(model_name)
        hidden_dim = self.esm.config.hidden_size  # 1280 for 650M

        # Apply LoRA
        lora_config = LoraConfig(
            task_type=TaskType.FEATURE_EXTRACTION,
            r=lora_rank,
            lora_alpha=lora_alpha,
            lora_dropout=dropout,
            target_modules=["query", "value"],
            bias="none",
        )
        self.esm = get_peft_model(self.esm, lora_config)

        # Regression head
        self.head = nn.Sequential(
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, 256),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(256, 64),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(64, 1)
        )

        # Print parameter summary
        self.esm.print_trainable_parameters()

    def forward(self, input_ids, attention_mask=None):
        # ESM-2 forward (LoRA-adapted)
        outputs = self.esm(
            input_ids=input_ids, attention_mask=attention_mask
        )
        hidden_states = outputs.last_hidden_state  # (B, L, D)

        # Mean pooling with attention mask
        if attention_mask is not None:
            mask = attention_mask.unsqueeze(-1).float()
            pooled = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
        else:
            pooled = hidden_states.mean(dim=1)

        # Predict T_m
        return self.head(pooled).squeeze(-1)

# Initialize model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = ThermostabilityPredictor(lora_rank=8, lora_alpha=16)
model = model.to(device)

# Count head parameters
head_params = sum(p.numel() for p in model.head.parameters())
total_trainable = sum(
    p.numel() for p in model.parameters() if p.requires_grad
)
print(f"\nHead parameters: {head_params:,}")
print(f"Total trainable (LoRA + head): {total_trainable:,}")
print(f"Device: {device}")

4. The Training Loop

The optimization strategy updates the LoRA adapter weights and the regression head together.

We train with AdamW (Adam with decoupled weight decay, which prevents the weight decay term from interfering with the adaptive learning rate), a cosine learning rate schedule, and Huber loss (a loss function that behaves like mean squared error for small residuals but switches to a linear penalty for large residuals, making it less sensitive to outlier melting temperature measurements that would otherwise dominate gradient updates). The Meltome dataset contains some extreme values (proteins from thermophilic organisms) that can destabilize training with standard mean squared error (MSE) loss.

from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from scipy.stats import spearmanr
import time

def train_epoch(model, loader, optimizer, scheduler, device):
    """Train for one epoch, return mean loss."""
    model.train()
    total_loss = 0
    loss_fn = nn.HuberLoss(delta=5.0)  # robust to T_m outliers

    for batch in loader:
        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels = batch["labels"].to(device)

        optimizer.zero_grad()
        predictions = model(input_ids, attention_mask)
        loss = loss_fn(predictions, labels)
        loss.backward()

        # Gradient clipping
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        optimizer.step()
        scheduler.step()

        total_loss += loss.item() * len(labels)

    return total_loss / len(loader.dataset)

@torch.no_grad()
def evaluate(model, loader, device):
    """Evaluate model, return loss, Spearman rho, and RMSE."""
    model.eval()
    all_preds, all_labels = [], []
    total_loss = 0
    loss_fn = nn.HuberLoss(delta=5.0)

    for batch in loader:
        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels = batch["labels"].to(device)

        predictions = model(input_ids, attention_mask)
        loss = loss_fn(predictions, labels)

        total_loss += loss.item() * len(labels)
        all_preds.extend(predictions.cpu().tolist())
        all_labels.extend(labels.cpu().tolist())

    all_preds = np.array(all_preds)
    all_labels = np.array(all_labels)

    avg_loss = total_loss / len(loader.dataset)
    rho, p_value = spearmanr(all_preds, all_labels)
    rmse = np.sqrt(np.mean((all_preds - all_labels) ** 2))

    return {
        "loss": avg_loss,
        "spearman_rho": rho,
        "rmse": rmse,
        "p_value": p_value
    }

# Training configuration
num_epochs = 5
learning_rate = 5e-4
weight_decay = 0.01

optimizer = AdamW(
    [p for p in model.parameters() if p.requires_grad],
    lr=learning_rate,
    weight_decay=weight_decay
)
total_steps = num_epochs * len(train_loader)
scheduler = CosineAnnealingLR(
    optimizer, T_max=total_steps, eta_min=1e-6
)

# Training loop
print(f"{'Epoch':>5s} | {'Train Loss':>10s} | {'Val Loss':>10s} | "
      f"{'Val rho':>8s} | {'Val RMSE':>9s} | {'Time':>6s}")
print("-" * 65)

best_rho = -1
for epoch in range(num_epochs):
    start = time.time()

    train_loss = train_epoch(
        model, train_loader, optimizer, scheduler, device
    )
    val_metrics = evaluate(model, val_loader, device)

    elapsed = time.time() - start

    print(f"{epoch+1:5d} | {train_loss:10.4f} | "
          f"{val_metrics['loss']:10.4f} | "
          f"{val_metrics['spearman_rho']:8.4f} | "
          f"{val_metrics['rmse']:9.2f} | {elapsed:5.0f}s")

    # Save best model
    if val_metrics["spearman_rho"] > best_rho:
        best_rho = val_metrics["spearman_rho"]
        torch.save(model.state_dict(), "best_thermostability_model.pt")

print(f"\nBest validation Spearman rho: {best_rho:.4f}")

Epoch | Train Loss |   Val Loss |  Val rho |  Val RMSE |   Time
-----------------------------------------------------------------
    1 |     3.8721 |     3.1245 |   0.5823 |      7.41 |   892s
    2 |     2.6543 |     2.5187 |   0.6547 |      6.83 |   885s
    3 |     2.1876 |     2.2341 |   0.6892 |      6.42 |   891s
    4 |     1.8432 |     2.1567 |   0.7034 |      6.18 |   888s
    5 |     1.6218 |     2.1823 |   0.6987 |      6.21 |   890s

Best validation Spearman rho: 0.7034

5. Evaluation and Error Analysis

A single number (Spearman rho) tells you the model works; error analysis tells you where it works and where it fails. The most informative analysis stratifies errors by organism, sequence length, and T_m range.

from scipy.stats import spearmanr
import numpy as np

# Load best model and evaluate on test set
model.load_state_dict(torch.load("best_thermostability_model.pt"))
test_metrics = evaluate(model, test_loader, device)

print(f"Test set results:")
print(f"  Spearman rho: {test_metrics['spearman_rho']:.4f} "
      f"(p = {test_metrics['p_value']:.2e})")
print(f"  RMSE: {test_metrics['rmse']:.2f} C")
print(f"  Huber loss: {test_metrics['loss']:.4f}")

# Collect predictions for error analysis
all_preds, all_labels = [], []
model.eval()
with torch.no_grad():
    for batch in test_loader:
        preds = model(
            batch["input_ids"].to(device),
            batch["attention_mask"].to(device)
        )
        all_preds.extend(preds.cpu().tolist())
        all_labels.extend(batch["labels"].tolist())

all_preds = np.array(all_preds)
all_labels = np.array(all_labels)
errors = all_preds - all_labels

# Error analysis by T_m range
print("\nError analysis by T_m range:")
print(f"{'Range':>15s} | {'Count':>6s} | {'RMSE':>7s} | "
      f"{'Bias':>7s} | {'Rho':>6s}")
print("-" * 55)

ranges = [
    (28, 40, "Cold (<40 C)"), (40, 50, "Mesophilic"),
    (50, 60, "Moderate"), (60, 70, "Thermophilic"),
    (70, 90, "Extreme (>70 C)")
]

for lo, hi, label in ranges:
    mask = (all_labels >= lo) & (all_labels < hi)
    if mask.sum() < 10:
        continue
    rho, _ = spearmanr(all_preds[mask], all_labels[mask])
    rmse = np.sqrt(np.mean((all_preds[mask] - all_labels[mask]) ** 2))
    bias = np.mean(all_preds[mask] - all_labels[mask])
    print(f"{label:>15s} | {mask.sum():6d} | {rmse:7.2f} | "
          f"{bias:+7.2f} | {rho:6.3f}")

Test set results:
  Spearman rho: 0.6923 (p = 0.00e+00)
  RMSE: 6.31 C
  Huber loss: 2.1934

Error analysis by T_m range:
          Range |  Count |    RMSE |    Bias |    Rho
-------------------------------------------------------
    Cold (<40 C) |    187 |    8.42 |   +3.21 |  0.412
     Mesophilic |    892 |    5.87 |   +0.43 |  0.623
       Moderate |    976 |    5.12 |   -0.28 |  0.681
   Thermophilic |    412 |    6.93 |   -1.87 |  0.598
 Extreme (>70 C) |    175 |    9.18 |   -4.32 |  0.389

6. Saving and Deploying the Adapter

The trained adapter can now be packaged for reuse and deployment.

The LoRA adapter weighs only a few megabytes; the 2.5 GB base model is shared across tasks and never duplicated. We save the adapter and regression head alone, then reload them onto a cached base ESM-2 at deployment time.

import os
import json

def save_predictor(model, save_dir: str, metadata: dict = None):
    """Save LoRA adapter and regression head separately."""
    os.makedirs(save_dir, exist_ok=True)

    # Save LoRA adapter (via PEFT)
    model.esm.save_pretrained(os.path.join(save_dir, "lora_adapter"))

    # Save regression head
    torch.save(
        model.head.state_dict(),
        os.path.join(save_dir, "head.pt")
    )

    # Save metadata
    meta = {
        "base_model": "facebook/esm2_t33_650M_UR50D",
        "lora_rank": 8,
        "lora_alpha": 16,
        "target_modules": ["query", "value"],
        "task": "thermostability_regression",
        "metric": "spearman_rho",
        "best_val_score": float(best_rho),
        "test_rmse": float(test_metrics["rmse"]),
    }
    if metadata:
        meta.update(metadata)

    with open(os.path.join(save_dir, "config.json"), "w") as f:
        json.dump(meta, f, indent=2)

    # Report sizes
    adapter_size = sum(
        os.path.getsize(os.path.join(dp, f))
        for dp, _, fns in os.walk(
            os.path.join(save_dir, "lora_adapter")
        )
        for f in fns
    )
    head_size = os.path.getsize(os.path.join(save_dir, "head.pt"))
    print(f"Saved adapter: {adapter_size / 1e6:.1f} MB")
    print(f"Saved head:    {head_size / 1e6:.1f} MB")
    print(f"Total:         {(adapter_size + head_size) / 1e6:.1f} MB")
    print(f"Base model:    ~2,500 MB (shared, not saved)")

save_predictor(model, "thermostability_predictor")

def load_predictor(save_dir: str, device="cpu") -> ThermostabilityPredictor:
    """Load a saved predictor with LoRA adapter and regression head."""
    with open(os.path.join(save_dir, "config.json")) as f:
        config = json.load(f)

    # Reconstruct model
    predictor = ThermostabilityPredictor(
        model_name=config["base_model"],
        lora_rank=config["lora_rank"],
        lora_alpha=config["lora_alpha"],
    )

    # Load LoRA weights
    from peft import PeftModel
    predictor.esm = PeftModel.from_pretrained(
        EsmModel.from_pretrained(config["base_model"]),
        os.path.join(save_dir, "lora_adapter")
    )

    # Load head weights
    predictor.head.load_state_dict(
        torch.load(
            os.path.join(save_dir, "head.pt"),
            map_location=device
        )
    )

    return predictor.to(device).eval()

# Verify round-trip
loaded_model = load_predictor("thermostability_predictor", device=device)
reloaded_metrics = evaluate(loaded_model, test_loader, device)
print(f"Reloaded model Spearman rho: "
      f"{reloaded_metrics['spearman_rho']:.4f}")
print(f"Match: {abs(reloaded_metrics['spearman_rho'] - test_metrics['spearman_rho']) < 1e-4}")

7. Integration with the Discovery Workbench

The thermostability predictor integrates into the Discovery Workbench (introduced in Chapter 6) as a service that accepts protein sequences and returns predicted melting temperatures.

from dataclasses import dataclass
from typing import List, Optional
import torch

@dataclass
class ThermostabilityResult:
    """Result from the thermostability prediction service."""
    sequence: str
    predicted_tm: float
    confidence_category: str  # "high", "medium", "low"

class ThermostabilityService:
    """Discovery Workbench service for protein thermostability prediction.

    Wraps the ESM-2 + LoRA predictor with batched inference,
    confidence estimation, and input validation.
    """

    def __init__(self, model_dir: str, device: str = "cuda"):
        self.model = load_predictor(model_dir, device)
        self.tokenizer = AutoTokenizer.from_pretrained(
            "facebook/esm2_t33_650M_UR50D"
        )
        self.device = device

        # Valid amino acid alphabet
        self.valid_residues = set("ACDEFGHIKLMNPQRSTVWY")

    def validate_sequence(self, sequence: str) -> Optional[str]:
        """Validate a protein sequence. Returns error message or None."""
        if len(sequence) < 10:
            return "Sequence too short (minimum 10 residues)"
        if len(sequence) > 1024:
            return "Sequence too long (maximum 1024 residues)"
        invalid = set(sequence.upper()) - self.valid_residues
        if invalid:
            return f"Invalid residues: {invalid}"
        return None

    @torch.no_grad()
    def predict(self, sequences: List[str],
                batch_size: int = 16) -> List[ThermostabilityResult]:
        """Predict melting temperatures for a list of proteins."""
        results = []
        self.model.eval()

        for i in range(0, len(sequences), batch_size):
            batch_seqs = sequences[i:i + batch_size]

            # Tokenize
            encoding = self.tokenizer(
                batch_seqs,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=1026
            ).to(self.device)

            # Predict
            predictions = self.model(
                encoding["input_ids"],
                encoding["attention_mask"]
            )

            for seq, pred in zip(batch_seqs, predictions):
                tm = pred.item()
                # Confidence heuristic: predictions near the training
                # distribution's center (40-65 C) are more reliable
                if 40 <= tm <= 65:
                    confidence = "high"
                elif 35 <= tm <= 75:
                    confidence = "medium"
                else:
                    confidence = "low"

                results.append(ThermostabilityResult(
                    sequence=seq,
                    predicted_tm=round(tm, 1),
                    confidence_category=confidence
                ))

        return results

# Usage example
service = ThermostabilityService("thermostability_predictor")

test_sequences = [
    # ubiquitin
    "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG",
    # GFP
    "MSKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTFSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK",
]

results = service.predict(test_sequences)
for r in results:
    print(f"Predicted T_m: {r.predicted_tm} C "
          f"({r.confidence_category} confidence) "
          f"[{len(r.sequence)} residues]")