Plug & play directed evolution of proteins with gradient-based discrete MCMC

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 Markov chain Monte Carlo 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 650 M 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.


Introduction
Engineering proteins to improve their productivity or catalyze new reactions requires scientists to navigate the complex landscape mapping a protein's amino acid sequence to its structure and function (Li et al., 2020).Directed evolution is a classic approach inspired by natural evolution where random mutations to a protein's sequence are screened in a wet lab until higher-performing variants are found, at which point the process repeats starting from these variants (Kuchner and Arnold, 1997).However, this becomes impractical when proteins cannot be assessed in a high-throughput fashion.The simplest approach, brute force search, is limited in practice to variants with one or two mutations.For a protein (b) Figure 1: a) Illustration of the machine-learning-based directed evolution pipeline.This paper proposes a plug and play sampler for proposing variants of improved fitness near a wild-type (WT) protein for each design round.b) We flexibly compose unsupervised and supervised protein sequence models as a product of experts distribution and derive a fast gradient-based discrete MCMC sampler to efficiently sample from this distribution.A product of experts places its probability mass at the intersection of the expert distributions.
The key idea of gradient-based discrete MCMC is to bias the mutation proposal distribution towards promising mutations within a local neighborhood around the current variant.Then, to avoid enumerating all mutations in this neighborhood, the proposal is approximated with a first-order Taylor series.An example MCMC trajectory of single point mutations starting from the WT is shown.
with 400 amino acids, there are ∼ 10 19 ways to make five single substitutions assuming the standard vocabulary of 20 amino acids.It is also remarkably difficult to find proteins with improved function.Most of protein space is non-functional and beneficial mutations are rare (Arnold, 1998).
As supervised machine learning methods for protein function prediction from primary sequence improve (Dallago et al., 2021;Hsu et al., 2022), machine-learning-based directed evolution has emerged, which offers a way to judiciously propose candidates for screening that ultimately reduces time spent in the wet lab (Yang et al., 2019;Wu et al., 2021;Biswas et al., 2021).Our work is concerned with improving the variant proposal step by increasing the chance of discovering a mutation that improves a target property of a known "wild type" (WT) protein (Figure 1a).We argue that the promising performance of recent unsupervised sequence models at mutation effect prediction (Meier et al., 2021;Hsu et al., 2022;Weinstein et al., 2022) has motivated a reconsideration of simple black-box optimization algorithms for searching sequence space.Black-box algorithms are appealing due to their simplicity, flexibility, compatibility with discrete spaces, and use of interpretable mutation operators.They offer a way to mix and match unsupervised and supervised models for proposing variants without requiring any model fine-tuning or re-training, i.e., in a "plug and play" manner.Combining both types of models is generally advantageous.While unsupervised models learn information that can steer search away from adversarial inputs which fool supervised models due to overestimation errors (Szegedy et al., 2014), supervised models learn specific information about beneficial mutations gleaned from assaylabeled data.However, black-box algorithms tend to be extremely inefficient at searching for variants with improved fitness.A practical plug and play framework for searching protein sequence space has remained elusive due to the difficulty of fast search in discrete, high-dimensional spaces.
This paper fills that gap by introducing Plug and Play Directed Evolution (PPDE), a practical plug and play sampling framework for the efficient discovery of functional variants directly in discrete protein space.PPDE flexibly combines unsupervised and supervised models with a product of experts distribution (Hinton, 2002) p(x) = i p i (x).Combining both types of protein models in this manner encourages variants to have both high sequence likelihood and high predicted function (e.g., activity).To sample efficiently from the highdimensional, discrete, and unnormalized distribution p(x), we derive a fast Markov chain Monte Carlo (MCMC) sampler that uses gradients of p(x) to propose mutations.See Figure 1b for a visual depiction the sampler.PPDE nearly maintains all of the flexibility of black-box algorithms-it only has to additionally assume that p(x) is differentiable at each discrete point of protein space.We theoretically characterize the efficiency of this sampler and empirically show that our framework works with a variety of pre-trained models without using continuous relaxations and without model retraining or fine-tuning.Although our formulation can be applied to a broad class of biological sequence design tasks, we focus on proteins due to the current widespread availability of pre-trained models.
A summary of our contributions: • We introduce a plug and play sampling framework that enables mixing and matching various unsupervised and supervised protein models without model re-training or finetuning.
• We demonstrate a novel application of gradient-based discrete MCMC for conducting search with a product of experts distribution.
• We conduct rigorous in silico directed evolution experiments on three target proteins with partially characterized wide (≥ 2 mutation) fitness landscapes.Our results provide strong evidence that our sampler offers a practical and effective approach for ML-based directed evolution beyond brute force and random search.

