SlowGuess/NeurIPS_2023_OCR huggingface dataset:
ProteinNPT: Improving Protein Property Prediction and Design with Non-Parametric Transformers
Computer Science,
Ruben Weitzman†
Debora S. Marks
Harvard Medical School
Protein design holds immense potential for optimizing naturally occurring proteins, with broad applications in drug discovery, material design, and sustainability. However, computational methods for protein engineering are confronted with significant challenges, such as an expansive design space, sparse functional regions, and a scarcity of available labels. These issues are further exacerbated in practice by the fact most real-life design scenarios necessitate the simultaneous optimization of multiple properties. In this work, we introduce ProteinNPT, a non-parametric transformer variant tailored to protein sequences and particularly suited to label-scarce and multi-task learning settings. We first focus on the supervised fitness prediction setting and develop several cross-validation schemes which support robust performance assessment. We subsequently reimplement prior top-performing baselines, introduce several extensions of these baselines by integrating diverse branches of the protein engineering literature, and demonstrate that ProteinNPT consistently outperforms all of them across a diverse set of protein property prediction tasks. Finally, we demonstrate the value of our approach for iterative protein design across extensive in silico Bayesian optimization and conditional sampling experiments.
Proteins, vital to most biological processes, have evolved over billions of years, shaping their structure through countless evolutionary experiments. The vastness of the protein sequence space offers immense opportunities for enhancing the properties of existing proteins, paving the way for advancements in healthcare, sustainability, and the development of novel materials [Arnold, 2018, Huang et al., 2016]. Recent progress in sequencing technologies and machine learning methods bear the promise to explore this space more efficiently and design functional proteins with the desired properties. A key step toward that goal is the ability to predict the fitness of novel protein sequences – i.e., learning a mapping between structure and fitness of proteins, commonly referred to as a 'fitness landscape' [Romero et al., 2013]. However, the protein space is sparsely functional and experimentally-collected labels about properties of interest are scarce relative to the size of protein space [Biswas et al., 2021a]. Additionally, real-life protein design scenario often seek to optimize multiple properties of the protein simultaneously, which further aggravates the label scarcity issues. Historically, learning fitness landscapes has been tackled by two different lines of research in the computational biology literature. Firstly, supervised methods are trained to approximate the
*Correspondence: pascal.notin@cs.ox.ac.uk, ruben.weitzman@cs.ox.ac.uk; † Equal contribution
fitness landscape by learning from experimental labels measuring a phenotype that is hypothesized to correlate with the function of interest. These methods have for instance been shown to lead to more efficient directed evolution pipelines [Dougherty and Arnold, 2009, Yang et al., 2019, Wittmann et al., 2020]. Given the scarcity of labels, they are typically constrained in terms of parameter count and expressivity of the underlying model (e.g., ridge regression, shallow neural network) [Heinzinger et al., 2019, Dallago et al., 2021, 2022, Stark et al., 2021] to avoid overfitting. Secondly, unsupervised fitness predictors have reached increasingly higher fitness prediction performance without being subject to the limitations and potential biases stemming from training on sparse labels [Hopf et al., 2017a, Riesselman et al., 2018, Laine et al., 2019, Frazer et al., 2021, Meier et al., 2022, Notin et al., 2022a,b, Marquet et al., 2022]. Recently, Hsu et al. [2022] proposed to combine these two lines of research into a unified architecture in which a ridge regression is fed both one-hot encodings of input protein sequences and unsupervised fitness scores, and then trained on experimental labels to reach increased performance across several fitness prediction benchmarks. Yet, one-hot encoding features offer a relatively contrived and brittle sequence representation. For instance, they do not allow to make predictions at unseen positions resulting in low performance when attempting to extrapolate across positions ( $\S$  4.3). Embeddings from protein language models offer a promising alternative but, since the resulting dimension of the full sequence embedding is typically too high, prior methods have resorted to dimensionality reduction techniques such as mean-pooling embeddings across the full sequence [Alley et al., 2019, Biswas et al., 2021a]. While this works reasonably well in practice, critical information may be lost in the pooling operation and, since not all residues may be relevant to a given task, we may want to be selective about which ones to consider. In this work, we introduce ProteinNPT ( $\S$  3), a non-parametric transformer [Kossen et al., 2022] variant which is ideally suited to label-scarce settings through an additional regularizing denoising objective, straightforwardly extends to multi-task optimization settings and addresses all aforementioned issues. In order to quantify the ability of different models to extrapolate to unseen sequence positions, we devise several cross-validation schemes ( $\S$  4.1) which we apply to all Deep Mutational Scanning (DMS) assays in the ProteinGym benchmarks [Notin et al., 2022a]. We then show that ProteinNPT achieves state-of-the-art performance across a wide range of protein property prediction and iterative redesign tasks. We summarize our contributions are as follows:
- We introduce ProteinNPT, a semi-supervised conditional pseudo-generative model for protein property prediction and design tasks (§ 3);
- We explore novel aspects of non-parametric transformers such as the use of auxiliary labels ( $\S 3.2$ ), the multiple property prediction setting ( $\S 4.5$ ), conditional sampling ( $\S 3.3$ ) and uncertainty quantification (Appendix G.1);
- We devise several cross-validation schemes to assess the performance of fitness predictors in the supervised setting and their ability to extrapolate across positions (§ 4.1);
- We implement several of the current top-performing baselines, illustrate certain of their limitations, then improve these baselines accordingly by integrating ideas from diverse branches of the protein fitness prediction and engineering literature (§ 4.2);
- We subsequently show that ProteinNPT outperforms all of these baselines across a wide spectrum of property prediction tasks (§ 4.3-4.5);
- We demonstrate the potential of ProteinNPT for protein design through a diverse set of in silico Bayesian optimization (§ 5.1) and conditional sampling (§ 5.2) experiments.
Protein fitness prediction Learning a fitness landscape is typically cast as a discriminative supervised learning task [Yang et al., 2018, Gelman et al., 2020, Freschlin et al., 2022]. Given the usually scarce number of labels relative to the dimensionality of the input space, the majority of approaches to date have relied on lightweight parametric models [Stärk et al., 2021, Heinzinger et al., 2019] of the primary protein structure and, occasionally, non-parametric models based on the tertiary protein structure [Romero et al., 2013]. Model features have revolved primarily around hand-crafted physico-chemical properties [Wittmann et al., 2020], one-hot-encodings of the amino acid sequence [Dallago et al., 2022] and, more recently, embeddings from deep generative models of evolutionary sequences [Alley et al., 2019, Brandes et al., 2021, Hesslow et al., 2022, Elnaggar et al., 2023a]. Deep generative models have also been successfully leveraged for zero-shot fitness
Figure 1: ProteinNPT architecture. (Left) The model takes as input the primary structure of a batch of proteins of length  $L_{seq}$  along with the corresponding  $L_{t}$  labels and, optionally,  $L_{a}$  auxiliary labels (for simplicity we consider  $L_{t} = L_{a} = 1$  here). Each input is embedded separately, then all resulting embeddings are concatenated into a single tensor. Several ProteinNPT layers are subsequently applied to learn a representation of the entire batch, which is ultimately used to predict both masked tokens and targets (depicted by question marks). (Right) A ProteinNPT layer alternates between tied row and column attention to learn rich embeddings of the labeled batch.
prediction [Hopf et al., 2017a, Riesselman et al., 2018, Meier et al., 2022, Notin et al., 2022a]. Their use is particularly compelling when experimental labels are not available or difficult to obtain for the property of interest. For instance, they enable accurate prediction of the effects of genetic mutations on human disease risks [Frazer et al., 2021, Brandes et al., 2022, Jagota et al., 2023], or can identify viral mutants with high escape potential at the beginning of a pandemic [Hie et al., 2020, Thadani et al., 2023].
Protein redesign The goal of protein redesign is to iteratively mutate a starting protein, typically a naturally occurring sequence, in order to enhance one or several of its properties, such as its stability, substrate specificity, or catalytic efficiency. Directed evolution [Arnold, 2018] is a widely-used experimental method in this field, emulating the way natural evolution works, but within a controlled lab environment. In each cycle, the starting protein sequences undergo random mutations, and the desired traits are then measured experimentally. The most promising variants are selected as the basis for the next round, and the process is repeated through several cycles until the desired design goals are achieved. Machine learning [Yang et al., 2019, Wu et al., 2019] has been shown to significantly enhance this process by exploring the combinatorial space of mutations more efficiently, for instance by reaching the same final performance with fewer experimental rounds [Biswas et al., 2021b].
Self-attention over multi-dimensional objects Popularized by the transformer architecture [Vaswani et al., 2017], self-attention enables to learn rich representations of complex structured objects, such as natural language sequences. One of its key limitations is the quadratic complexity of the attention mechanism with respect to the length of the input. Axial attention [Ho et al., 2019]
addresses that issue for multi-dimensional objects such as images by decomposing the traditional attention mechanism along each axis: separate self-attention mechanisms are successively applied on each axis, thereby drastically reducing computational complexity. The MSA Transformer [Rao et al., 2021] leverages axial attention to learn representations of a large and diverse set of Multiple Sequence Alignments (MSA) – two-dimensional objects comprised of homologous protein sequences aligned in the same coordinate system. As MSAs encapsulate rich evolutionary and structural information about the corresponding proteins, the MSA Transformer has lead to strong performance for several tasks, such as zero-shot contact prediction, fitness prediction [Meier et al., 2022] and protein engineering [Sgarbossa et al., 2023]. Viewing tabular data as two-dimensional objects, Non-Parametric Transformers [Kossen et al., 2022] applies axial attention to learn a representation of an entire labeled training dataset in a supervised setting. Self-attention is thus used to model relationships both between features and labels (across columns), and between labeled objects (across rows).
3.1 ProteinNPT model
Model architecture ProteinNPT (Fig. 1) is a semi-supervised conditional pseudo-generative model that learns a joint representation of protein sequences and associated property labels. The model takes as input both the primary structure representation of the proteins along with the corresponding labels for the property of interest. Let  $(X^{\mathrm{full}}, Y^{\mathrm{full}})$  be the full training dataset where  $X^{\mathrm{full}} \in \{1, 20\}^{N.L_s}$  are protein sequences (with N the total number of labeled protein sequences and L the sequence length), and  $Y^{\mathrm{full}} \in \mathbb{R}^{N.T}$  the corresponding property labels (where T is the number of distinct such labels, including  $L_t$  true targets and  $L_a$  auxiliary labels, as discussed in § 3.2). Depending on whether we are at training or inference time, we sample a batch of B points and mask different parts of this combined input as per the procedure described later in this section. We separately embed protein sequences and labels, concatenate the resulting sequence and label embeddings (each of dimension d) into a single tensor  $Z \in \mathbb{R}^{(B.(L_s + T).d)}$ , which we then feed into several ProteinNPT layers. A ProteinNPT layer (Fig. 1 - right) learns joint representation of protein sequences and labels by applying successively self-attention between residues and labels for a given sequence (row-attention), self-attention across sequences in the input batch at a given position (column-attention), and a feedforward layer. Each of these transforms is preceded by a LayerNorm operator  $LN(.)$  and we add residual connections to the output of each step. For the multi-head row-attention sub-layer, we linearly project embeddings for each labeled sequence  $n \in \{1, B\}$  for each attention head  $i \in \{1, H\}$  via the linear embeddings  $Wr_i^K, Wr_i^Q$  and  $Wr_i^V$  respectively. Mathematically, we have:
$$
\operatorname {R o w} - \operatorname {A t t} (Z) = Z + \operatorname {T i e d} - \operatorname {R o w} - \operatorname {M H S A} (L N (Z)) = Z + \operatorname {c o n c a t} \left(O _ {1}, O _ {2}, \dots , O _ {H}\right). W ^ {O} \tag {1}
$$
where the concatenation is performed row-wise,  $W_{O}$  mixes outputs  $O_{i}$  from different heads, and we use tied row-attention as defined in Rao et al. [2021] as the attention maps ought to be similar across labeled instances from the same protein family:
$$
\text {T i e d - R o w - A t t} \left(\left(Q _ {n}, K _ {n}, V _ {n}\right) _ {n \in \{1, B \}}\right) = \operatorname {s o f t m a x} \left(\sum_ {n = 1} ^ {B} \frac {Q _ {n} . K _ {n} ^ {T}}{\sqrt {B . d}}\right). V _ {n} \tag {2}
$$
We then apply column-attention as follows:
$$
\operatorname {C o l} - \operatorname {A t t} (Z) = Z + \operatorname {M H S A} (L N (Z)) = Z + \operatorname {c o n c a t} \left(P _ {1}, P _ {2}, \dots , P _ {H}\right). W ^ {P} \tag {3}
$$
where the concatenation is performed column-wise,  $W_{P}$  mixes outputs  $P_{i}$  from different heads, and the standard self-attention operator  $\mathrm{Att}(Q,K,V) = \mathrm{softmax}(Q.K^{T} / \sqrt{d}).V$ . Lastly, the feedforward sub-layer applies a row-wise feed-forward network:
$$
\operatorname {F F} (Z) = Z + \operatorname {r F F} (L N (Z)) \tag {4}
$$
In the final stage, the learned embeddings from the last layer are used to predict both the masked tokens and targets: the embeddings of masked targets are input into a L2-penalized linear projection to predict masked target values, and the embeddings of masked tokens are linearly projected then input into a softmax activation to predict the corresponding original tokens.
Model training Our overall architecture is trained with a two-stage semi-supervised procedure. First, due to the scarcity of labeled instances relative to the size of the underlying protein sequence space, we leverage embeddings from protein language models that have been pre-trained on large quantities of unlabeled natural sequences. Throughout training, the parameters of the model used to obtain these protein embeddings are frozen to prevent overfitting. This also helps alleviating several practical challenges pertaining to GPU memory bottlenecks at train time (Appendix C.3). Second, similarly to what was proposed in Non-Parametric Transformers [Kossen et al., 2022], we train our architecture on a composite task combining input denoising and target prediction. In our case, the former is equivalent to a masked language modeling objective on amino acid tokens [Devlin et al., 2019]. We randomly mask a fraction of amino acids and of input labels across all sequences in the mini batch (15% for both), and seek to predict these masked items. Our training loss is thus comprised of two components: 1) a reconstruction loss over masked amino acids 2) a target prediction loss over masked targets:
$$
\mathcal {L} ^ {\text {t o t a l}} = \alpha_ {t}. \mathcal {L} ^ {\mathrm {A A} \text {r e c o n s t r u c t i o n}} + (1 - \alpha_ {t}). \mathcal {L} ^ {\text {t a r g e t} \text {p r e d i c t i o n}} \tag {5}
$$
where  $\mathcal{L}^{\mathrm{target~prediction}}$  is target prediction loss (e.g., Mean Squared Error between predictions for the masked labels and corresponding true labels for continuous targets),  $\mathcal{L}^{\mathrm{AA~reconstruction}}$  the cross-entropy loss for the masked input tokens, and  $\alpha_{t}$  balances out the two objectives. As suggested by Kossen et al. [2022],  $\alpha_{t}$  is progressively annealed to increasingly focus on target prediction during training. The amino acid denoising objective entics the network to learn a representation of the full dataset, which in turn acts as a regularization mechanism which is beneficial in label scarce settings.
Inference At inference, we form input batches by concatenating row-wise a fraction of test sequences (with unknown, masked labels) with training instances (with known, unmasked labels). Predictions are therefore dependent on both the learned relationships between tokens and targets, as well as homology across labeled instances. When the size of the training set is large enough, we cannot use the entire training set at inference due to memory bottlenecks, and have to resort to sampling training instances. We explore the impact of the number of training sequences sampled at inference in § 6.
Advantages of tri-axial attention To obtain the initial protein sequence embeddings, we experimented with several best-in-class protein language models and obtained better results with the MSA Transformer (Appendix C.2). As a result, our overall architecture operates tri-axial self-attention to output predictions: across residues and labels, across homologous sequences from the retrieved MSA (in MSA Transformer layers only) and across labeled examples (in ProteinNPT layers only). There are several advantages conferred by this architecture. First, it seamlessly adapts to multiple target predictions, by concatenating as many label columns in its input as required. Second, it naturally captures correlations between these targets and automatically handles – at training and at inference – observations with partially available labels. Third, to avoid overfitting to sparse label sets, prior semi-supervised approaches typically apply a pooling operator (e.g., average pooling across the full sequence length) to reduce the dimensionality of input features. This pooling operation potentially destroys valuable information for the downstream task. In contrast, no such pooling operation is applied in the ProteinNPT architecture, and the network leverages self-attention to learn dependencies between labels and the embeddings of specific residues in the sequence.
3.2 Auxiliary labels
The effectiveness of our architecture is significantly enhanced by the incorporation of auxiliary labels at training and inference. We define auxiliary labels as additional inputs which are both 1) easy to obtain for all or for a subset of train and test instances 2) known to be correlated with the target of interest. They are concatenated column-wise to the input batch and are handled by ProteinNPT layers in the forward pass just like any other target. However, they are excluded from the loss computation and are only meant to increase performance on the target prediction task by instilling complementary information into the model. In the fitness prediction tasks discussed in § 4, we use as auxiliary labels the unsupervised fitness predictions obtained with the underlying protein language model used to extract the input sequence embeddings (i.e., MSA Transformer). As shown in ablations (Appendix C.2), and consistent with the findings from Hsu et al. [2022], this helps significantly increase our prediction performance, in particular when performing inference at positions where no mutation was observed during training.
Diverging from prior works focusing on supervised protein property predictions (§ 2 and § 4.2), ProteinNPT learns a joint representation of both protein sequences and labels. Since it is trained with a masked-language model objective, its output at a given masked position is a distribution over the amino acid vocabulary at that position conditioned on non-masked tokens, including labels. As such, it is a conditional pseudo-genative model which can be used to sample new sequences iteratively. Prior works have investigated sampling with protein masked-language models in the unsupervised setting [Sgarbossa et al., 2023]. However, a distinctive feature of our architecture is its ability to condition on values of the property of interest. This approach allows to steer the generation of new sequences to align with specific protein design goals, as detailed in § 5.2.
4 Protein property prediction
4.1 ProteinGym benchmark
The fitness landscape of naturally occurring proteins is the result of an intricate set of overlapping constraints that these proteins are subjected to in an organism. It is thus challenging to identify a single molecular property that is both easy to measure experimentally and that reflects that complexity. Experimental measurements are imperfect ground truth measurements of the underlying fitness. Furthermore, there is no such thing as fitness in the absolute but rather fitness at a specific temperature, pH, or in a given cell type. To make the conclusions that we seek to draw in the subsequent sections more meaningful, it is critical to consider a large and diverse set of experimental assays. To that end, we conducted our experiments on ProteinGym [Notin et al., 2022a], which contains an extensive set of Deep Mutational Scanning (DMS) assays covering a wide range of functional properties (e.g., thermostability, ligand binding, viral replication, drug resistance). We include an additional 13 DMS assays to ProteinGym, bringing the total to 100 DMS assays (Appendix A).
We develop 3 distinct cross-validation schemes to assess the ability of each model to extrapolate to positions not encountered during training. In the Random scheme, commonly-used in other supervised fitness prediction benchmarks [Rao et al., 2019, Dallago et al., 2022], each mutation is randomly allocated to one of five distinct folds. In the Contiguous scheme, the sequence is split into five contiguous segments along its length, with mutations assigned to each segment based on the position they occur in the sequence. Lastly, the Modulo scheme uses the modulo operator to assign mutated positions to each fold. For example, position 1 is assigned to fold 1, position 2 to fold 2, and so on, looping back to fold 1 at position 6. This pattern continues throughout the sequence. We note that there is no inherent issue with using a Random cross-validation scheme to estimate the performance of predictive models. However, the conclusions drawn and the generalizability claims based on it require careful consideration.