ProteinOPD: Towards Effective and Efficient Preference Alignment for Protein Design

Abstract: Designing proteins with desired functions or properties represents a core goal in synthetic biology and drug discovery. Recent advances in protein LLMs (PLMs) have enabled the generation of highly designable protein sequences, while preference alignment provides a promising way to steer designs toward desired functions and properties. Nevertheless, they often trigger catastrophic forgetting of pretrained knowledge, degrading basic designability and failing to balance multiple competing objectives. To address these issues, we draw inspiration from On-Policy Distillation (OPD), an advanced post-training method renowned for mitigating catastrophic forgetting through its mode-seeking nature. In this work, we propose ProteinOPD, a multi-objective preference alignment framework that can effectively balance multiple preference objectives while maintaining the inherent designability of PLMs. ProteinOPD adapts a pretrained PLM into preference-specific teachers and distills their knowledge into a shared student via token-level OPD on the student's own trajectories. During this process, the student is aligned to a unique normalized geometric consensus of weighted teachers while ensuring bounded optimization under conflicts. This bridges the gap for OPD in multi-objective/teacher alignment. Extensive experiments show that ProteinOPD achieves substantial gains on target preference objectives without compromising the designability, with an 8x training speedup over RL-based alignment competitors.

• The paper introduces an innovative on-policy distillation method to achieve multi-objective preference alignment in protein design, effectively mitigating catastrophic forgetting.
• It demonstrates remarkable trade-off improvements with a 34.8%-67.6% surge in Pareto efficiency metrics, alongside significant reductions in perplexity and boosts in thermostability and solubility.
• The framework offers an 8x training speedup with minimal data requirements, providing a scalable and efficient solution for property-driven protein engineering.

ProteinOPD: A Multi-Objective Preference Alignment Framework for Protein Design

Overview and Motivation

ProteinOPD addresses a central challenge in protein engineering: generating protein sequences exhibiting specific, desired properties while retaining structural plausibility and diversity. While pretrained protein LLMs (PLMs) have demonstrated strong capacities for generating designable proteins, aligning these generative models with explicit multi-property objectives (e.g., foldability, solubility, thermostability) remains unresolved. Existing preference alignment paradigms—such as supervised fine-tuning (SFT), reinforcement learning (RL), and test-time steering—suffer from specific weaknesses: SFT is mode-covering and can drive catastrophic forgetting of the underlying PLM, RL is data/compute intensive and prone to policy drift, and test-time steering provides weak preference optimization and fails in multi-objective scenarios.

ProteinOPD introduces a principled post-training framework extending On-Policy Distillation (OPD) to the multi-objective regime, enabling efficient and effective preference alignment in protein sequence design. Through structured teacher-student distillation with a product-of-experts consensus, ProteinOPD preserves the foundational designability of the PLM while maximizing multiple property-driven performance metrics.

Methods and Framework

Single-Objective OPD

OPD leverages a pretrained PLM adapted into a preference-specific teacher via lightweight SFT using a highly filtered subset of high-quality sequences based on oracle scoring. The student model then aligns with the teacher through token-level KL/JSD minimization over rollouts generated by the student itself, addressing exposure bias and reducing catastrophic forgetting. The dense feedback per token accelerates convergence compared to sparse reward RL. Empirically, OPD achieves property alignment comparable to RL while maintaining sequence naturalness and diversity.

Generalized OPD for Multi-Objective Alignment

A significant methodological innovation of ProteinOPD is the extension of OPD to multiple, potentially conflicting objectives. For $M$ teacher policies $\{p_i\}$ with property-specific weights $\{w_i\}$, the student’s consensus distribution is a normalized product-of-experts (PoE) over the teacher distributions:

$$p_{PoE}(\cdot | x, y_{<n}) = \text{softmax}\left(\sum_{i=1}^M w_i \log p_i(\cdot | x, y_{<n})\right)$$

This geometric consensus robustly captures jointly supported sequence patterns, avoids over-generalization typical of arithmetic mixtures, and provides a bounded optimization signal even in the presence of teacher disagreement. Notably, the normalization constant quantifies conflict among teachers, offering a low-cost introspective metric for training stability. Student training uses OPD with this composite consensus, balancing property gains and preserving model priors.

Efficiency and Practicalities

ProteinOPD’s protocol is highly efficient:

• Teacher adaptation requires only SFT on ≈100–200 protein sequences per preference.
• Student OPD distillation dramatically reduces compute time relative to RL, reaching comparable property alignment within one-eighth of the wall-clock time.
• Offline data cost is minimized due to the small number of high-scoring sequences needed to construct competent teachers.

Experimental Results

Multi-Objective Performance and Trade-off

Extensive benchmarking on both unconditional (ProtGPT2) and conditional (ProLLaMA) generation tasks demonstrates that ProteinOPD substantially advances the achievable Pareto front between designability and property alignment. Notable quantitative results (unconditional setting):

• pLDDT (foldability): improved by 14.8% over base ProtGPT2 and by 6% over the best strong baseline.
• Solubility: improved by 16.9% over base, matching or exceeding all baselines.
• Thermostability: increased by 54.2%, outperforming the closest baseline by 8.8%.
• Sequence plausibility (PPL): reduced by 83.7%, with increased novelty, indicating the retention of "natural" protein features and diversity.
• Hypervolume (HV): improved by 34.8% over strong multi-objective baselines, reflecting a better global trade-off.

In the conditional setting (e.g., lysozyme-like family), ProteinOPD improved ProTrek score, pLDDT, pAE, and thermostability over both PLM and preference-alignment competitors.

Ablation and Analysis

1. Pareto Expansion: ProteinOPD consistently expands the trade-off surface, especially for high-novelty proteins, which is crucial for de novo design in unexplored sequence space.
2. Catastrophic Forgetting: OPD mitigates the loss in designability observed with SFT-derived teachers, confirmed by minimal reduction in sequence novelty and improved condition faithfulness in conditional tasks.
3. Acceleration and Sample Efficiency: Eightfold reduction in alignment time vs. RL, and significant gains with limited teacher data support broad practical utility for expensive protein objectives.

Implications and Future Prospects

The introduction of multi-teacher OPD in ProteinOPD marks a notable step toward scalable, controllable protein design systems. By enabling the simultaneous and efficient optimization of multiple properties without sacrificing the generative fluency and diversity of PLMs, ProteinOPD sets a new paradigm for practical protein engineering with artificial intelligence.

Practical implications include:

• Synthetic Biology and Drug Discovery: Rapid generation of highly functional proteins meeting complex multi-property constraints (e.g., stability, manufacturability, function).
• Generalized Model Alignment: The geometric-consensus OPD framework is applicable to other domains requiring multi-property or multi-preference controllable sequence generation (e.g., chemical design, regulatory genomics).
• Scalability/Automation: The demonstrated training speed and data-parsimonious approach are critical for real-world high-throughput applications, including iterative wet-lab validation loops.

Key avenues for future research:

• Integration with explicit structure prediction or generative modeling to close the loop with 3D design.
• Incorporation of uncertainty quantification in preference oracles and weights.
• Experimental validation of generated proteins in wet-lab settings to confirm computational gains translate to empirical functionality.

Conclusion

ProteinOPD demonstrates that multi-objective on-policy distillation, with a product-of-experts geometric consensus, enables protein LLMs to achieve strong, balanced property alignment with minimal compromise to sequence plausibility and novelty and at greatly reduced computational cost. This work substantiates OPD as an advanced paradigm for alignment tasks in molecular generative modeling, providing both methodological and empirical advances in AI-driven protein design (2605.10189).