g-DPO: Scalable Preference Optimization for Protein Language Models

AI summary

Overview

Research area: Machine learning for protein engineering — specifically preference optimization (DPO) applied to protein language models, combined with data-efficient clustering and batched likelihood estimation.

Technical level: Intermediate. Readers need familiarity with language model fine-tuning and the DPO objective; the protein-engineering and clustering components are explained in accessible terms.

Scope: The paper introduces g-DPO, a method that prunes redundant preference pairs from experimentally labeled protein variant datasets via sequence-space clustering and amortizes likelihood computations across groups of sequences, and evaluates it in silico and in vitro on three protein engineering tasks.

What This Paper Is About

Protein language models are fine-tuned with Direct Preference Optimization to favor protein variants that perform better in experiments, but building preference pairs from scalar assay labels requires comparing variants with each other, and the number of possible pairs grows quadratically with the number of labeled sequences. This makes training slow even on modest datasets. The paper's goal is to prune the redundant comparisons and compute likelihoods more cheaply, so that DPO-style alignment on protein data converges faster without losing model quality.

Key Contributions

1. Scalable preference sampling. A sequence-space clustering procedure (union mask clustering) that greedily groups variants with overlapping mutational positions, discarding redundant pairs and keeping comparisons local in sequence space where mutational effects are more informative.

2. Efficient training via grouped likelihood amortization. Using the shared union mask of a group, all pairwise preferences among g sequences are evaluated from a single forward pass instead of one pass per pair; the authors note that with g=4, 6 pairs come from one forward pass instead of 6.

3. Empirical validation across three protein mutational landscapes, showing in silico metrics and in vitro assay results statistically indistinguishable from standard DPO while converging 1.7× to 5.4× faster, with the advantage expected to grow with dataset size.

4. Ablations isolating the effects of the clustering threshold τ and the group size g, and showing that grouping only works when clustering keeps the union mask small.

Main Findings

• Convergence speedup of 1.7× to 5.4×. Training speedups over standard DPO on a single NVIDIA A100 GPU: 1.68× training / 1.66× total on anti-SARS-CoV-2 VHH, 5.40× / 4.83× on anti-SARS-CoV-2 VHH L, 3.22× / 2.89× on Trastuzumab scFv, and 2.93× / 2.88× on haloalkane dehalogenase. Wall-clock training times ranged from 18m46s (g-DPO, Trastuzumab) to 7h22m09s (DPO, VHH L).

• Quality preserved in silico. Both DPO and g-DPO improve rank correlation relative to the evo-tuned reference model, and both shift the distribution of beam-search-generated sequences toward better predicted function. Two-sample Kolmogorov-Smirnov tests show statistically significant but very small effect sizes between DPO and g-DPO distributions (for example, D = 0.0290 on DhaA), which the authors interpret as nearly identical in practice. Trends hold for both Spearman ρ and Kendall's τ.

• Quality preserved in vitro. Assays on thermostability of DhaA and expression of Trastuzumab showed comparable outcomes between DPO and g-DPO, with no significant differences between their distributions. Candidates were down-selected by averaging performance across the top-k predicted sequences with k=3, following the Monte Carlo strategy of DiscoBAX.

• Moderate clustering is the sweet spot. With g=2, sweeping τ from 0.5 down to 0.1 left performance unchanged until roughly τ ≈ 0.3; below that, performance declined as clusters became too tight and useful signal was lost.

• Grouping alone fails at high mutation span. Without clustering, larger groups create union masks spanning too much of the sequence, likelihood approximations degrade, and model performance drops.

• Clustering plus grouping gives the best trade-off. With g=4, performance matched g=2 at τ ≈ 0.3 while converging significantly faster; the performance drop beyond τ ≈ 0.3 was less severe for g=4, which the authors attribute to within-group preferences being processed in one batch with coupled likelihoods and updates.

• Dataset sizes were modest. Experiments covered 76–474 variants (VHH 462, VHH L 1833, Trastuzumab scFv 76, DhaA 474) with mutation breadth of 13.1% to 92.4% of positions.

Methodology in Plain English

Training proceeds in three stages. First, the base model — ESM-2-650M, a masked protein language model — is fine-tuned without labels on evolutionarily related sequences of the wild-type protein retrieved with MMseqs2 searches against ColabFold databases, a step the authors call evo-tuning that supplies general evolutionary context rather than assay supervision.

