Be Your Own Teacher: Steering Protein Language Models via Unsupervised Reward Optimization

Abstract: Protein LLMs (PLMs) have emerged as powerful tools for controllable biomolecular design, yet their post-training adaptation typically relies on costly wet-lab validation or curated preference datasets. To overcome this supervision bottleneck, we introduce unsupervised reward optimization of PLMs, a comprehensive framework for steerable protein generation without ground-truth labels. Our key insight is that task-agnostic rewards, which combine intrinsic model uncertainty with extrinsic semantic consistency informed by protein representation models, exhibit strong correlation with controllability measures across base models and temperature regimes. Building upon this discovery, we propose two offline algorithms: Soft Reward Optimization (SRO) and Binarized Reward Optimization (BRO), which effectively maximize the classical RLHF objective induced by these proxy rewards. Extensive experiments on compositional out-of-distribution prompts demonstrate that both methods significantly outperform competitive baselines (DPO, KTO), while approaching oracle performance across multiple sampling temperatures, model scales and protein families. Moreover, PLMs fine-tuned with unsupervised rewards can achieve consistently higher coverage compared to their base model in pass@k evaluations. By enabling self-improvement of PLMs through their own generated experience, our framework provides a scalable pathway toward controllable biomolecular design in settings where labeled preferences or experimental feedback are scarce or unavailable.

• The paper introduces unsupervised reward optimization to steer PLMs, demonstrating significant improvements in keyword and structure recovery compared to baseline methods.
• It develops a hybrid approach using intrinsic (entropy-based) and extrinsic (semantic distance) rewards, optimized via soft and binarized strategies under offline settings.
• Empirical evaluations on Pfam and AlphaFoldDB datasets show robust performance, with methods reaching up to 61.3% pass@1 recovery on challenging out-of-distribution prompts.

Steering Protein LLMs via Unsupervised Reward Optimization

Motivation and Context

The adaptation of Protein LLMs (PLMs) for controlled biomolecular design is currently bottlenecked by the cost of supervision, relying extensively on either expensive wet-lab validation or large hand-curated preference data. This challenge is particularly acute as model capacity scales and as application domains exceed the bounds of current experimental throughput or readily available expertise. The study "Be Your Own Teacher: Steering Protein LLMs via Unsupervised Reward Optimization" (2606.18961) proposes a methodology for unsupervised post-training of PLMs, leveraging proxy rewards defined without ground-truth labels, and develops algorithmic strategies for robust reward maximization under noisy, task-agnostic supervision. The principal goal is to endow PLMs with steerability and controllable generalization in regimes where annotation scarcity and offline constraints preclude classical supervised or RLHF approaches.

Methodological Innovations

The framework introduced centers around two core contributions: (1) the design of hybrid unsupervised reward signals correlated with controllability, and (2) reward optimization algorithms compatible with offline, proxy-labeled datasets.

Task-Agnostic Proxy Rewards

The authors systematically construct intrinsic and extrinsic reward functions:

• Intrinsic Uncertainty-Based Rewards: Predictive negative log-likelihood (entropy) and its normalized variant are used as proxies for model confidence, with high confidence (lower entropy) hypothesized to correspond to higher controllability.
• Extrinsic Consistency Rewards: Semantic distances between generated proteins (via embeddings from pretrained protein representation models, notably ESMC) are employed as extrinsic rewards. The negative of the distance is used to encourage generations that are consistent under the representation model, penalizing statistical outliers and hallucinations.

A crucial empirical finding is the correlation between these proxy rewards and ground-truth controllability metrics, such as keyword/structure recovery, across PLMs, temperatures, and compositional prompts.

Multi-Temperature Sampling

Diversity in both reward signals and generative outputs is ensured via multi-temperature sampling, counteracting degeneracy at low and incoherency at high generation temperatures. The reward function is adaptively chosen per temperature to maximize correlation with controllability, thus constructing an offline training set conducive to robust post-training.

Offline Reward Optimization Algorithms

Two offline algorithms are developed:

