Improving protein optimization with smoothed fitness landscapes

We denote the starting set of $N$ proteins as $\mathcal{D}=(X,Y)$ where $X=\{x_{1},\dots,x_{N}\}\subset\mathcal{V}^{M}$ are the sequences and $Y=\{y_{1},\dots,y_{N}\}$ are corresponding real-valued scalar fitness measurements. Each sequence $x_{i}\in\mathcal{V}^{M}$ is composed of $M$ residues from a vocabulary $\mathcal{V}$ of 20 amino acids. Subscripts refer to different sequences. Note our method can be extended to other modalities, e.g. nucleic acids.

For in-silico evaluation, we denote the set of all known sequences and fitness measurements as $\mathcal{D}^{*}=(X^{*},Y^{*})$. We assume there exists a unknown black-box function $g:\mathcal{V}^{M}\rightarrow\mathbb{R}$ such that $g(x^{*})=y^{*}$. In practice, $g$ needs to be approximated by a evaluator model, $g_{\phi}$, trained with weights $\phi$ to minimize prediction error on $\mathcal{D}^{*}$. $g_{\phi}$ poses a limitation to evaluation since the true fitness needs to be verified with biological experiments. Nevertheless, an in-silico approximation provides a accessible way for evaluation and is done in all prior works. The starting dataset is a strict subset of the known dataset $\mathcal{D}\subset\mathcal{D}^{*}$ to simulate fitness optimization scenarios. Given $\mathcal{D}$, our task is to generate a set of sequences with higher fitness than the starting set.

Our goal is to develop a model of the sequence-to-fitness mapping that can be utilized when sampling higher fitness sequences. Unfortunately, the high-dimensional sequence space coupled with few data points and noisy labels can result in a noisy model that is prone to sampling false positives or getting stuck in local optima. To address this, we use smoothing techniques from graph signal processing.

The smoothing process is depicted in Figure 2. First, we train a noisy fitness model $f_{\tilde{\theta}}:\mathcal{V}^{M}\rightarrow\mathbb{R}$ with weights $\tilde{\theta}$ on the initial dataset $\mathcal{D}$ using Mean-Squared Error (MSE). $\mathcal{D}$ is usually very small in real-world scenarios. We augment the dataset by using $f_{\tilde{\theta}}$ to infer the fitness of neighboring sequences which we do not have labels for – known as transductive inference. Neighboring sequences are generated by randomly applying point mutations to each sequence in $X$. The augmented and original sequences become nodes, $V$, in our graph while their fitness labels are node attributes. Edges, $\mathcal{E}$, are constructed with a $k$-nearest neighbor (kNN) graph around each node based on the Levenshtein distance³³3Defined as the minimum number of mutations between two sequences.. The graph construction algorithm can be found in Algorithm 4.

The following borrows techniques from Isufi et al. (2022). The smoothness of the fitness variability in our protein graph is defined as the sum over the square of all local variability,

$\mathsf{TV}$ refers to Total Variation and $\Delta y_{i}$ is the local variability of node $i$ that measures local changes in fitness. Using $\mathsf{TV}_{2}$ as a regularizer, we solve the following optimization problem, known as Tikhunov regularization (Zhou & Schölkopf, 2004), for a new set of smoothed fitness labels,

With abuse of notation, we represent $Y$ as a vector with each node’s fitness. $\gamma$ is a hyperparameter set to control the smoothness; too high can lead to underfitting. We experiment with different $\gamma$’s in Section 4. Since eq. 2 is a quadratic convex problem, it has a closed form solution, $\hat{Y}=(\mathds{I}+\gamma L)^{-1}Y$ where $L$ is the graph Laplacian and $\mathds{I}$ is the identity matrix. The final step is to retrain the model on the sequences in the graph and their smoothed fitness labels. The result will be a model $f_{\theta}$ with lower $\mathsf{TV}_{2}$ than before and thus improved smoothness. The smoothing algorithm is in Algorithm 2.

Equipped with model $f_{\theta}$ from fig. 2, we apply it in a procedure to sample mutations that improve the starting sequences’ fitness. $f_{\theta}$ can also be viewed as an energy-based model (EBM) that defines a Boltzmann distribution $\log p(x)=f_{\theta}(x)-\log Z$ where $Z$ is the normalization constant. Higher fitness sequences will be more likely under this distribution, while sampling will induce diversity and novelty. To sample from $p(x)$, we use Gibbs With Gradients (GWG) Grathwohl et al. (2021) which has attracted significant interest due to its simplicity and state-of-the-art performance in discrete optimization. In this section, we describe the GWG procedure for protein sequences. GWG uses Gibbs sampling with approximations of locally informed proposals (Zanella, 2020):

With slight abuse of notation, we use the one-hot sequence representation $x\in\{0,1\}^{M\times|\mathcal{V}|}$ where $x_{i}\in\{0,1\}^{|\mathcal{V}|}$ represents the $i$th index of the sequence with 1 at its amino acid index and 0 elsewhere. $\odot$ is the element wise product. $H(x)=\{y\in\mathcal{V}^{M}:d_{\mathrm{Hamming}}(x,y)\leq 1\}$ is the 1-ball around $x$ using Hamming distance. The core idea of GWG is to use $d_{\theta}(x)_{i}$ as the first order approximation of a continuous gradient of the change in likelihood from mutating the $i$th index of $x$ to a different amino acid. The quality of the proposals in eq. 3 rely on the smoothness of the energy $f_{\theta}$ (Theorem 1 in Grathwohl et al. (2021)). If the gradients, $\nabla f_{\theta}$, are noisy, then the proposal distributions are ineffective in sampling better sequences. Hence, smoothing $f_{\theta}$ is desirable (see section 4).