Second, the labeled mutant dataset is clustered. Each cluster tracks a "union mask": the set of sequence positions where members of the cluster differ from one another. Starting with each sequence in its own cluster, the algorithm repeatedly merges the pair of clusters whose union mask would grow the least, stopping once the smallest available merge would push the mask beyond a fraction τ of the sequence length. Large datasets are first coarse-clustered with MMseqs2, then union-mask clustered within those buckets. The reported complexity is O(n² log n + n²L).

Third, groups of g sequences are sampled uniformly without replacement from each cluster. Because all sequences in a group are scored under the same masked context, one forward pass on a jointly masked input yields logits for every differing position, and the log-likelihood of each sequence is approximated by summing log-probabilities over the union mask. This is a mean-field approximation that assumes masked positions are conditionally independent given the observed tokens; the error scales with the size of the difference mask, so keeping masks small via clustering is what makes it reliable. The DPO loss is then applied over all pairwise comparisons inside each group, and sampling continues until every sequence has appeared in at least one group, defining one epoch.

Training used SGD with a learning rate of 7×10⁻⁴, no weight decay, a 300-step linear warmup, β = 0.04, batch size 64, a maximum union mask size of 0.3L, and groups of g=4. Validation loss was monitored every 250 steps with checkpointing and relative early stopping requiring at least 1% improvement with a patience of 3 validation intervals.

Why This Matters

The work targets a practical bottleneck in applying preference optimization to experimental protein data: pairwise preference construction is quadratic in dataset size, so alignment becomes computationally expensive exactly when datasets get large enough to be scientifically useful. g-DPO shows that much of that cost is redundant, and that cheaper training need not cost model quality.

Real-world applications:

• Antibody thermostability engineering — the anti-SARS-CoV-2 VHH datasets optimize thermostability, relevant to making biologic drugs more stable and easier to manufacture and store.
• Antibody expression optimization — the Trastuzumab scFv dataset targets expression, a common failure mode in developing antibody therapeutics.
• Biocatalysis and bioremediation — haloalkane dehalogenase catalyzes hydrolysis of halogenated compounds by cleavage of the carbon-halogen bond, an enzyme class relevant to industrial chemistry and environmental cleanup.
• General assay-driven protein design campaigns — any wet-lab loop where variants are measured and the results must be fed back into model fine-tuning under time and compute constraints.

Industry relevance: The authors are affiliated with Cradle (Zürich, Switzerland), indicating direct commercial interest in making protein design campaigns faster and cheaper. Reducing convergence time by 1.7× to 5.4× on a single GPU translates into more design-build-test cycles per unit of compute, which matters for any organization iterating on protein variants at scale.

Future Directions

• Scaling evaluation to larger, more diverse benchmarks. The authors note their experiments cover three mutational landscapes with 76–474 variants and call for testing on multiple wild types and broader mutation spans; they expect g-DPO's advantage to grow in such settings.
• Rank-based objectives. Moving beyond the pairwise DPO loss to objectives that use the full ranking of variants within a cluster, leveraging more of the available ordering information.
• Multi-modal foundation models. The clustering and grouping strategies are described as modality-agnostic, making sequence-structure or sequence-function models a natural extension.
• Hyperparameter generalization. The authors state that g=4 with a 0.3L union mask cap was the final configuration after sweeps, but expect the optimal group size to depend on dataset scale and clustering level — leaving open how to choose these automatically.

Target Audience

Machine learning researchers working on preference optimization and alignment beyond NLP, protein engineers and computational biologists building assay-driven design pipelines, and practitioners at biotech and pharmaceutical companies who need to fine-tune protein language models within realistic compute budgets. Readers with a background in language model fine-tuning will get the most from the methodological sections, while the ablation results are useful to anyone deciding how to configure clustering and grouping in a production setting.

Authors’ abstract

Direct Preference Optimization (DPO) is an effective approach for aligning protein language models with experimental design goals. However, DPO faces a scalability bottleneck: the number of possible training pairs grows quadratically with the number of labeled sequences, leading to prohibitive training times even for modestly sized datasets. We introduce g-DPO, a framework that (i) uses sequence space clustering to prune redundant pairs while preserving training signal, and (ii) amortizes likelihood computations with group-based approximations. Across three protein engineering tasks, g-DPO maintains in silico and in vitro performance that is statistically indistinguishable from standard DPO, while converging 1.7x to 5.4x times faster, with speedups that scale with dataset size and the structure of the underlying mutational landscape.