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

2026-06-17

AI Summary

Protein language models face a supervision bottleneck in controllable design due to reliance on costly wet-lab experiments or human preference labels. This work proposes an unsupervised reward optimization framework that constructs task-agnostic proxy rewards by integrating the model’s intrinsic uncertainty with semantic consistency from protein representation models. It introduces two offline optimization algorithms—SRO and BRO—enabling self-guided optimization without any ground-truth labels. Leveraging only data generated by the model itself, the approach achieves controllable generation and significantly outperforms baselines such as DPO and KTO on out-of-distribution compositional prompts, approaching oracle-level performance. Moreover, it consistently improves pass@k coverage across varying temperatures, model scales, and protein families.

Abstract

Protein language models (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.

Problem

• protein language models
• controllable biomolecular design
• supervision bottleneck
• unsupervised reward optimization
• preference datasets

Innovation

• unsupervised reward optimization
• protein language models
• controllable generation
• self-improvement
• proxy rewards