Steering Generative Models with Experimental Data for Protein Fitness Optimization

Yisong Yue (Caltech, Asari AI) · Wenda Chu (California Institute of Technology) · Jason Yang (Caltech) · Daniel Khalil (California Institute of Technology) · Raul Astudillo (MBZUAI) · Bruce Wittmann (Microsoft) · Frances Arnold (California Institute of Technology)

adaptive sequence selectionbayesian optimization classifier guidance diffusion models discrete diffusion models generative models labeled data language modelsplug-and-play guidanceposterior samplingprotein fitness optimizationreinforcement learningsequence-fitness pairssurrogate rewardsthompson sampling

Protein fitness optimization involves finding a protein sequence that maximizes desired quantitative properties in a combinatorially large design space of possible sequences. Recent advances in steering protein generative models (e.g., diffusion models and language models) with labeled data offer a promising approach. However, most previous studies have optimized surrogate rewards and/or utilized large amounts of labeled data for steering, making it unclear how well existing methods perform and compare to each other in real-world optimization campaigns where fitness is measured through low-throughput wet-lab assays. In this study, we explore fitness optimization using small amounts (hundreds) of labeled sequence-fitness pairs and comprehensively evaluate strategies such as classifier guidance and posterior sampling for guiding generation from different discrete diffusion models of protein sequences. We also demonstrate how guidance can be integrated into adaptive sequence selection akin to Thompson sampling in Bayesian optimization, showing that plug-and-play guidance strategies offer advantages over alternatives such as reinforcement learning with protein language models. Overall, we provide practical insights into how to effectively steer modern generative models for next-generation protein fitness optimization.