Steering Generative Models with Experimental Data for Protein Fitness Optimization

TL;DR

This paper evaluates methods for guiding generative models to optimize protein fitness using limited labeled data, providing practical insights into effective strategies for real-world applications.

Contribution

It systematically compares guidance strategies like classifier guidance and posterior sampling for protein sequence generation with small datasets.

Findings

Guidance methods outperform reinforcement learning in protein optimization.

Plug-and-play guidance strategies are more effective and practical.

Adaptive sequence selection improves optimization efficiency.

Abstract

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.