g-DPO: Scalable Preference Optimization for Protein Language Models

TL;DR

This work tackles the quadratic scaling of Direct Preference Optimization (DPO) for protein language models by introducing g-DPO, which combines sequence-space clustering (union-mask clustering) to prune redundant preference pairs with grouped likelihood amortization to share forward passes across multiple sequences. The method begins with evo-tuning on evolutionarily related sequences, then clusters mutants by shared mutational positions, and finally updates the model using group-based DPO losses computed within each cluster. Across three protein mutational landscapes, g-DPO matches the in silico and in vitro performance of standard DPO while delivering 1.7× to 5.4× faster convergence, with larger gains expected on bigger datasets. This scalable approach enables more efficient protein design with PLMs and points to future directions such as rank-based objectives and extensions to multi-modal, sequence-structure models.

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.

Figures (8)

• Figure 1: Clustering of mutants with overlapping mutations.
• Figure 2: In silico evaluation of g-DPO. (a) Spearman correlation ($\rho$) of evo-tuned, DPO, and g-DPO models. (b) Convergence speedup of g-DPO vs. DPO (higher better), measured on a single NVIDIA A100 GPU; wall-clock times are in Table \ref{['tab:runtime_comparison']}. (c) Predicted property distributions of sequences generated from the evo-tuned, DPO, and g-DPO models using beam search. KS tests confirm that g-DPO and DPO yield nearly identical, and improved distributions over the reference.
• Figure 3: In vitro validation. Distribution of assay measurements for sequences designed with g-DPO and DPO on (a) thermostability of DhaA and (b) expression of Trastuzumab. Distributions show that both models yield comparable outcomes.
• Figure 4: (a) Effect of clustering threshold $\tau$ on performance and training pairs. (b) Effect of grouping size $g$ on performance and convergence speed without clustering.
• Figure 5: Likelihood computation comparison between g-DPO and DPO.