Plug & Play Directed Evolution of Proteins with Gradient-based Discrete MCMC

Abstract

A long-standing goal of machine-learning-based protein engineering is to accelerate the discovery of novel mutations that improve the function of a known protein. We introduce a sampling framework for evolving proteins in silico that supports mixing and matching a variety of unsupervised models, such as protein language models, and supervised models that predict protein function from sequence. By composing these models, we aim to improve our ability to evaluate unseen mutations and constrain search to regions of sequence space likely to contain functional proteins. Our framework achieves this without any model fine-tuning or re-training by constructing a product of experts distribution directly in discrete protein space. Instead of resorting to brute force search or random sampling, which is typical of classic directed evolution, we introduce a fast MCMC sampler that uses gradients to propose promising mutations. We conduct in silico directed evolution experiments on wide fitness landscapes and across a range of different pre-trained unsupervised models, including a 650M parameter protein language model. Our results demonstrate an ability to efficiently discover variants with high evolutionary likelihood as well as estimated activity multiple mutations away from a wild type protein, suggesting our sampler provides a practical and effective new paradigm for machine-learning-based protein engineering.

31 pages, 8 figures. To appear in the Machine Learning: Science & Technology (ML:S&T) journal. Code is available at https://github.com/pemami4911/ppde. A short version of this work appeared at the NeurIPS 2022 Machine Learning in Structural Biology Workshop

References in corpus (10)

• The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables
• Tranception: protein fitness prediction with autoregressive transformers and inference-time retrieval
• Generating and designing DNA with deep generative models
• RITA: a Study on Scaling Up Generative Protein Sequence Models
• AdaLead: A simple and robust adaptive greedy search algorithm for sequence design
• Model Inversion Networks for Model-Based Optimization
• Deep Extrapolation for Attribute-Enhanced Generation
• ReLSO: A Transformer-based Model for Latent Space Optimization and Generation of Proteins
• A Langevin-like Sampler for Discrete Distributions
• Discrete Langevin Sampler via Wasserstein Gradient Flow