SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jeffrey Ouyang-Zhang
jozhang@utexas.edu
Adam R. Klivans
Daniel J. Diaz
danny.diaz@utexas.edu
Philipp Krahenbuhl
philkr@cs.utexas.edu
Stabilizing proteins is a foundational step in protein engineering. However, the evolutionary pressure of all extant proteins makes identifying the scarce number of mutations that will improve thermodynamic stability challenging. Deep learning has recently emerged as a powerful tool for identifying promising mutations. Existing approaches, however, are computationally expensive, as the number of model inferences scales with the number of mutations queried. Our main contribution is a simple, parallel decoding algorithm. Our Mutate Everything is capable of predicting the effect of all single and double mutations in one forward pass. It is even versatile enough to predict higher-order mutations with minimal computational overhead. We build our Mutate Everything on top of ESM2 and AlphaFold, neither of which were trained to predict thermodynamic stability. We trained on the MegaScale cDNA proteolysis dataset and achieved state-of-the-art performance on single and higher-order mutations on S669, ProTherm, and ProteinGym datasets. Our code is available at https://github.com/jozhang97/MutateEverything.
Protein engineering is the discipline of mutating a natural protein sequence to improve properties for industrial [5, 78] and pharmaceutical applications [2, 25, 41]. However, evolution simultaneously optimizes several properties of a protein within its native environment, resulting in proteins with marginal thermodynamic stability ( $\sim$  5-15 kcal/mol) [39] which become non-functional in an industrial setting. Therefore, repurposing a natural protein for biotechnological applications usually begins with identifying non-deleterious mutations that stabilize the structure. With a stabilized structure, downstream engineering goals, which often require exploring destabilizing mutations, become tractable. Historically, this process has been bottlenecked by the requirements of extensive laboratory characterization of rational designs [33, 79] or directed evolution libraries [4, 12, 21, 22]. The recent explosion in biological data [9, 29, 72, 73] has enabled deep learning frameworks to accelerate the identification of stabilizing mutations. A successfully stabilized protein often requires several mutations. However, current frameworks do not take into account epistatic interactions between multiple mutations. Thus, to significantly accelerate protein engineering, it is critical to efficiently navigate the epistatic landscape of higher-order mutations [10, 69]. However, due to its combinatorial nature, thorough exploration quickly becomes computationally prohibitive.
In this paper, we introduce Mutate Everything to directly predict changes in thermodynamic stability  $(\Delta \Delta G)$  for all single and higher-order mutations jointly. Mutate Everything is a parallel decoding algorithm that works in conjunction with a sequence model for predicting thermodynamic stability.
Figure 1: Mutate Everything efficiently predicts  $\Delta \Delta G$ , the change in thermodynamic stability of folding, for over a million mutations (e.g. all single, double mutations) in a single inference step. This helps identify and prioritize stabilizing mutations ( $\Delta \Delta G < 0$ ) in protein engineering efforts. The notation for a mutation is  $AA_{from}$ , pos,  $AA_{to}$  (e.g. "P1K" mutates from P to K at position 1)
Prior models for predicting thermodynamic stability work by producing an embedding or representation of the input sequence. Our decoder takes this representation and uses a linear network to create further representations, one for every possible single mutation. These mutation-level representations are then further aggregated to form representations for higher-order mutations. We feed these higher-order representations into a lightweight multi-layer perception (MLP) head to output predicted  $\Delta \Delta G$  measurements. Since the mutation-level representations are computed only once, we are able to scale to millions of (higher-order) mutations, as the aggregation and MLP-head computations are inexpensive. As such, given a fixed computational budget, our model evaluates millions more potential mutations than prior methods.
Mutate Everything estimates the effects of all single and double mutations for a single protein in seconds on one GPU. This can efficiently compute the stability of all single and double mutants across all  $\sim 20,000$  proteins in the human proteome within a short time frame. To the best of our knowledge, Mutate Everything is the first tool that renders the computational analysis of double mutants across the entire proteome tractable. Mutate Everything can be used in conjunction with any model that generates representations for protein sequences. In this paper, we use AlphaFold [31] as a backbone architecture and fine-tune it on a dataset labeled by experimentally derived  $\Delta \Delta G$  measurements. This results in the first accurate model of stability prediction based on AlphaFold.
We evaluate the performance of our Mutate Everything on well-established benchmarks for predicting experimentally-validated  $\Delta \Delta G$  measurements of mutations: ProTherm [48], S669 [54], and ProteinGym [49]. On ProTherm high-order mutations (PTMul), our model achieves a Spearman of 0.53, compared to 0.50 in the next best method. On S669 single mutations, our model achieves a Spearman of 0.56, outperforming the state-of-the-art at 0.53. On ProteinGym, the Mutate Everything outperforms state-of-the-art methods from 0.48 to 0.49. Our model makes finding the most stabilizing double mutations computationally tractable. We show on the double mutant subset of the cDNA-proteolysis dataset (cDNA2) [72], where only  $3\%$  of known mutations are stabilizing, that Mutate Everything effectively ranks stabilizing mutations ahead of destabilizing ones.
2.1 Protein Engineering
Protein engineering is the process of mutating a natural protein to improve particular phenotypes, such as thermodynamic stability and function. The field has historically relied on rational design and stochastic methods, such as error-prone PCR [1], DNA shuffling [71], and directed evolution (DE) [4, 27], to identify gain-of-function mutations. Rational design is limited to proteins with solved structures and requires an extensive understanding of biochemistry and specific knowledge of the particular protein to select mutations highly likely to improve the target phenotypes [33, 79]. Directed evolution requires little to no knowledge of the protein and instead generates a library of protein variants that are then screened for the target phenotypes (e.g. fluorescence brightness, antibiotic
resistance, stability, activity) [4, 27]. The library can be generated via site-saturated mutagenesis of a handful of positions in the sequence [12] or DNA shuffling, in which the gene is digested into random fragments and reassembled into full-length sequences [71]. After screening the library, the most "fit" variant is then selected as the initial sequence for the next round of directed evolution. This iterative process repeats until it obtains a protein variant with the desired phenotypes.
Machine learning has demonstrated its ability to augment rational design and accelerate the stabilization of a variety of proteins [20, 28, 40, 52, 68]. Separately, machine learning-guided directed evolution (MLDE)[75] has been shown to improve the likelihood of obtaining the global fitness maximum by 81-fold compared to traditional DE [76]. MLDE has accelerated the engineering of several proteins, such as the enantioselectivity of enzymes for kinetic resolution of epoxides [23] and the activity and expression of a glutathione transferase [46]. Mutate Everything empowers the experimentalist to accelerate the stabilization of a protein for both rational design and MLDE.
2.2 Machine Learning for Protein Structure Prediction
Recent advances in machine learning have led to remarkable progress in protein structure prediction. AlphaFold [31] has demonstrated that deep learning is highly effective at predicting protein structures from a sequence by using evolutionary history via a multiple sequence alignment (MSA). AlphaFold passes the MSA and a pairwise representation of the sequence into Evoformer to capture the co-evolutionary patterns between residues. The Evoformer output is then processed by the Structure Module, which predicts the protein's structure. We challenge prior works postulating that AlphaFold cannot be used for stability prediction and show that fine-tuning these rich co-evolutionary and structural features yield highly performant stability predictors [53].
Evolutionary Scale Modeling (ESM) [38, 62] has shown that protein structure prediction can be performed without MSAs and specialized architectures by leveraging large transformers pre-trained on masked token prediction. Other works extended this masked pre-training framework, including MSA-Transformer [60] which incorporates a sequence's MSA as input, and Trancection [49], which develops a hybrid convolution and attention based autoregressive architecture. We show that fine-tuning these evolutionary representations excels at protein stability assessment without MSAs.
2.3 Protein Stability Assessment
Traditional approaches to protein stability assessment relied on a combination of physics-based methods, statistical analysis, and traditional machine learning techniques. Physics-based tools, such as FoldX, Rosetta, and SDM, utilize energy functions and statistical patterns to assess how mutations affect a protein's stability [32, 67, 77]. DDGun [44] directly inferred  $\Delta \Delta G$  from heuristics, including Blosum substitution scores, differences in interaction energy with neighboring residues, and changes in hydrophobicity. Many traditional machine learning approaches used support vector machines and decision trees with physics-based feature engineering [11, 14, 15, 17, 37, 66]. Others ensemble existing machine learning methods [35, 56, 59, 63, 64].
Recently, deep learning-based approaches have recently begun to outperform existing physics and traditional machine learning approaches [13, 19, 34]. ACDC-NN [6] trains an asymmetric convolutional neural network for predicting forward and reverse mutation effects. Thermonet [36] voxelizes and feeds both the wild-type and mutant protein structures into a 3D-CNN to regress a  $\Delta \Delta G$  value. PROSTATA feeds the wild-type and mutant protein sequence into a pre-trained ESM2 model and then regresses a  $\Delta \Delta G$  value [74]. Stability Oracle takes the native and mutant amino acid type along with the local structure as input and regresses  $\Delta \Delta G$  [18]. Several deep learning frameworks [64, 80, 81] model multiple mutations in addition to single mutations. In this paper, we develop a framework that models the protein to enable efficient enumeration of all mutation candidates.
Problem Setup. A protein  $w = (w_{1},\dots,w_{L})$  is a sequence of amino acids  $w_{l}\in AA$ , where  $AA = \{A,C,\ldots ,Y\}$  are the 20 different amino acid types encoded by genetic information. Let  $\mu = (p,t)$  denote a mutation substituting the amino acid  $w_{p}$  at position  $p\in [1,\dots,L]$  to amino acid type  $t\in AA$ . Our goal is to determine the change in thermodynamic stability  $\Delta \Delta G\in \mathbb{R}$  for a protein  $w$  under a large number of mutation sets  $(M_1,\dots,M_N)$  where a mutation set  $M_{i} = \{\mu_{k}\}_{k = 1}^{K_{i}}$ .
(a) Backbone
(b) Mutate Single
Figure 2: Mutate Everything decodes any mutation. The model takes a protein sequence of length  $L$  as input. (a) A pre-trained feature extractor (e.g. ESM [38], AlphaFold [31]) computes per-position features for each amino acid in the sequence (yellow). Each token is reshaped by an MLP into an array of all possible mutations at that location. (b) A single mutation is directly decoded to its  $\Delta \Delta G$ . (c) For a mutation set containing multiple mutations, selected mutations are aggregated and decoded to obtain a predicted  $\Delta \Delta G$ .
(c) Mutate Multiple
The mutation sets include higher-order mutations  $(K > 1)$ , which is beyond the scope of state-of-the-art single mutation predictors  $(K = 1)$  [18, 36, 74]. These higher-order mutations are strictly more challenging than single mutations. The model must understand the effect of each single mutation and also the epistatic interactions that arise between individual mutations. We introduce a method that captures these interactions for any mutation set and decodes a corresponding  $\Delta \Delta G$ .
Furthermore, existing paradigms cannot efficiently scale to a large number of mutations. Higher-order mutation models [16, 35, 44, 64] take a protein and its mutant as input and predict the change in stability. These methods would need to enumerate and sort over a million model evaluations to find the most stabilizing mutations for a small protein of 75 amino acids. Our model runs a fine-tuned backbone only once and a lightweight decoder once for each mutation in parallel.
Feature Extraction. We fine-tune AlphaFold and other pre-trained models to alleviate the scarcity of labeled data. AlphaFold's Evoformer learns co-evolutionary patterns, such as residue interactions, and the Structure Module learns structural information, such as solvent accessibility or residue-residue distance. We demonstrate this is connected to a protein's stability and how mutations affect stability.
Each element of the protein sequence is tokenized by its amino acid type  $w_{l}$  and positionally encoded. This produces a set of high-dimensional tokens for the transformer to take as input. The tokens are fed into a pre-trained model, which extracts per-position features  $(x_{l})_{l = 1}^{L}$  where  $x_{l}\in \mathbb{R}^{d}$ . The transformer feature extractor is finetuned with our training objective. This is illustrated in Figure 2a.
Our model uses per-position features  $(x_{l})_{l = 1}^{L}$  to determine the stability of a mutation set  $M$ . We distinguish between single point mutations  $K = |M| = 1$ , and higher-order mutations  $K > 1$ .
Both single and higher-order mutations share the same backbone architecture and encoding scheme, see Figure 2b. For each mutation  $\mu = (p,t)$  at position  $p$  to amino acid  $t$ , we compute a latent mutation representation  $z(\mu) \in \mathbb{R}^d$ , which captures the effect of the mutation on the protein. We decompose the mutation representation  $z(\mu) = f^{t}(x_{p}) + h^{t}$ , where sequence-dependent features  $f^{t} \in \{f^{A}, \dots, f^{Y}\}$  project the per-position feature according to the mutated amino acid type, and sequence-independent features  $h^{t} \in \{h^{A}, \dots, h^{Y}\}$  are embeddings unique to the amino acid types  $A, C, \dots, Y$ . Intuitively,  $f^{t}$  captures a contextualized feature at the position, and  $h^{t}$  captures information about the mutated amino acid type. This mutation representation builds up a feature representing the amino acid substitution.
Decoding single mutations. Figure 2b illustrates single mutation decoding. The change in thermodynamic stability  $\Delta \Delta G$  for a single mutation  $\mu$  is read out from its mutation representation  $z(\mu)$ . A lightweight linear head  $g^{1}:\mathbb{R}^{d}\to \mathbb{R}$  decodes this representation into the change in thermodynamic stability:  $\Delta \Delta G^{pr} = g(z(\mu))$ .
Figure 3: Mutate Everything decodes many mutations. Each token output by the backbone  $x_{l}$  is expanded to 20 tokens  $z(\mu)$ , each corresponding to one unique amino acid type (yellow with amino acid types). For each mutation set  $M$ , we first aggregate the features corresponding to each mutation's position and amino acid being mutated to (arrows). A lightweight head  $g$  then decodes the aggregated features to obtain the predicted  $\Delta \Delta G$ . These heads are illustrated in Figure 2b and Figure 2c. Parallel feature aggregation and decoding enable efficient stability prediction for millions of mutations.
Decoding higher-order mutations. Figure 2c illustrates higher-order mutation decoding. The  $\Delta \Delta G$  for a higher-order mutation  $M = \{\mu_k\}_{k=1}^K$  is read out from its  $K$  corresponding single point mutation representations  $\{z(\mu_k)\}_{k=1}^K$ . Mutate Everything aggregates the representations for each single point mutation within a higher-order mutation. By aggregating the single mutations in the latent space, the model learns to account for the interactions between them. We aggregate with summation for its simplicity. A lightweight head  $g: \mathbb{R}^d \to \mathbb{R}$  then decodes the higher-order mutation  $\Delta \Delta G^{pr} = g\left(\sum_k z(\mu_k)\right)$ . This formulation allows us to seamlessly parameterize any higher-order mutation. For higher-order mutations (e.g.  $K > 1$ ), we empirically find that it is beneficial to predict the residual to the sum of single mutation  $\Delta \Delta G_s$  from our single mutation decoder (see Table 5d).
Decoding a million mutations in one forward pass. This decoding scheme efficiently computes  $\Delta \Delta G$  for millions of mutations, whether single or higher-order (see Figure 3). Our insight is to precompute all mutation features and reuse them for all predictions. First, we compute all  $L \times 20$  mutation representations  $Z = \{z((p,t)): p \in [1,\dots,L], t \in AA\}$ . This expands each token  $x_{l}$  to 20 representations that each capture the effect of mutating to one specific amino acid (including the wild-type amino acid). The backbone and this expansion are executed once for a protein sequence.
A lightweight head decodes  $\Delta \Delta G$  for any higher-order mutation (a mutation set). For each mutation set, we index and add the relevant mutation representations in  $Z$ . Index and sum operations are parallelized with outer-summation for  $K = 2$  and the "gather" operation for  $K > 2$ . Finally, the lightweight decoder  $g$  computes the  $\Delta \Delta G$  for each mutation set in parallel.
Training. Our model is fine-tuned using supervised learning. We optimize the Huber loss.
$$
\mathcal {L} \left(w, \left(M _ {i}, \Delta \Delta G _ {i} ^ {e x p}\right) _ {i = 1} ^ {N}\right) = \frac {1}{N} \sum_ {i = 1} ^ {N} \left| \left| \Delta \Delta G ^ {p r} \left(w, M _ {i}\right) - \Delta \Delta G _ {i} ^ {e x p} \right| \right| _ {H} \tag {1}
$$
where  $w$  is a protein sequence and  $(M_i, \Delta \Delta G_i^{exp})_{i=1}^N$  are the mutation sets for this protein with experimentally validated  $\Delta \Delta G$ . This approach efficiently backpropagates the loss through all labeled mutations for a given protein in a single pass, unlike prior work which backpropagates through only one mutation in a single pass. We train on single and double mutants with on average 2000 mutations for each protein. When training on double mutants, we learn the residual to the sum of the experimental (ground truth)  $\Delta \Delta G$  for the constituent single mutations. Note that we use the experimental  $\Delta \Delta G$  of the single mutants to train double mutants but during testing, we use the model's predicted  $\Delta \Delta G$  of the single mutants to infer double mutant  $\Delta \Delta G$ . Thus, we do not need experimental data for single effects to predict combined effects.
Only  $3\%$  of the double mutants in our training set are stabilizing  $(\Delta \Delta G < -0.5\mathrm{kcal / mol}$  [8]). To improve our performance on stabilizing mutations, we subsample the number of destabilizing double mutations to at most 8 times the number of stabilizing double mutations.
Method	Stabilizing rs	nDCG	DetPr	AUC	MCC
Mean	0.00	0.02	0.03	0.50	0.00
MSA	0.04	0.06	0.02	0.60	0.01
ESM2 [38]	0.04	0.08	0.02	0.65	0.00
PROSTATA [74] (Additive)	0.08	0.15	0.05	0.75	0.05
Mutate Everything (Additive)	0.08(0.03)	0.25(0.01)	0.10(0.01)	0.81(0.01)	0.25(0.02)
Mutate Everything (Ours)	0.14(0.02)	0.43 (0.02)	0.16 (0.01)	0.84 (0.01)	0.27 (0.01)
Table 1: Stabilizing Double Mutation Results on cDNA2. The metrics evaluate the model's performance on stabilizing mutations. Our additive baseline naively adds predicted single mutation  $\Delta \Delta G_{s}$ . Our model is state-of-the-art at finding stabilizing mutations. The proteins in cDNA2 val have at most  $36\%$  sequence similarity to those in the training set. Parenthesis indicates standard error across 7 training runs.
Our primary feature extractor is AlphaFold, implemented in OpenFold [3, 31]. The multiple sequence alignment (MSA) is computed using Colabfold [43]. Sequence features from the Evoformer and Structure Module are aggregated as input to the decoder. An adapter maps the backbone hidden dimension to  $D = 128$ . Our amino acid projections and single mutation decoders are implemented with a linear layer. The higher-order mutation decoder transforms the previous embedding with a 2-layer MLP. These representations are aggregated and fed into a 3-layer MLP to predict  $\Delta \Delta G$ .
We train on the cDNA display proteolysis dataset [72], which leverages a high throughput proteolysis screen enabled by next generation sequencing (NGS) to extract noisy  $\Delta \Delta G$  values for 100 mini-proteins totaling over 100,000 single and double mutations. To evaluate generalization on unseen proteins, we remove proteins from our training set that have high sequence similarity using MMSeqs2 [70] to any protein in our evaluation benchmarks. The short protein lengths ( $< 100$  amino acids) keep the memory requirement small for the higher-order decoder. See Section A.1 for more information about cDNA and other datasets.
We fine-tune a pre-trained backbone on single mutations for 20 epochs. Then, we finetune the model on both single and double mutations for 100 epochs using a cosine learning rate schedule with 10 warmup epochs. We use a batch size of 3 proteins due to the high memory requirements of AlphaFold. We use a learning rate of  $3 \times 10^{-4}$  and weight decay of 0.5. Training takes 6 hours on 3 A100 GPUs.
5.2 Finding Stabilizing Double Mutations
We evaluate our model's ability to find the stabilizing mutations on cDNA2. cDNA2 is our validation split of the cDNA double mutations dataset, consisting of 18 mini-proteins totaling 22,000 double mutations. The proteins in the validation set have at most  $36\%$  homology with those in the cDNA training set. Of these mutations, only 198 or  $0.8\%$  are stabilizing with  $\Delta \Delta G < -0.5 \, \text{kcal/mol}$  [8]. Like protein engineering, this scenario closely resembles the challenge of identifying a small number of stabilizing mutations amidst a much larger pool of destabilizing or neutral mutations.
We evaluate different methods on stabilizing mutations in Table 1. Stabilizing  $\mathbf{r}_s$  is the Spearman coefficient on the experimentally stabilizing subset introduced in [65]. Normalized discounted cumulative gain (nDCG) [30] measures the quality of the ranked mutation list by taking into account its experimentally validated  $\Delta \Delta G$  and position in the list [59]. Detection Precision (DetPr) is the proportion of experimentally validated stabilizing mutations among the top  $K = 30$  predictions. We adapt previous state-of-the-art single mutation predictors by naively adding both single mutation predictions. A detailed description of the metrics is found in Section A.2.
Our model demonstrates exceptional performance in prioritizing stabilizing double mutations over destabilizing ones, achieving a significantly higher normalized discounted cumulative gain of 0.43 compared to 0.25, as well as a superior detection precision of 0.16 compared to 0.10. Our model additionally improves classification metrics Matthews Correlation Coefficient (MCC) and Area under Precision-Recall Curve (AUC) by 0.02 and 0.03, respectively.