Related Work
Black-box approaches to directed evolution: A popular paradigm for ML-based directed evolution has been to learn a surrogate for the sequence-to-function landscape from a small labeled dataset and then to use a black-box optimizer or sampler such as an evolutionary algorithm or simulated annealing to search for mutants (Hansen and Ostermeier, 1996;Angermüller et al., 2020;Sinai et al., 2020;Biswas et al., 2021).Black-box approaches are plug and play in that they can easily combine unsupervised and supervised protein models without any model fine-tuning or re-training.Navigating the high-dimensional discrete space is approached by either using a random walk, which is highly inefficient, or by using gradient ascent via a continuous relaxation of the space, which can bias the search process towards poor local optima.These are the main baselines for our method and will be described in more detail in Section 5. Reinforcement learning (RL) has also been used as a black-box optimization framework for training autoregressive models to generate highfitness variants (Angermueller et al., 2019).However, this approach has less flexibility than plug and play methods, in that it does not support easy mixing of unsupervised pre-trained models with the learned RL policy.
Protein design by generative modeling: An alternative to directed evolution is to cast protein design as inference, i.e., fitting a conditional generative model to assay-labeled proteins showing high activity (Brookes and Listgarten, 2018;Brookes et al., 2019;Fannjiang and Listgarten, 2020;Gupta and Zou, 2018;Zhang et al., 2022a;Jain et al., 2022).Since high performing variants are rare, this task amounts to density estimation of a rare event which is a difficult statistical inference problem in its own right.Another popular alternative is latent space optimization, in which gradient ascent is performed on a surrogate function trained directly in the latent space of a generative model.These generative models use regularization to shape their latent space favorably so that optimization remains constrained to the manifold of viable proteins (Killoran et al., 2017;Gómez-Bombarelli et al., 2018;Kumar and Levine, 2020;Linder et al., 2020;Chan et al., 2021;Trabucco et al., 2021;Castro et al., 2022).The Deep Manifold Sampler (DMS) (Gligorijević et al., 2021) more explicitly constrains optimization by alternating between steps of gradient ascent and reprojection of a noised version of the sequence onto the latent space of a custom denoising autoencoder (DAE).Our plug and play approach works with a variety of unsupervised models and can combine multiple such models if desired.
At the expense of trust in candidate proteins, unconditional generative models can also be used to hallucinate proteins from distant, unexplored regions of protein space.For example, training a generative model on sequences from a target protein family has been used to generate functional variants (Costello and Martin, 2019;Hawkins-Hooker et al., 2021;Shin et al., 2021).Unconditional sampling from protein language models trained on unaligned (Madani et al., 2020;Rives et al., 2021;Hesslow et al., 2022;Lin et al., 2022;Ferruz et al., 2022;Nijkamp et al., 2022;Yang et al., 2022) andaligned (Rao et al., 2021;Notin et al., 2022) sequences has recently been explored.While this is a promising new direction for protein engineering, our focus is on directed evolution from known proteins.
Plug and play text generation: Controlled generation of text has similarly been approached by combining pre-trained unsupervised models (e.g., large language models) and supervised models (e.g., sentiment classifiers) as a discrete product of experts (Holtzman et al., 2018;Dathathri et al., 2020;Qin et al., 2022).Unlike controlled text generation, where the goal is to guide a sampling process to transform a sentence by altering entire words and phrases, directed evolution seeks to accumulate a few precise changes to an initial sequence by modifying a small number of amino acids (i.e., individual characters).Moreover, plug and play approaches for controlled text generation intentionally avoid solving a combinatorial search problem by instead restricting the class of suitable unsupervised experts to left-toright autoregressive language models.This enables the use of sequential sampling algorithms such as left-to-right beam search decoding.Our use of gradient-based discrete MCMC for plug and play directed evolution allows us to efficiently search the combinatorial space of mutations near the wild type protein and to use a more general class of unsupervised experts that includes orderless protein sequence models which capture epistatic relationships between any pair of amino acids.The Metropolis-adjusted Langevin Algorithm (MALA) used for plug and play image generation (Nguyen et al., 2017) is a simple gradient-based MCMC sampler for continuous spaces, and is closely related to locally-balanced discrete MCMC (Zhang et al., 2022b;Sun et al., 2022a).Therefore, we use MALA with a handcrafted continuous relaxation as a baseline for plug and play sampling from discrete product of experts in our experiments.

Background
Before presenting our sampler, we define our search problem, briefly review Metropolis-Hastings (MH) MCMC, and introduce gradient-based discrete MCMC, a class of MH samplers able to efficiently explore high-dimensional discrete distributions.
Problem definition: The variant proposal step of directed evolution involves searching for mutations that improve one or more target properties of a given wild type (WT) protein.This search problem is formally defined over discrete protein sequences x := {x 0 , . . ., x L−1 }, x ∈ X, of length L with each x i taking on a value in a vocab of size V (typically V = 20 for the 20 standard amino acids).We assume that each x i is one-hot encoded.The search is initialized at the WT protein x W T and terminates when a predefined condition is met (e.g., a maximum allowable number of search steps is reached).
Metropolis Hastings: The MH algorithm (Metropolis et al., 1953;Hastings, 1970) defines the following algorithm for drawing samples from a distribution over protein variants p(x).Let f (x) be the unnormalized log probability of x such that log p(x) = f (x) − log Z where Z = x∈X exp(f (x)) is a normalizing constant.Given the current state x, draw a candidate next state x from proposal distribution q(x | x).Accept the proposed state with probability otherwise reject the transition and stay at x.This probabilistic acceptance/rejection criterion is desirable since it provides MH samplers with theoretical convergence guarantees and does not contain any hyperparameters, simplifying the implementation of MH in practice.
Gradient-based discrete MCMC: Uninformed MH proposals such as the uniform distribution are often inefficient for sampling from high-dimensional discrete distributions since candidate states x are proposed "blindly".In general, the efficiency of MH is highly dependent on the choice of proposal distribution.MH with locally-balanced informed pro-posals (Zanella, 2020) use distributions of the form where 1(x ∈ N (x)) is an indicator function.The key idea is to bias the proposal distribution towards local state transitions within a neighborhood N (x) that incur an increase in likelihood.We take the square root of the exponential in this proposal as the choice of local balancing function w(t) = tw(1/t), ∀t > 0. This function "balances" the acceptance and rejection probabilities in the local neighborhood N (x) to achieve a high acceptance rate.
The square root w(t) = √ t was empirically validated as a good default option in Zanella (2020).Since enumerating all local moves in N (x) in a discrete high-dimensional spaces is infeasible, an efficient alternative is available for functions f (x) whose gradient can be evaluated at the discrete state x (Grathwohl et al., 2021).A gradient-based locally-balanced informed proposal is the first-order Taylor-series approximation of Equation 1 around x: (2) When N (x) is the 1-Hamming ball, this proposal amounts to a tempered softmax over single changes to one dimension of x.One forward pass and one backwards pass is required to compute the forward q(x |x) and reverse q(x|x ) approximate proposals.
A concern with this sampler is that it is susceptible to getting trapped in poor local optima since it cannot perform large jumps across the search space at each step of MCMC (Sun et al., 2022b).One way to address this is to increase the Hamming window size, U , to U > 1.While this enables the sampler to propose changes to multiple dimensions of x simultaneously, it also incurs a large increase in computation since N (x) now contains O((LV ) U ) states to evaluate.Alternatively, a path proposal sequentially samples R ∼ Unif(1, U ) single changes to apply to the current state x (Sun et al., 2022b).For example, if x is the sequence "AAA", then a sampled path of length 3 could look like "AAA" → "AAB" → "AAC" → "ABC".
The approximate path proposal distribution starting at The r th path state x r is sampled from q x r |x r−1 .The sampler decides whether to accept or reject the terminal path state x := x R (e.g., "ABC"), i.e., the accumulation of all R changes applied to x.The reverse proposal qR (x|x ) is computed using the gradient at the terminal state ∇ x f (x ) and the reversed sequence of states (x R , x R−1 , . . ., x 0 ).The number of forward and backwards passes required to compute the forward and reverse path proposals is still two.In the next section, we will introduce and characterize a gradient-based path proposal for product of experts distributions.