The choice of $H(\cdot)$ as the 1-Hamming ball limits $x^{\prime}$ to point mutations from $x$ and only requires $\mathcal{O}\left(M\times|\mathcal{V}|\right)$ compute to construct. Let the point mutation where $x$ and $x^{\prime}$ differ be defined by the residue location, $i^{\mathrm{loc}}\in\{1,\dots,M\}$, and amino acid substitution, $j^{\mathrm{sub}}\in\{1,\dots,|\mathcal{V}|\}$. By limiting $x^{\prime}$ to point mutants $(i^{\mathrm{loc}},j^{\mathrm{sub}})$, sampling $q(x^{\prime}|x)$ is equivalent to sampling the following,

where $\tau$ is the sampling temperature and $d_{\theta}(x)_{i,j}$ is the logits of mutating to $(i,j)$. The proposal sequence $x^{\prime}$ is constructed by setting its $i^{\mathrm{loc}}$ residue to $j^{\mathrm{sub}}$ and equal to $x$ elsewhere. Each proposed sequence is accepted or rejected using Metropolis-Hasting (MH),

We provide the GWG algorithm in Algorithm 3.

We develop a set of tasks based on two well-studied protein systems: Green Fluoresent Protein (GFP) and Adeno-Associated Virus (AAV) (Sarkisyan et al., 2016; Bryant et al., 2021). These were chosen due to their relatively large amount of measurements, 56,806 and 44,156 respectively, with sequence variability of up to 15 mutations from the wild-type. Other datasets are either too small or do not have enough sequence variability. GFP’s fitness is its fluorescence properties as a biomarker while for AAV’s is the ability to package a DNA payload, i.e. for gene delivery. We found GFP and AAV to suffice in demonstrating how prior methods fail to extrapolate.

One measure of difficulty is the number of mutations required to achieve the highest known fitness; this assesses a method’s exploration capability. We designate the set of optimal proteins, $X^{\text{99th}}$, as any sequence in the 99th fitness percentile in the entire dataset⁴⁴4This may differ from the true optimal protein found in nature. Unfortunately, we must work with existing datasets since every possible protein cannot be experimentally measured.. Quantitatively, we compute the minimum number of mutations required from the training set to achieve the optimal fitness:

A high mutational gap would require the method discovering many mutations in a high dimensional space. A second measure of difficulty is the fitness range of the starting set of sequences. Starting with a low range of fitness requires the method to learn from barely functional proteins and exploit limited knowledge to find mutations that confer higher fitness. Appendix A shows Gap and starting rate are necessary as we found the previous GFP benchmark (Trabucco et al., 2022) as too “easy” by only requiring one mutation to achieve the optimal fitness.

GGS substantially outperforms all unsmoothed baselines, consistently achieving a improvement in fitness from the starting range of fitness in each difficulty. However, the smoothed baselines (lines with + GS) demonstrated a up to three fold improvement for CbAS, AdaLead. We find larger improvements in GFP where the sequence space is far larger than AAV – suggesting the GFP fitness landscape is harder to optimize over.

The most difficult task is clearly hard difficulty on GFP where all the baselines without smoothing cannot achieve fitness higher than the training set. With smoothing, GGS achieves the best fitness since the sampling procedure uses gradient-based proposals that benefit from a smooth model. Section C.2.1 presents results on additional difficulties to analyze GGS beyond hard..

We observe GGS is able to achieve the highest fitness while exhibiting respectable diversity and novelty. Notably, GGS’s novelty falls within the range of the mutational gap in each difficulty, suggesting it is extrapolating an appropriate amount for each task. Our sampling procedure, GWG, fails to perform without smoothing which agrees with its theoretical requirements of requiring a smooth model for good performance. We conclude smoothing is a beneficial technique not only for GGS but also for some baselines. GGS is able to achieve state-of-the-art results in our benchmark.

We analyze the effect of varying the following hyperparameters: number of nodes $N_{\text{nodes}}$ in the protein graph, smoothness weight $\gamma$ in eq. 2, and number of sampling rounds $R$ during GWG sampling. For space, we leave the analysis of the sampling temperature $\tau$ in section C.1. Figure 3 presents the results of running GGS with different hyperparameters on the hard difficulty of GFP and AAV. Along the X-axis, we plot the median performance of the sequences during each round of GWG where $r=0$ is initialization and $r=15$ are the sequences and the end of GWG. The Y-axis shows the predicted fitness of the smoothed model in blue while the fitness scored with our is shown in red. Interestingly, we find in the majority of cases the smoothed model’s predictions are highly correlated with the evaluator along the sampling trajectory. This is despite the model being trained on 4% of the data with the hard filtering. Section C.2.2 shows the prediction error where we find smoothing greatly improves in predicting the fitness of unseen sequences despite having higher train error.