• Soft Reward Optimization (SRO): Directly maximizes the expected proxy reward with KL-regularization against the base policy, employing continuous supervision signals, with efficient closed-form solutions derived and implemented via importance-weighted Monte Carlo estimation.
• Binarized Reward Optimization (BRO): Adapts to noisy or ambiguous reward signals by thresholding the continuous rewards into positive and negative domains, recasting the objective as a binary classification loss in the policy space, framed game-theoretically as maximizing the empirical win-rate of the improved policy versus the reference.

This approach generalizes and subsumes methods such as DPO and KTO, which are limited to pairwise or strictly ordered preferences, making SRO/BRO broadly applicable.

Empirical Evaluation

Experiments are performed on compositional out-of-distribution (OOD) prompt sets derived from Pfam families and AlphaFoldDB proteins, spanning both sequence-to-function (Func2Seq) and structure-to-sequence (Struct2Seq) tasks.

Notable empirical results include:

• Significant Outperformance Over Baselines: SRO and BRO consistently outperform competitive post-training methods such as DPO and KTO across model sizes, sampling temperatures, and protein families, with average pass@1 recovery rates notably higher (see Table 1 and Table 10 in the paper for details). For the 764M parameter ProGen2, the BRO method achieves a 61.3% average recovery rate on challenging OOD prompts, compared with 60.3% for an oracle model fine-tuned with ground-truth rewards.
• Robustness to Noisy Supervision: Across ablation studies on reward types and sampling strategies, both SRO and BRO show stable improvement, underscoring the effectiveness of unsupervised reward and data selection.
• Steerability and Coverage: Post-trained models using unsupervised rewards exhibit higher pass@k coverage and sustained ability to prioritize controllable generations for small k, indicating not only improved raw accuracy on prompt recovery but also enhanced reliability in practical top-k selection protocols.
• Generalization: Methods are robust across Pfam families and OOD compositions, supporting claims of increased compositional generalization and operational creativity in sequence generation.

The study also finds a "critical temperature" in sampling where proxy rewards are maximally predictive of controllability, a phenomenon stable across model architectures and tasks.

Theoretical Implications

This paper advances the theory and practice of offline, unsupervised RL for generative sequence models by formalizing the equivalence between reward maximization and divergence minimization to an optimal policy proved in Theorems 3.1 and 3.2. The introduction of a game-theoretic interpretation of BRO further elucidates the dynamics of unsupervised policy improvement, highlighting the limits and guarantees of improvement given noisy, unlabeled rewards.

The hybridization of intrinsic and extrinsic signals in reward design echoes themes in recent work on LLM alignment, yet the empirical study demonstrates effective transfer and adaptation to the protein generation domain, where biological meaning is deeply decoupled from surface sequence probability.

Practical Implications and Future Directions

From a practical perspective, this framework enables scalable self-improvement and steerability of PLMs in biologically relevant domains and compositional tasks where labeled training signals are sparse, delayed, or expensive. Applications span enzyme engineering, therapeutic design, and synthetic biology, where controllable generation is paramount but wet-lab feedback is a bottleneck.

Several limitations are acknowledged:

• The framework currently underperforms structure-based metrics for inverse folding/structure prediction tasks.
• All results are in silico; translation to wet-lab or clinical validations is needed for practical impact.
• Reliance on proxy rewards may induce an upper-bound on performance, especially where weak alignment with biological utility is present.
• The approach is dual-use: improved PLM controllability may have both beneficial and adverse implications, necessitating biosecurity consideration.

Future work directions include the development of more sophisticated, uncertainty-aware or adaptive rewards, the integration of limited supervision or weak oracle signals to reduce proxy reward approximation error, and extending the approach to multimodal biomolecular design where sequence, structure, and function must be jointly optimized.

Conclusion

This work establishes unsupervised reward optimization as a viable paradigm for post-training and steering of protein LLMs, with strong empirical evidence for enhanced controllability and generalization in the absence of ground-truth labels. The introduced proxy rewards, multi-temperature sampling, and SRO/BRO algorithms are shown to provide robust gains over contemporary frameworks. The study positions unsupervised, task-agnostic reward optimization as a cornerstone for future advances in BioAI, particularly in settings of data and supervision scarcity. It also highlights several open challenges at the intersection of self-supervised RL, biological sequence modeling, and AI safety for biological systems.