Plug and Play Directed Evolution
Our approach for proposing mutations that improve a target property of the WT protein is to sample from a distribution constructed by combining unsupervised and supervised sequence models without using any continuous relaxations, model re-training, or fine-tuning.This distribution places its probability mass on proteins that have both high evolutionary density f (x) (i.e., high likelihood of being a naturally occurring protein) and high predicted function g(x) (e.g., activity or stability) by taking the product of multiple pre-trained "expert" distributions: Each f i (x) is an unsupervised model, each g j (x) is a supervised model, and is the unknown normalizing constant.Typically, f i (x) has been trained to do density estimation on unlabeled yet aligned sequences (e.g., a collection of evolutionarily related sequences provided by a multiple sequence alignment (MSA)) or unaligned sequences.While in this work we assume the g j (x) are an ensemble of nonlinear models trained to regress activity from a labeled dataset of mutants, this formulation can easily be extended to the multi-objective case where, for e.g., g 1 predicts activity and g 2 predicts stability.The unsupervised experts p f (x) ∝ i exp (f i (x)) act as a soft constraint that keeps the sampler near regions of high evolutionary density and away from, e.g., adversarial local optima of the supervised models.Examples of unsupervised models that provide an evolutionary density score include the EVmutation Potts model (Hopf et al., 2017), the ESM protein masked language models (Rives et al., 2021;Lin et al., 2022), denoising autoencoders (DAEs), energy based models (EBMs), autoregressive protein language models (e.g., ProGen2 (Nijkamp et al., 2022)), and normalizing flows (non-exhaustive list).The supervised expert p g (x) acts as a soft constraint that guides sampling towards proteins that have high activity, where p g (x) ∝ j exp (λ g j (x)) assigns high probability to sequences with high activity.The hyperparameter λ ≥ 0 allows us to balance the contribution of the unsupervised and supervised experts; for example, we can emphasize the "realism" of the protein (the evolutionary density) by setting λ to 0. However, recent evidence indicates that evolutionary density scores also positively correlate with protein fitness in a "zeroshot" manner (i.e., without using any labels) (Meier et al., 2021;Hsu et al., 2022;Weinstein et al., 2022).

Product of experts gradient-based discrete MCMC
Sampling from the product of experts (Equation 4) is difficult since the normalization constant Z is assumed unknown and is intractable to compute in practice.Although we have a good initialization for an MCMC sampler-x W T , the wild type protein of verified viability-traditional MCMC based on random walk exploration is too inefficient to discover variants with high predicted activity in reasonable time.
Our solution is to use fast gradient-based discrete MCMC.We need only additionally assume that each expert is a continuous function that is differentiable at each discrete x ∈ X (e.g., as is the case when the experts are neural networks).In detail, assume we have M continuously differentiable unsupervised experts and N continuously differentiable supervised experts.During each step of MCMC, we use the gradients of the M + N experts to approximate the change in product of experts likelihood due to making R point mutations to the current protein variant.This allows us to bias the proposal distribution towards the most promising R mutations.The gradient-based Taylor approximation of our MCMC proposal for log p(x) We use this path proposal to sample R single amino acid substitutions, which we apply to the current variant x 0 := x at each step of MCMC.The terminal state of the path x R := x is the variant that results from the accumulation of the R substitutions.To avoid computing extra forward and backwards passes through the product of experts for the intermediate path proposals q(x r |x r−1 ), following Sun et al. (2022b) we re-use the gradient taken with respect to the path origin x 0 instead of recomputing gradients at intermediate states.The same is done for the reverse path proposals q(x r−1 |x r ) with respect to the terminal state x R .In the next section, we characterize the efficiency of our sampler to understand how composing Taylor approximations for multiple experts affects the theoretical convergence rate.
Algorithm 1 shows pseudo-code for our fast MCMC sampler for plug and play directed evolution of proteins.The sampler follows the basic structure of MH MCMC.At each sampler step, we first compute the forward path proposal distribution qR (x |x) which we use to sample the proposed protein x .Then, we compute the reverse path proposal distribution qR (x|x ).We use these distributions to compute an acceptance criterion for determining whether to accept or reject x , after which the process repeats until termination (e.g., a predetermined number of MCMC steps is reached).

Sampler Analysis
Since we are using a gradient-based approximation of the product of experts proposal distribution, a natural question is whether this approximation reduces the sample efficiency of our sampler.It turns out that the choice of unsupervised and supervised experts plays a key role in determining the theoretical sampler efficiency.In particular, the following corollary to Theorem 3 from Sun et al. (2022b) relates the smoothness of each expert's gradient to our sampler's ability to efficiently explore protein space.
the max path length is U , and a 1-Hamming ball neighborhood N (x).Let Q R (x, x ) and QR (x, x ) be the Markov transition kernels induced by our sampler with the product of experts proposal q R (x |x) and with its approximation qR (x |x), respectively.These transition Algorithm 1 Plug and Play Directed Evolution (PPDE) input one-hot encoded wild-type protein x W T , unsupervised experts f i , supervised experts g j , scale λ, max path length U output evolved protein x * while still searching do define // compute the forward path proposal distribution 2 // sample a single amino acid substitution and apply it to x r−1 end while kernels are related by See Appendix B for the proof.
Remark: This result bounds the efficiency of the sampler by a product of exponential functions of each expert's gradient's Lipschitz constant K i .This means that if just one expert has a gradient with a large Lipschitz constant (e.g., the unsupervised experts are highly nonlinear protein language models), it is possible that this expert greatly reduces the overall efficiency of the sampler.For supervised experts (possibly an ensemble), which are typically shallow neural networks whose gradients have small Lipschitz constants, we can expect that the first-order approximation will not greatly reduce the sampler's efficiency.However, Equation 6 is a fairly loose bound.We will empirically compare how the sampler fares with different types of experts in practice (see Section 5.2.3).

Experiments
In this section, we present the results of multiple synthetic experiments.First, we validate that our proposed sampler is able to discover a diverse set of good optima for product of experts distributions in high-dimensional discrete spaces using a toy MNIST-based task.Then, with a realistic in silico directed evolution experimental setup, we characterize advantages of combining unsupervised and supervised models and compare our sampler with appropriate baseline plug and play algorithms.The baselines are: • Simulated annealing: This is a simple random-mutation-based MCMC-style algorithm from Biswas et al. (2021) for sampling from the global optimum of a discrete Boltzmann distribution p = 1/Z exp (−y/T ) with temperature T and y = − log p(x).
We apply this to sample from our product of experts ( Equation 4).At each sampling step, a new sequence is proposed by randomly sampling m ∼ Poisson(µ − 1) + 1 uniformly sampled amino acid substitutions with rate µ ∼ Uniform(1, 2.5).Proposed variants are accepted with probability min(1, exp((ŷ − y)/T )) with T annealed over time to encourage the sampler to become more exploitative.Variants with higher log p(x) are always accepted whereas sequences that decrease log p(x) are accepted with a probability proportional to the difference scaled by the temperature.
• Random sampling: We provide a simple baseline that does not use search but rather uniformly samples variants around the WT protein, with aims of quantifying the advantage of using search to accumulate multiple promising mutations.Given a budget of N variants, we use the same mutation proposal algorithm as simulated annealing to sample and apply a single mutation (which could consist of multiple amino acid substitutions) to the WT.After sampling N such variants, we sort them by log p(x) and return the top K.
• MALA-approx : This baseline employs a continuous relaxation of discrete protein space to use stochastic gradient-based continuous optimization inspired by Nguyen et al. (2017).During each step, MALA-approx samples where is a step size and ∇ x log p(round(x)) is a straight-through estimate of the gradient (Bengio et al., 2013) of the product of experts with respect to the current relaxed variant.The continuous relaxation is based on a temperature-controlled relaxation of categorical variables to continuous vectors on the simplex (Jang et al., 2016;Maddison et al., 2016) with temperature τ .That is, we interpret each sequence position of protein x to be a rounded sample from a relaxed categorical distribution.We run this sampler in "logit space" of these distributions since the support of logit space extends across the entire real line, whereas enforcing valid probability distributions at each sequence position would require solving a difficult constrained optimization problem.Given the one-hot WT protein x W T , the sampler is initialized at logit values given by a convex combination of a uniform distribution over all amino acids and the wild type: log((1 To round to discrete values for evaluating log p(round(x))), we sample from the relaxed categorical distributions with logits x and then take the argmax at each sequence position.
• CMA-ES: The CMA-ES algorithm (Hansen and Ostermeier, 1996) is a powerful evolutionary strategies optimizer that accelerates search by maintaining an estimate of higher order relational statistics between each dimension of the search space.This algorithm is designed for continuous spaces with tens or hundreds of dimensions to eventually concentrate around the global optimum.It iteratively updates a Gaussian distribution (both mean and full covariance matrix) by sampling and evaluating a population of candidates.Like MALA-approx , it requires using a continuous relaxation of protein space.Simply, we "flatten" a one-hot protein x into a vector of dimension LV which we then re-interpret as a continuous vector in R LV .To undo this relaxation, we reshape the vector into a length L sequence and interpret the V entries at each sequence position as scores, i.e., we take the argmax.
We use multiple metrics to compare samplers that measure not only activity but also whether designed proteins are evolutionarily plausible.Our results provide strong evidence that combining unsupervised evolutionary models with supervised models produces variants more likely to appear in nature and with higher fitness, especially when using unsupervised models fit on aligned sequences.Moreover, we find that composing multiple unsupervised models trained on aligned and unaligned sequences has the best performance.We also verify that, by and large, our gradient-based discrete MCMC sampler for product of experts achieves better sampling performance than the random-walk-based and evolutionary-strategies-based baselines.The objective is to evolve the wild type image (right) to maximize the sum of the two digits.

MNIST-Sum
Before we present results on ML-based directed evolution, we evaluate the samplers on a toy search problem where optimized designs are easy to interpret.Setup: We construct a binary MNIST task that emulates a typical ML-based directed evolution problem.We are given two binary MNIST images x 1 and x 2 where x 2 is the wild type image.Each image is treated as a length L = 784 one-hot sequence where V = {0, 1}.The goal is to evolve x 2 , keeping x 1 fixed, to maximize the sum of the two digits (Figure 2).Here, the equivalent of an amino acid substitution is flipping a binary pixel.This task is difficult for methods that optimize in discrete space, since flipping most pixels causes the image to "fall off" the underlying low dimensional manifold of MNIST digits.To obtain supervised experts p g (x) for a product of experts log p(x) (Equation 4), we use a labeled dataset of 50K pairs of binary MNIST digits with sum ≤ 10 (the "assay-labeled proteins") to train an ensemble of three shallow siamese ConvNets that regresses the sum x 1 + x 2 .We then use an unlabeled datasest of 50K binary MNIST images (the "unlabeled protein database", such as UniRef) to train a deep EBM and to train a DAE for unsupervised experts p f (x) (see appendix for model details).
PPDE is compared against simulated annealing, MALA-approx , and CMA-ES.All methods use a single EBM unsupervised expert and ensemble of three CNN-based su- x 1 = 1, x 2 = 5 (which has maximum sum 10 since x 2 is evolved while x 1 is held fixed) for 10K steps.We test each sampler on the out-of-distribution image pair x 1 = 9, x 2 = 1 with maximum possible sum 18.In this case, "solving" the task is equivalent to evolving the wild type "1" digit into a "9".Each method evolves a population of size 128 and each sample in the population is optimized for a budget of 20K steps.To estimate "ground truth" sums for optimized image pairs, we train an ensemble of three siamese ConvNets to near perfect accuracy on pairs of digits whose sums are ≤ 18.
Results: The qualitative and quantitative results in Figure 3 show that composing the supervised experts with an unsupervised expert (the EBM or the DAE) is necessary to solve the task; otherwise, the PPDE sampler is unable to avoid adversarial inputs and each run of the sampler converges to white noise images.We also observe that PPDE samples a highly diverse population of digits with noticeable semantic differences.We see that PPDE (DAE) outperforms PPDE (EBM), which we suspect could be caused by the DAE's gradients being more informative due to its denoising training objective.Simulated annealing and MALA-approx fail in this high-dimensional space due to the difficulty of using random walks to find search directions that increase the product of experts probability.CMA-ES is outperformed by PPDE despite its use of covariance information to accelerate search in high-dimensional spaces.PPDE's evolved images are considerably more realistic than those produced by CMA-ES.We also observe that on this task, CMA-ES produces images that show low semantic diversity (the images differ at the pixel-level but the by and large look highly similar).We speculate that PPDE gains an advantage over CMA-ES by using gradients to select which pixels to mutate instead of approximate high-order statistics.

In silico Directed Evolution
Datasets: We use three benchmark proteins with partially characterized wide fitness landscapes and MSAs from Hsu et al. (2022) for in silico directed evolution experiments: the Poly(A)-binding protein (PABP) dataset of variants measuring binding activity (95 residues, each variant has ≤ 2 mutations), the ubiquitination factor E4B (UBE4B) protein dataset measuring ligase activity (103 residues, each variant has ≤ 6 mutations), and GFP protein dataset measuring fluorescence (238 residues, each variant has ≤ 15 mutations).To emulate a realistic protein engineering setup with reasonable amounts of data for training our supervised experts g(x), we use the "2-vs-rest" mutation train/test split as suggested in Dallago et al. (2021)-that is, after splitting each dataset 80/20, we keep only sequences with two or fewer mutations for training and subsample 10% of these.For PABP, both the train and test sets only have variants with at most 2 mutations.This amounts to 3K sequences for PABP, 2K for UBE4B, and 1K for GFP.
Metrics: Quantitative evaluation of optimized variants involves measuring the activity (log fitness), the likelihood to appear in nature (evolutionary density), population diversity, and ability to explore the fitness landscape around the WT (average number of mutations per variant).To compute log fitness, we take inspiration from model-based optimization benchmarks (Trabucco et al., 2022) and simulate an expensive wet lab verification with a powerful "oracle" model that we train with 80% of all variants (with all numbers of mutations) for each protein.Log fitness is scored relative to WT with the Augmented EVmutation Potts model from Hsu et al. (2022) as the oracle.To compute an evolutionary density score relative to WT, we use a different transformer, the MSA Transformer (Rao et al., 2021), conditioned on 500 randomly subsampled sequences from each MSA.For both log fitness and evolutionary density, positive values mean the variant outscored the WT while negative values are worse than WT.We run each sampler 128 times for 10K steps each, initialized at the WT, and select the sample with the maximum log p(x) per run for the final population.As plug and play baselines we again use simulated annealing, CMA-ES, and MALA-approx , as well as random sampling.
Experts: We use the EVmutation Potts and ESM2 family of protein language models as pre-trained unsupervised experts.All baselines use a single Potts expert and are directly compared with PPDE (Potts).Unless stated otherwise, we use the 150M parameter version of ESM2.We also run our sampler with both Potts and ESM2 (i.e., we take the sum of the Potts and ESM2 evolutionary density scores) (PPDE (Potts+ESM2)).For the supervised expert we use an ensemble of three shallow CNNs introduced as a baseline in Dallago et al. (2021)  architecture details and details on how we implement the evolutionary density scoring for the Potts and protein language model scoring functions.
Choosing λ: Multiplying the supervised experts with the hyperparameter λ ≥ 0 trades off how much the unsupervised and supervised experts influence which variants are sampled from the product of experts distribution.A larger value for λ inflates the supervised expert scores which causes the sampler to have more confidence in variants with high predicted activity.We use a simple heuristic to select λ for the three benchmark proteins and for each of the Potts and ESM2 unsupervised experts.The heuristic roughly assigns equal emphasis to the unsupervised and supervised expert scores.In detail, we estimate the {min, max, mean} statistics of expert scores computed from a small representative sample of labeled protein variants and then picks λ such that the statistics for the supervised expert scores are close to those of the unsupervised expert scores.We sample 100 protein variants from the training set that have lower activity than WT and 100 proteins with higher activity higher than WT and then evaluate the {min, max, mean} of the unsupervised expert scores and λ×{min, max, mean} of the supervised expert scores for λ ∈ {0.5, 1, 3, 5, 15}.When using the Potts unsupervised expert we use λ = 5, 0.5, 15 and when using the ESM2 expert we use λ = 5, 3, 1 for the PABP, UBE4B, and GFP proteins, respectively.Less manual approaches that tune λ over a grid of values or learn λ over a held-out validation set of labeled variants could be used instead if desired (Holtzman et al., 2018).

Main results
Figure 4 shows the 80 th percentile metrics for each evolved population colored by the population's average number of mutations.The 50 th and 100 th percentile results for all samplers are provided in a table in the appendix (Table 2).Diversity scores are in Table 1.We make the following observations.
Diversity: Across all three proteins, PPDE-based samplers achieve the best diversity (highest percentage of unique variants), implying superior exploration of good optima around the WT protein.
Evolutionary density, log fitness, and mutation counts: Out of the samplers using the Potts unsupervised expert, PPDE (Potts) discovers variants with higher log fitness and average number of mutations than random search, simulated annealing, and MALA-approx .In terms of evolutionary density, PPDE (Potts) achieves the highest scores on PABP, and on the more challenging UBE4B and GFP proteins, we suspect that the slightly lower 80 th percentile evolutionary density scores compared to the baselines (except for CMA-ES, which scored poorly) are in part due to PPDE discovering variants with higher average mutations (2 − 4+ mutations vs. ∼ 1 mutation for the baselines) and higher log fitness.We found that the CMA-ES sampler had inconsistent performance across the three proteins.Among samplers using the Potts unsupervised expert, it achieves the highest fitness scores on two proteins (PABP and UBE4B), but does so with diversity scores of 0.8% and 3.1% compared to 85.2% and 12.5% for PPDE (Potts) (Table 1).CMA-ES finds variants with high numbers of mutations (∼ 10 − 17) but low evolutionary density scores (50 th percentile scores of 3.47 on PABP, for PPDE; see Table 2).It finds just one protein variant for PABP that is 17 mutations from WT with an evolutionary density score of 3.47, and it seems to have significant difficulty with the larger proteins UBE4B and GFP; e.g., on GFP the average 50 th percentile log fitness is -2.50 compared to -0.04 for PPDE. Figure 5 shows sampler trajectories, with which we see that across all three proteins, PPDE is the superior approach for efficiently sampling from the high-dimensional discrete product of experts distributions.We see here again that CMA-ES performs worse as the protein sequence length and fitness landscape complexity increases.It appears that a major limitation of CMA-ES is its continuous relaxation, which projects each sequence in the population to a variant far from the WT.For GFP, the population average product of experts probability starts and stays around log p(x) ≈ −50.We conclude that CMA-ES can be recommended for use only when a single protein variant is desired and the length of the protein is relatively small (e.g., ≤ 95 residues).Swapping unsupervised experts: Without using an unsupervised expert, PPDE discovers variants with worse evolutionary density scores than even random sampling, confirming our observations from the previous MNIST-Sum experiment that supervised experts alone are susceptible to adversarial inputs.We also examined PPDE without any supervised expert (PPDE (Potts only)) by setting λ = 0.The observation that PPDE (Potts only) achieves a higher log fitness on PABP and only slightly worse log fitness on UBE4B than PPDE (Potts) is not surprising, as recent evidence suggests unsupervised evolutionary density scores are (at least moderately) predictive of mutation effects (Meier et al., 2021;Hsu et al., 2022;Weinstein et al., 2022).However, the sampler performance degrades significantly by all metrics on GFP.When using the ESM2 unsupervised expert, PPDE discovers variants with lower evolutionary density scores than when using the Potts.This is not altogether unexpected since the MSA Transformer, which is used to score the variants, is conditioned on the same MSA as used to fit the Potts model.Although, this result also suggests that using unsupervised experts fit to aligned sequences is more promising for ML-based directed evolution.We observe a promising compositional effect from combining unsupervised experts trained on unaligned sequences (ESM2) and aligned sequences (Potts); e.g., across all three proteins the PPDE (Potts+ESM2) sampler achieves higher 100 th percentile fitness and density scores than when using the Potts or ESM2 experts alone, and increases the diversity and average number of mutations as well.

Qualitative results
We provide a qualitative analysis of the mutations discovered by the supervised only, simulated annealing, PPDE (Potts) and PPDE (Potts+ESM2) samplers in Figure 6.Many The marginal frequency of a mutation site is the fraction of protein variants in the population with a mutation at a particular sequence position.We visualize site frequencies as a heatmap on the 3D structure of the WT protein.For example, nearly 100% of UBE4B variants discovered by PPDE (Potts+ESM2) have a mutation at position 1145.We annotate only a few high frequency sequence positions per protein for clarity; many sequence positions with a low frequency are not annotated.
of the variants found by PPDE tend to have the same 1-3 "highly plausible" mutations in addition to a few "rare" mutations that show up in variants at a much lower rate.Figure 7: Swapping in different protein language model experts.Full populations (128 variants) colored by log fitness with the median evolutionary density annotated.We swap out different unsupervised experts for the in silico directed evolution of UBE4B using the same supervised expert.This includes a 650M parameter protein language model (ESM2-650M).Combining unsupervised experts trained on aligned and unaligned sequences (Potts+ESM2-150M) discovers variants with highest density and higher predicted fitness than WT (arrow).

Comparing protein language model experts
or use language models fine-tuned for specific protein families.As the number of ESM2 parameters increase, we observe a trend where variants with higher predicted activity and marginally higher evolutionary density are discovered.A similar observation in Nijkamp et al. ( 2022) also suggests that model size positively correlates with zero-shot fitness prediction on wide fitness landscapes.When directly comparing supervised only and PPDE (Potts), we can see a large improvement in evolutionary density when the Potts expert is used.ESM2, which is trained on unaligned proteins, achieves lower evolutionary density scores than the Potts, suggesting it acts as a much softer search constraint.Combining the Potts and ESM2 experts results in improving the top percentile of discovered variants compared to those found by using the Potts or ESM2 experts alone.

Path length sensitivity analysis
PPDE's proposal distribution samples a path of R ∼ Unif(1, U ) single substitutions at each sampler step to rapidly navigate through protein space.We conducted a sensitivity analysis on the maximum path length U , where U = 2X − 1 and X ∈ {1, 2, 5, 10} (Figure 8c) on the UBE4B protein.In this experiment, we found that path lengths longer than 3 did not significantly improve sampler efficiency.This is possibly because we already have a good initialization for the sampler (the WT) and short path lengths are sufficient for discovering high quality variants.It is also possible that, in our setting, the accumulation of errors from the first order Taylor approximation along a long path leads to worse sampler efficiency despite better exploration.Moreover, using a longer path length leads to discovering variants with a higher number of mutations, whose activity could be inaccurately characterized by the supervised experts.

Run time comparison
We compare wall clock run times for PPDE and the baseline samplers under the Potts (Figure 8a) and ESM2 (Figure 8b) experts using the PABP protein, 1K sampler steps, and a population size of 16 variants.Samplers are run on a single NVIDIA Tesla cloud GPU; see appendix for complete hardware details.CMA-ES is significantly slower than the other samplers.When using the Potts expert, PPDE's run time is relatively comparable to simulated annealing and MALA-approx .The cost of computing the gradient of ESM2 is much higher than for the Potts model, so in this setting simulated annealing (which does not use gradients) has lower run time than MALA-approx and PPDE.

Limitations
In this section, we discuss limitations of the proposed sampling method related to computational costs and handling insertion/deletion mutations.Each step of the PPDE sampler requires two backwards passes-one for computing the forward proposal and one for the reverse proposal.When PPDE is used with a large protein language model whose backwards pass has a large memory footprint, this becomes a bottleneck if sufficient computing resources are unavailable.Fortunately, MCMC samplers are embarrassingly parallel.For example, with a budget of K available GPUs, we can evolve a single protein per GPU in parallel.
The wall-clock time incurred by running PPDE depends on the ability to rapidly evaluate and differentiate through the product of experts distribution.For certain classes of unsupervised experts, this is difficult.For example, evaluating the log probability of a protein with a variational autoencoder (e.g., DeepSequence (Riesselman et al., 2018)) requires computing a large-sample Monte Carlo estimate of the evidence lower bound, which is highly computationally expensive.
PPDE also does not currently support inserting and deleting amino acids (indels).We necessarily assume the product of experts is locally smooth in a neighborhood around the current protein to compute our proposal distribution for MCMC.This assumption would need to still hold despite proposing length changes to the sequence.While we do not believe lack of support for indels is a fundamental limitation of the sampler, supporting indel mutations appears to be nontrivial.

Conclusions
In this study, we have shown how to flexibly combine unsupervised models of evolutionary density, including protein language models, and supervised models of protein function and how to efficiently sample from the resulting distribution to discover proteins that maximize a desired function while avoiding poor local optima.This strategy leverages the vast amounts of unlabeled data that are available for unsupervised (or self-supervised) pre-training to improve designed sequences, even when relatively few labelled data are available for training the fitness function.Our empirical results suggest our framework offers a practical and effective new paradigm for machine-learning-based directed evolution.Although our framework can potentially be applied to a broader range of biological sequence design tasks, we focused solely on proteins due to the variety of available pre-trained models.
Based on our findings, we recommend constructing the product of experts distribution with unsupervised experts that have been pre-trained on aligned protein sequences (e.g., MSAs) when available.We found that these models more strongly constrain search to the evolutionarily plausible mutations than unsupervised models pre-trained on unaligned protein sequences.However, combining unsupervised models of both types produced designs superior to using one or the other alone.
One important capability that we did not demonstrate is the incorporation of additional hard constraints, such as a maximum allowable number of mutations per variant or the preservation of particular regions of the WT sequence.In our framework, we can easily accomplish this by forcing the mutation proposal probabilities at the sequence positions in question to zero, which can be done by setting their logits to negative infinity in the softmax in Algorithm 1. Future work may also extend this framework to larger problems in biological design.For instance, the simultaneous engineering of several sequences in multimeric enzyme complexes, or incorporating substrate structure in evaluating the likelihood of enzyme-substrate complexes.

A. Deriving the Gradient-based Approximate Proposal
We provide a short derivation that illustrates how to obtain a gradient-based approximate locally-balanced informed proposal (Equation 2) from the locally-balanced informed proposal of Equation 1.The goal is to avoid evaluating f (x ) for all x ∈ N (x) when computing the proposal distribution q(x |x) ∝ exp(f (x ) − f (x)) 1 2 1(x ∈ N (x)).The basic idea from Grathwohl et al. (2021) is to assume that f is a continuously differentiable function, which allows us to approximate f (x ) − f (x) with a first-order Taylor series.
In detail, the first-order Taylor series of the function f (x) at the point a is By subtracting f (a) from both sides, we get To arrive at the gradient-based proposal, suppose a is the current state of our MCMC sampler and let x be any point in the neighborhood of a, x ∈ N (a).When we plug in our first-order approximation,

B. Proof of Corollary 1
The basic idea of the proof is to use the triangle inequality to obtain a bound on the approximation error of a sum of experts which assumes that each expert has sufficiently smooth gradients.Definition A function f : R N → R has K-Lipschitz continuous gradient when for all x, x ∈ R N .For convenience, we reproduce a pertinent result here from Nesterov (1998) (Lemma 1.2.3).
Lemma 1.2.3 Nesterov (1998) If f : R N → R has an L-Lipschitz gradient, then for any x, x ∈ R N we have: Lemma 2 For functions f and g with K f -Lipschitz and K g -Lipschitz gradients respectively, the sum-composition h = f + g has (K f + K g )-Lipschitz gradient.
Proof : By the triangle inequality: Lemma 3 Suppose f i , i = 1, . . ., M are functions with K i -Lipschitz gradient.Then for any x, x ∈ R N we have Proof: Let g = M i=1 f i where each f i , i = 1, . . ., M has K i -Lipschitz gradient.By Lemma 1 we can see that g has a M i=1 K i -Lipschitz gradient.Then Lemma 2 follows immediately by applying Lemma 1.2.3 from Nesterov (1998) to g.
The proof of Corollary 1 proceeds by bounding the approximation error between two consecutive states x r−1 and x ∈ N (x r−1 ) in a path of length R ∼ Unif(1, U ).For simplicity we assume N (x) is the 1-Hamming ball, i.e., x r − x r−1 2 = 1. For Then an upper bound for g(x ) − g(x r−1 ) is Following similar steps, we also have The remainder of the proof for Corollary 1 exactly follows Equations 64-72 in the proof of Theorem 3 in Sun et al. (2022b) with g(x) which has K-Lipschitz gradient.

C. Experiment details
C.1 Neural architectures MNIST denoising autoencoder: The encoder and decoder neural networks are constructed out of residual blocks, where a single block has two 3x3 2D Conv layers, each followed by a BatchNorm layer and Swish activation.The first Conv layer in the block uses a stride of 2 and a 1x1 shortcut Conv layer is added to its output.The encoder consists of one 3x3 2D Conv layer, followed by two blocks, then another 3x3 2D Conv layer, followed by a fully connected layer that projects the flattened output features into a 16-dimensional latent space.The encoder and decoder are symmetric, and the decoder is implemented with transposed Conv layers and the per-pixel output is interpreted as the logits of a Bernoulli distribution.All Conv layers use 64 channels.To train the DAE, we corrupt the input 28 × 28 binary MNIST image by randomly flipping p% of the pixels, where p ∼ Unif(0,15), and minimize the reconstruction error.We use the AdamW optimizer with a learning rate of 1e-4 to train the model for 10K steps using a mini-batch size of 100.MNIST energy-based model: This model follows the residual EBM architecture from Grathwohl et al. ( 2021) closely.The model consists of one 3x3 2D Conv layer, followed by 2 residual blocks (the same blocks used by the DAE), and then an additional 6 residual blocks with all Conv layers using a stride of 1.The output features are flattened with global average pooling and then mapped to a scalar with a fully connected layer.See Grathwohl et al. (2021) for details about the contrastive divergence training algorithm and training hyperparameters.Protein ConvNet: We use the ConvNet baseline from the FLIP benchmark (Dallago et al., 2021) to regress activity.This model takes in a mini-batch of one-hot encoded protein sequences of shape B × L × V , applies a 1D Conv with kernel size 5 and ReLU activation to project V to dimension L, followed by a fully connected layer to expand the feature dimension to 2L.Then, max pooling is used to reduce the the length of the sequence to 1, and a linear layer projects the feature dimension from 2L to 1.The model is trained with AdamW optimizer with a learning rate of 1e-3 using a mini-batch size of 256.

C.2 Unsupervised expert score functions
We provide details for computing i f i (x) with each type of unsupervised expert.MNIST denoising autoencoder: For a given image x, we define f (x) to be the sum of the binary cross entropy between the input pixel and the predicted pixel logits over all pixel locations.MNIST energy-based model: We use the scalar output as f (x), since it can be interpreted as an unnormalized log probability for the input.This only requires a single forward pass to compute.Protein EVmutation Potts: We use the difference in the Potts Hamiltonian between the protein variant x and the wild type such that f (x) = H(x) − H W T (x).The Hamiltonian for a one-hot protein x is defined as x T i J ij x j where J ∈ R L×L×20×20 and h ∈ R L×20 are the Potts parameters.Protein ESM2: We use the difference in the sum of the per-amino-acid log probabilities between the protein variant x and the wild type for f (x).For computational efficiency, we compute the score with a single forward pass and therefore do not use any masking of input sequence positions.In detail, for a one-hot protein x let φ(x) be the L × V matrix of logits computed by a single unmasked forward pass of ESM2.Then the score is x ij log softmax(φ(x) i ) j − L i=1 V j=1 x W T ij log softmax(φ(x W T ) i ) j .

C.3 Hardware details
We ran our experiments on the Summit supercomputer on a single NVIDIA V100 GPU.
When running with ESM2, we speed up the sampler by parallelizing the population of 128 proteins across multiple GPUs.For example, we ran each of the 128 MCMC trajectories on a separate GPU in parallel when we used the 650M parameter ESM2 model.Wall clock run time measurements were taken after-the-fact using a single NVIDIA Tesla T4 GPU available via the Google Colab free tier.

C.4 Aligning the amino acid vocabulary across models
Not all pre-trained protein sequence models order the amino acids in their vocabulary in the same way.Also, some models, such as masked language models, have extra tokens such as <mask> beyond the standard 20 amino acids.To enable combining a diverse array of models, we first decide on a canonical ordering for the standard amino acids.For any pre-trained model that uses a different ordering than the canonical one, we pre-compute a permutation matrix that, when applied to the columns of the one-hot encoded protein x, will put the amino acids in the order which the model expects.We also pad the permuted x with extra columns of zeros if the pre-trained model has a vocabulary larger than 20. in particular discover variants with higher predicted fitness and average # of mutations than Random search, Simulated annealing, and MALA-approx .Out of the samplers using the Potts expert, PPDE (Potts) achieves the highest evolutionary density scores on PABP.We suspect that the slightly lower 50 th percentile evolutionary density scores compared to the baselines (except CMA-ES) on the more challenging UBE4B and GFP proteins are in part due to PPDE discovering variants with more average mutations (2 − 4+ mutations vs. ∼ 1 mutation per variant) and higher fitness.The 100 th percentile performance of PPDE (Potts+ESM2) shows impressively high log fitness and evolutionary density, however.CMA-ES sampler had inconsistent performance across the three proteins-it finds variants with high numbers of mutations (10 − 17) but low evolutionary density scores (3.47 on PABP, -94.76 on UBE4B and -62.43 on GFP compared to 6.92, -4.79, and -5.98 for PPDE (Potts)).It finds just one protein variant for PABP that is 17 mutations from WT with an evolutionary density score of 3.47, and it seems to have significant difficulty with the larger proteins UBE4B and GFP; e.g., on GFP the average 50 th percentile log fitness is -2.50 compared to -0.04 for PPDE.We conclude that CMA-ES can be recommended for use only when a single protein variant is desired and the length of the protein is relatively small (e.g., ≤ 95 residues).

Figure 2 :
Figure 2: MNIST-Sum:The objective is to evolve the wild type image (right) to maximize the sum of the two digits.

Figure 3 :
Figure 3: MNIST-Sum results: a) Random samples from the optimized digit populations.The samples from PPDE (DAE) have the highest diversity and also closely resemble MNIST images of a "9".PPDE (None) does not use an unsupervised expert which causes sampler trajectories to become trapped in adversarial regions of the space where samples resemble white noise.The random mutation proposals employed by simulated annealing and MALA-approx cause these samplers to fail due to the high-dimensionality of the space.CMA-ES's samples have low diversity yet converge onto a population of images that somewhat resemble a "9".b) The supervised experts' predicted sum (orange) vs. the oracle's predicted sum (blue) for the sampler trajectories.Each line plot is the 70 th percentile sum with the shaded region showing the 50 th and 90 th percentiles.PPDE (DAE) reliably maximizes the sum of the digits ("18") as measured by the oracle model.

Figure 4 :
Figure 4: In silico directed evolution 80 th percentile results.Positive log fitness and evolutionary (evo.)density scores are better than the WT while negative values are worse than WT.The red dashed lines intersect at WT performance (0, 0).Each sampler is colored by the average number of mutations per variant (higher is better).On PABP (left), PPDE (Potts+ESM2) and PPDE (Potts only) (arrows) achieve the highest evo.density and log fitness-both better than WT-and with relatively high numbers of mutations.On UBE4B (center), PPDE (Potts) and PPDE (Potts only) (arrows) appear to best trade off evo.density, log fitness, and mutation count.PPDE (ESM2) achieves the best log fitness but with slightly worse evo.density scores.On GFP (right), only PPDE (Potts) and PPDE (Potts+ESM2) (arrows) discover samples with log fitness better than WT.CMA-ES is omitted for viewing clarity due to outlier results on UBE4B (log fitness: 2.54, evo.density: −94.76) and GFP (log fitness: −1.13, evo.density: −22.85).Best viewed in color.

Figure 5 :
Figure5: PPDE is most efficient at finding good optima of the high-dimensional product of experts.Cumulative maximum product of experts log probability averaged across the population (higher is better).All samplers use the same product of experts target distribution with Potts unsupervised expert (other unsupervised experts not shown for clarity of presentation).CMA-ES is not visible in c) because the optimization starts and remains around log p(x) ≈ -50.

Figure 6 :
Figure6: Visualization of mutation site frequency.Best viewed in color and zoomed in.The marginal frequency of a mutation site is the fraction of protein variants in the population with a mutation at a particular sequence position.We visualize site frequencies as a heatmap on the 3D structure of the WT protein.For example, nearly 100% of UBE4B variants discovered by PPDE (Potts+ESM2) have a mutation at position 1145.We annotate only a few high frequency sequence positions per protein for clarity; many sequence positions with a low frequency are not annotated.

Figure 7
Figure7highlights the ability of PPDE to successfully sample from the 35M, 150M, and 650M parameter ESM2 language models without any extra fine-tuning or re-training.This flexibility makes it easy to leverage continued improvement in protein language models,

Figure 8 :
Figure 8: Wall clock run time comparisons using the Potts (a) and ESM2 (b) unsupervised experts.Run times are averaged over 7 trials using the PABP protein, 1K sampler steps, and a population of size 16.c) Comparing PPDE on UBE4B with various max path lengths U = 2X − 1 where X ∈ {1, 2, 5, 10}.

Table 1 :
which regress activity directly from one-hot encoded proteins.See the appendix for Population diversity (% of variants that are unique).