Title: ProteinNPT: Improving Protein Property Prediction and Design with Non-Parametric Transformers
Venue: NeurIPS 2023
Accepted: yes
Keywords: Non-Parametric Transformers

protein design

protein property prediction

fitness prediction

Bayesian optimization

ProteinGym
Abstract: 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.

ProteinNPT: Improving Protein Property Prediction and Design with Non-Parametric Transformers

Abstract

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.

1 Introduction

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

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, Stärk 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 (§ 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 (§ 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 (§ 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 (§ 3.2), the multiple property prediction setting (§ 4.5), conditional sampling (§ 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.

2 Related work

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

Target

Masked-tokenspredictions

predictions

Linear projection of last layerembeddings for

masked tokens,then softmax activation

L2-penalized linear

projection

ProteinNPT Layer

Last

ndano

Layer

B:

B:

tensor

Embed.

(Lseg +2).d

Split

ProteinNPTLayers

Feedforward

There are 5 consecutive PNPT layers. Each PNPT layer performs axial attention:

Row-wise: across amino acid tokens, targets and auxiliary labels

LayerNorm

Column-wise: across labeled sequences in the mini batch

KP

Column Attention

Input

Embed.

B:

Attention labeled sequences in the sampled mini batch

LayerNorm

(L+2).d

Concatenate

Tied RowAttention

Sequenceembedder

Target

Aux.label

Attention across amino acid tokens, targets and aux.labels

MSA Transformer performs axial attention:

embedder

embedder

Average atttention maps across sequences in mini-batch

Row-wise: across amino acid tokens

Learned linear

Learned linear

Column-wise: across sequences from a MSA

embedding

embedding

LayerNorm

Input

B:

？

B :

Input

batch

？

tensor

Lsea

4→

4

4→

L。

(Lseg +2).d

Protein sequences

Targets

Aux.labels

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 Methods

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 full , Y full ) be the full training dataset where X full ∈ { 1 , 20 } N.L s are protein sequences (with N the total number of labeled protein sequences and L the sequence length), and Y full ∈ 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 ∈ 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 ∈ { 1 , B } for each attention head i ∈ { 1 , H } via the linear embeddings Wr K i , Wr Q i and Wr V i respectively. Mathematically, we have:

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:

We then apply column-attention as follows:

where the concatenation is performed column-wise, W P mixes outputs P i from different heads, and the standard self-attention operator Att ( Q,K,V ) = softmax ( Q.K T / √ d ) .V . Lastly, the feedforward sub-layer applies a row-wise feed-forward network:

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:

where L target prediction is target prediction loss (e.g., Mean Squared Error between predictions for the masked labels and corresponding true labels for continuous targets), L AA reconstruction the cross-entropy loss for the masked input tokens, and α t balances out the two objectives. As suggested by Kossen et al. [2022], α t is progressively annealed to increasingly focus on target prediction during training. The amino acid denoising objective entices 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.

3.3 Conditional sampling

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-generative 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.

4.2 Baselines

We start from the top-performing model variant from [Hsu et al., 2022], namely a ridge regression on one-hot encodings of protein sequences, augmented with unsupervised fitness predictions from DeepSequence [Riesselman et al., 2018]. Building on recent progress from the protein property prediction and engineering literature, we then improve this baseline in several ways. First, we show that augmenting the model from Hsu et al. [2022] with the latest state-of-the-art unsupervised fitness predictors helps further boost the predictive performance, with the performance boost being consistent with the zero-shot performance of the corresponding unsupervised models (Appendix D). Then, considering the performance across the various cross-validation schemes described in § 4.1, we show the limits of One-Hot-Encoded (OHE) features to generalize across positions and propose to use instead embeddings from large protein language models. There is a rich literature investigating this idea, and we test out a wide range of model architectures that have been proposed to learn a mapping between embeddings and labels: L2-penalized regression, shallow Multi-Layer Perceptron (MLP), Convolutional Neural Network (CNN), ConvBERT [Elnaggar et al., 2023b] and Light Attention [Stärk et al., 2021]. Combining these different ideas together, we obtain the main baselines we compare against in subsequent sections: models that both leverage embeddings from pretrained language models, and are augmented with unsupervised fitness predictions. We carry out thorough ablations and hyperparameter search (Appendix C.2) to identify the best performing combinations.

4.3 Single mutant property prediction

Table 1: Fitness prediction performance . Spearman's rank correlation & MSE between model predictions & experimental measurements, averaged across all assays in ProteinGym. DS and MSAT are shorthand for DeepSequence and MSA Transformer respectively. "Aug." indicates models augmented with unsupervised predictions. The Spearman for MSAT in the zero-shot setting is 0.41 (independent of CV split). Assay-level performance & standard errors are reported in Appendix D.

Experimental setup We first focus on the task of predicting fitness for single substitution mutants. For each assay in ProteinGym and cross-validation scheme introduced in § 4.1, we perform a 5-fold cross-validation, selecting the first 4 folds for training, and using the remaining one as test set. Intuitively, since the Random cross-validation scheme is such that mutations in the test set may occur at the same positions as certain mutations observed in the training data, model performance will typically be higher in that regime. Similarly, the Contiguous scheme is intuitively the most difficult regime in which we assess the ability of the different models to extrapolate on positions relatively 'far' in the linear sequence from any mutation being part of the training data.

Results We find that ProteinNPT markedly outperforms all baselines, across all cross-validation schemes, and both in terms of Spearman's rank correlation and Mean Squared Error (MSE) (Table 1). Our results also confirm that a model that would learn solely based on one-hot-encoded inputs would be unable to make any valuable prediction at positions that were not observed at training time (Spearman's rank correlation near zero on the Contiguous and Modulo schemes). Augmenting these models with predictions from unsupervised predictors helps to generally address that issue, although the performance on the Contiguous and Modulo schemes is similar to the zero-shot performance of unsupervised models (the average Spearman's rank correlation across all assays for the MSA Transformer in the zero-shot setting is 0.41). The baseline combining a transform of embeddings from the MSA Transformer, augmented with unsupervised fitness predictions ('Embeddings - Augmented (MSAT) in Table 1) performs better across the board than the one-hot encoding equivalent, yet its performance is far from that of ProteinNPT, in particular in the Random cross-validation scheme. We interpret the latter by the fact the column-wise attention model in the ProteinNPT layer (Fig. 1) is more beneficial to performance when a mutation at that position has been observed at train time.

4.4 Multiple mutants property prediction

Experimental setup We next consider all assays in the extended ProteinGym benchmark that include multiple mutants. Given the challenges to assign multiple mutants to non-overlapping position sets, we only consider the Random cross-validation scheme in this analysis. We report the Spearman's rank correlation in Fig. 2 and MSE in Appendix E (conclusions are consistent for the two metrics).

Results The performance lift of ProteinNPT observed in the single mutants analyses translates to multiple mutants: ProteinNPT has superior aggregate performance and outperforms all other baselines in 14 out of the 17 assays with multiple mutants (Fig. 2, left). The performance lift is also robust to the mutational depth considered (Fig. 2, right).

4.5 Multiple property prediction

ProteinGym contains a subset of protein families for which several experimental assays have been carried out, offering the possibility to test out the ability of various models to predict multiple properties of interest simultaneously. We report the performance on the Random cross-validation scheme in Fig. 9, and other schemes are provided in Appendix E. ProteinNPT outperforms all other baselines in that regime as well, consistently across all settings analyzed.

Models

Models

Zero-shot (MSAT)

Embeddings - Aug. (MSAT)

-Zero-shot(MSAT)

-- Embeddings -Aug. (MSAT)

OHE

ProteinNPT

-OHE

--ProteinNPT

OHE - Aug. (MSAT)

○-OHE-Aug.(MSAT)

CAPSD_AAV2S_Sinai_substitutions_2021

C

0.8

Q6WV13_9MAX1_Somermeyer_2022

Q8WTC7_9CNID_Somermeyer_2022

0.7

D7PM05_CLYGR_Somermeyer_2022

A4_HUMAN_Seuma_2021

GRB2_HUMAN_Faure_2021 -

0.6

DLG4_HUMAN_Faure_2021

GFP_AEQVI_Sarkisyan_2016

YAP1_HUMAN_Araya_2012

0.5

HIS7_YEAST_Pokusaeva_2019

CAPSD_BPMS2_Brauer_2022

F7YBW8_MESOW_Aakre_2015

0.4

BAC7_Held_2022 -

GCN4_YEAST_Staller_induction_2018

SPG1_STRSG_Wu_2016 -

Average

0.0

0.2

0.4

0.6

0.8

1.0

2

3

5+

Spearman

Mutational Depth

5 Protein redesign

5.1 In silico iterative protein redesign

Experimental setup Our goal is to emulate in silico a real-world protein engineering scenario where we iteratively redesign a known protein sequence across several experimental rounds to increase its fitness. Performing multiple rounds of experiments sequentially allows us to harness the information acquired in prior rounds, as is commonly done in directed evolution. In a standard wet lab setting, there is no constraint on the subset of mutated sequences that can be experimentally tested. However, in our in silico setup, we are constrained to select mutants that belong to a predefined pool of sequences which have been experimentally measured in the DMS assays we consider, such as the ones from ProteinGym. We provide a detailed algorithm for our procedure in Fig. 3 (left), which we cast as a batch Bayesian optimization task. At the beginning of each in silico experiment, this pool is entirely unlabeled. In each subsequent round, we score all possible variants in the unlabeled pool and select from it a subset B of items which are deemed optimal by the Upper Confidence Bound (UCB) acquisition function:

where µ is the predicted fitness value, σ quantifies uncertainty in the prediction, and λ controls the exploration-exploitation trade-off. We develop and test various uncertainty quantification methods for non-parametric transformer architectures, based on Monte Carlo dropout [Gal and Ghahramani, 2015] and resampling inference batches with replacement (see details in Appendix G.1). We find that a hybrid scheme in which we both sample model parameters via Monte Carlo dropout and resample inference batches works best. After each acquisition cycle, all models are retrained from scratch for 2k training steps. On Fig. 3, we plot the proportion of acquired points that are above a predefined threshold (top 3 deciles of all measurements in the DMS assays) as a function of the number of acquisition rounds performed, averaged across all assays. The baselines considered are the same as for the property prediction experiments. Individual plots for each assay are provided in Appendix G.3.

Results We find that ProteinNPT vastly outperforms all other baselines in aggregate across all assays, and achieves superior performance on the large majority of individual assays. Furthermore, we observe that the performance lift is more pronounced for assays for which we start with a larger pool of unlabeled sequences. This suggests that in a real-world redesign scenario, in which choices are not limited to a predefined pool as in our in silico setting, ProteinNPT might exhibit even greater performance gains relative to the baseline models.

0.8

Algorithm 1 Iterative protein redesign

ProteinNPT

Embeddings - Aug. (MSAT)

0.7

OHE - Aug. (MSAT)

points

OHE - Aug. (DeepSequence)

0.6

OHE

fitness

Zero-shot (MSAT)

0.5

Random

Recall rate of highf

0.4

0.3

0.2

0.1

0

2

4

6

8

10

Acquisition cycle

5.2 Conditional sampling

As a supervised pseudo-generative model of labeled protein sequences, ProteinNPT has the ability to generate novel sequences conditioned on particular values of these properties. This is in contrast to prior generative models that have been used to sample natural-like proteins with no bias towards properties of interest [Repecka et al., 2021, Sgarbossa et al., 2023], or conditioned on broad taxonomic labels characterizing full protein families [Madani et al., 2023] as opposed to the sequence-specific properties we consider here. In this experiment, we illustrate the ability of ProteinNPT to sample novel protein sequences with desired property from the Green Fluorescent Protein (GFP), but the approach can be applied to any other protein. We first identify the sequence with the highest measured property in the GFP DMS assay [Sarkisyan et al., 2016]. We then form an input batch by randomly selecting other labeled sequences, mask a fixed number of tokens (5 in our experiment) in the fittest sequence, obtain the resulting log softmax over the masked positions with a forward pass and sample from these to obtain new sequences. Critically, rather than selecting the masked positions at random, we sample them from the positions with the highest row-wise attention coefficient with the target (averaged across heads) in the last layer of ProteinNPT. This helps ensure we sample new mutations at the positions with the highest impact on the target of interest. We generate 1k new sequences with that process and measure the corresponding fitness with an independent zero-shot fitness predictor (ESM-1v). Finally, we compare the resulting fitness distribution across all samples with distributions obtained with two baselines: the first in which mutations are randomly selected, and the second in which mutations are selected to further minimize the fitness of the less fit protein. We report the corresponding results in Figure 4, and observe that proteins sampled with ProteinNPT have substantially higher fitness values relative to the two baseline sampling strategies and relative to sequences in the original DMS assay.

6 Discussion

Row-wise attention captures dependencies between residues and function. Since we learn a joint representation of sequence and labels, the attention maps of ProteinNPT capture meaningful correlations between residues and the property labels we train on. We represent the row-wise attention coefficients for a ProteinNPT model trained on a DHFR protein assay [Thompson et al., 2020a] in Fig. 5 (last layer, first attention head). The analysis recapitulates certain known facts about important residues for the assayed property (e.g., the highest attention score in red corresponds to a core substrate binding site), and may help identify new residues involved the protein function of interest.

0.35

Highfitness sampling

Random sampling

Lowfitness sampling

Best of labeled set

FOE'O

.....Worst of labeled set

: 1

: 1

0.25

: 1

：1

：1

0.15-

0.10-

0.05-

0.00

12.5

-10.0

7.5

-5.0

2.5

0.0

2.5

5.0

7.5

Fitness Value(ESM-1v zero-shot fitness score)

Column-wise attention is critical to reach top performance. We test out the impact of columnattention on model performance. All analyses are conducted on the same subset of 8 DMS assays used in other ablations (Appendix B.2). We first re-train ProteinNPT models without the column-wise attention layers, and observe that the downstream performance drops significantly (Table 2, left). At inference, we find that predicting with 100 randomly-selected labeled sequences captures most of the performance lift, and observe no benefit beyond 1k batch points (Table 2, right).

7 Conclusion

This work introduced ProteinNPT, a novel conditional pseudo-generative architecture which reaches state-of-the-art performance across a wide range of fitness prediction and iterative redesign settings, covering a diverse set of protein families and assayed properties (e.g, thermostability, binding, enzymatic activity). The architecture straightforwardly extends to the multiple property prediction setting, can leverage auxiliary labels, supports uncertainty quantification, and allows to sample new sequences conditioned on desired property values. While sequence embeddings from the MSA Transformer lead to the best performance in our ablations, the architecture of ProteinNPT is agnostic to the model used to obtain these embeddings. As protein language models continue to improve, ProteinNPT will also benefit from the resulting improved embeddings. Future work will investigate additional strategies for conditional sampling, extension of the architecture to handle insertions and deletions, and generalization across protein families.

References

Frances H. Arnold. Directed Evolution: Bringing New Chemistry to Life. Angewandte Chemie International Edition , 57(16):4143-4148, 2018. ISSN 1521-3773. doi: 10.1002/anie. 201708408. URL https://onlinelibrary.wiley.com/doi/abs/10.1002/anie.201708408 . _eprint: https://onlinelibrary.wiley.com/doi/pdf/10.1002/anie.201708408.

Po-Ssu Huang, Scott E. Boyken, and David Baker. The coming of age of de novo protein design. Nature , 537(7620):320-327, September 2016. ISSN 1476-4687. doi: 10.1038/nature19946. URL https://www. nature.com/articles/nature19946 . Number: 7620 Publisher: Nature Publishing Group.

Philip A. Romero, Andreas Krause, and Frances H. Arnold. Navigating the protein fitness landscape with Gaussian processes. Proceedings of the National Academy of Sciences , 110(3):E193-E201, January 2013. doi: 10.1073/pnas.1215251110. URL https://www.pnas.org/doi/10.1073/pnas.1215251110 . Publisher: Proceedings of the National Academy of Sciences.

Surojit Biswas, Grigory Khimulya, Ethan C. Alley, Kevin M. Esvelt, and George M. Church. Low-N protein engineering with data-efficient deep learning. Nature Methods , 18(4):389-396, April 2021a. ISSN 1548-7091, 1548-7105. doi: 10.1038/s41592-021-01100-y. URL http://www.nature.com/articles/ s41592-021-01100-y .

Michael J Dougherty and Frances H Arnold. Directed evolution: new parts and optimized function. Current Opinion in Biotechnology , 20(4):486-491, August 2009. ISSN 0958-1669. doi: 10.1016/j.copbio.2009.08.005. URL https://www.sciencedirect.com/science/article/pii/S0958166909000986 .

Kevin K. Yang, Zachary Wu, and Frances H. Arnold. Machine-learning-guided directed evolution for protein engineering. Nature Methods , 16(8):687-694, August 2019. ISSN 1548-7105. doi: 10.1038/s41592-019-0496-6. URL https://www.nature.com/articles/s41592-019-0496-6 . Number: 8 Publisher: Nature Publishing Group.

Bruce J. Wittmann, Yisong Yue, and Frances H. Arnold. Machine Learning-Assisted Directed Evolution Navigates a Combinatorial Epistatic Fitness Landscape with Minimal Screening Burden, December 2020. URL https://www.biorxiv.org/content/10.1101/2020.12.04.408955v1 . Pages: 2020.12.04.408955 Section: New Results.

Michael Heinzinger, Ahmed Elnaggar, Yu Wang, Christian Dallago, Dmitrii Nechaev, Florian Matthes, and Burkhard Rost. Modeling aspects of the language of life through transfer-learning protein sequences. BMC Bioinformatics , 20(1):723, December 2019. ISSN 1471-2105. doi: 10.1186/s12859-019-3220-8. URL https://doi.org/10.1186/s12859-019-3220-8 .

Christian Dallago, Konstantin Schütze, Michael Heinzinger, Tobias Olenyi, Maria Littmann, Amy X. Lu, Kevin K. Yang, Seonwoo Min, Sungroh Yoon, James T. Morton, and Burkhard Rost. Learned Embeddings from Deep Learning to Visualize and Predict Protein Sets. Current Protocols , 1(5):e113, 2021. ISSN 2691-1299. doi: 10.1002/cpz1.113. URL https://onlinelibrary.wiley.com/doi/abs/10.1002/cpz1.113 . _eprint: https://onlinelibrary.wiley.com/doi/pdf/10.1002/cpz1.113.

Christian Dallago, Jody Mou, Kadina E. Johnston, Bruce Wittmann, Nick Bhattacharya, Samuel Goldman, Ali Madani, and Kevin K. Yang. FLIP: Benchmark tasks in fitness landscape inference for proteins. January 2022. URL https://openreview.net/forum?id=p2dMLEwL8tF .

Hannes Stärk, Christian Dallago, Michael Heinzinger, and Burkhard Rost. Light attention predicts protein location from the language of life. Bioinformatics Advances , 1(1):vbab035, January 2021. ISSN 2635-0041. doi: 10.1093/bioadv/vbab035. URL https://doi.org/10.1093/bioadv/vbab035 .

Thomas A. Hopf, John B. Ingraham, Frank J. Poelwijk, Charlotta P. I. Schärfe, Michael Springer, Chris Sander, and Debora S. Marks. Mutation effects predicted from sequence co-variation. Nature Biotechnology , 35(2): 128-135, February 2017a. ISSN 1546-1696. doi: 10.1038/nbt.3769.

Adam J. Riesselman, John B. Ingraham, and Debora S. Marks. Deep generative models of genetic variation capture the effects of mutations. Nature Methods , 15(10):816-822, October 2018. ISSN 1548-7105. doi: 10.1038/s41592-018-0138-4. URL https://www.nature.com/articles/s41592-018-0138-4 . Number: 10 Publisher: Nature Publishing Group.

Elodie Laine, Yasaman Karami, and Alessandra Carbone. GEMME: A Simple and Fast Global Epistatic Model Predicting Mutational Effects. Molecular Biology and Evolution , 36(11):2604-2619, November 2019. ISSN 0737-4038. doi: 10.1093/molbev/msz179. URL https://doi.org/10.1093/molbev/msz179 .

Jonathan Frazer, Pascal Notin, Mafalda Dias, Aidan Gomez, Joseph K. Min, Kelly Brock, Yarin Gal, and Debora S. Marks. Disease variant prediction with deep generative models of evolutionary data. Nature , 599(7883):91-95, November 2021. ISSN 1476-4687. doi: 10.1038/s41586-021-04043-8. URL https: //www.nature.com/articles/s41586-021-04043-8 . Number: 7883 Publisher: Nature Publishing Group.

Joshua Meier, Roshan Rao, Robert Verkuil, Jason Liu, Tom Sercu, and Alexander Rives. Language models enable zero-shot prediction of the effects of mutations on protein function. January 2022. URL https: //openreview.net/forum?id=uXc42E9ZPFs .

Pascal Notin, Mafalda Dias, Jonathan Frazer, Javier Marchena Hurtado, Aidan N. Gomez, Debora Marks, and Yarin Gal. Tranception: Protein Fitness Prediction with Autoregressive Transformers and Inference-time Retrieval. In Proceedings of the 39th International Conference on Machine Learning , pages 16990-17017. PMLR, June 2022a. URL https://proceedings.mlr.press/v162/notin22a.html . ISSN: 2640-3498.

Pascal Notin, Lood Van Niekerk, Aaron W. Kollasch, Daniel Ritter, Yarin Gal, and Debora Susan Marks. TranceptEVE: Combining Family-specific and Family-agnostic Models of Protein Sequences for Improved Fitness Prediction. December 2022b. URL https://openreview.net/forum?id=l7Oo9DcLmR1 .

Céline Marquet, Michael Heinzinger, Tobias Olenyi, Christian Dallago, Kyra Erckert, Michael Bernhofer, Dmitrii Nechaev, and Burkhard Rost. Embeddings from protein language models predict conservation and variant effects. Human Genetics , 141(10):1629-1647, October 2022. ISSN 1432-1203. doi: 10.1007/ s00439-021-02411-y.

Chloe Hsu, Hunter Nisonoff, Clara Fannjiang, and Jennifer Listgarten. Learning protein fitness models from evolutionary and assay-labeled data. Nature Biotechnology , 40(7):1114-1122, July 2022. ISSN 1087-0156, 1546-1696. doi: 10.1038/s41587-021-01146-5. URL https://www.nature.com/articles/ s41587-021-01146-5 .

Ethan C Alley, Grigory Khimulya, Surojit Biswas, Mohammed AlQuraishi, and George M Church. Unified rational protein engineering with sequence-only deep representation learning. bioRxiv , page 589333, 2019.

Jannik Kossen, Neil Band, Clare Lyle, Aidan N. Gomez, Tom Rainforth, and Yarin Gal. Self-Attention Between Datapoints: Going Beyond Individual Input-Output Pairs in Deep Learning, February 2022. URL http://arxiv.org/abs/2106.02584 . arXiv:2106.02584 [cs, stat] version: 2.

Kevin Kaichuang Yang, Zachary Wu, and Frances H. Arnold. Machine learning in protein engineering. arXiv: Biomolecules , 2018. URL https://api.semanticscholar.org/CorpusID:53789384 .

Sam Gelman, Sarah A Fahlberg, Pete Heinzelman, Philip A. Romero, and Anthony Gitter. Neural networks to learn protein sequence-function relationships from deep mutational scanning data. Proceedings of the National Academy of Sciences of the United States of America , 118, 2020. URL https://api.semanticscholar. org/CorpusID:226228756 .

Chase R. Freschlin, Sarah A Fahlberg, and Philip A. Romero. Machine learning to navigate fitness landscapes for protein engineering. Current opinion in biotechnology , 75:102713, 2022. URL https: //api.semanticscholar.org/CorpusID:248078341 .

Nadav Brandes, Dan Ofer, Yam Peleg, Nadav Rappoport, and Michal Linial. Proteinbert: a universal deeplearning model of protein sequence and function. Bioinformatics , 38:2102 - 2110, 2021. URL https: //api.semanticscholar.org/CorpusID:235219208 .

Daniel Hesslow, Niccoló Zanichelli, Pascal Notin, Iacopo Poli, and Debora S. Marks. Rita: a study on scaling up generative protein sequence models. ArXiv , abs/2205.05789, 2022. URL https://api.semanticscholar. org/CorpusID:248722069 .

Ahmed Elnaggar, Hazem Essam, Wafaa Salah-Eldin, Walid Moustafa, Mohamed Elkerdawy, Charlotte Rochereau, and Burkhard Rost. Ankh: Optimized protein language model unlocks general-purpose modelling. bioRxiv , 2023a. URL https://api.semanticscholar.org/CorpusID:255941875 .

Nadav Brandes, Grant Goldman, Charlotte H. Wang, Chun Jimmie Ye, and Vasilis Ntranos. Genome-wide prediction of disease variants with a deep protein language model. bioRxiv , 2022. URL https://api. semanticscholar.org/CorpusID:251910813 .

Milind Jagota, Chengzhong Ye, Carlos Albors, Ruchir Rastogi, Antoine Koehl, Nilah M. Ioannidis, and Yun S. Song. Cross-protein transfer learning substantially improves disease variant prediction. Genome Biology , 24, 2023. URL https://api.semanticscholar.org/CorpusID:259077035 .

Brian L. Hie, Ellen D. Zhong, Bonnie Berger, and Bryan D. Bryson. Learning the language of viral evolution and escape. Science , 371:284 - 288, 2020. URL https://api.semanticscholar.org/CorpusID: 220516287 .

Nicole N. Thadani, Sarah Gurev, Pascal Notin, Noor Youssef, Nathan J. Rollins, Chris Sander, Yarin Gal, and Debora S. Marks. Learning from prepandemic data to forecast viral escape. Nature , 622:818 - 825, 2023. URL https://api.semanticscholar.org/CorpusID:263906690 .

Zachary Wu, Sun Kan, Russell D. Lewis, Bruce J. Wittmann, and Frances H. Arnold. Machine learning-assisted directed protein evolution with combinatorial libraries. Proceedings of the National Academy of Sciences , 116:8852 - 8858, 2019. URL https://api.semanticscholar.org/CorpusID:67770057 .

Surojit Biswas, Grigory Khimulya, Ethan C. Alley, Kevin M. Esvelt, and George M. Church. Low-N protein engineering with data-efficient deep learning. Nature Methods , 18(4):389-396, April 2021b. ISSN 1548-7105. doi: 10.1038/s41592-021-01100-y.

Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Neural Information Processing Systems , 2017. URL https://api.semanticscholar.org/CorpusID:13756489 .

Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial Attention in Multidimensional Transformers, December 2019. URL http://arxiv.org/abs/1912.12180 . arXiv:1912.12180 [cs].

Roshan M. Rao, Jason Liu, Robert Verkuil, Joshua Meier, John Canny, Pieter Abbeel, Tom Sercu, and Alexander Rives. MSA Transformer. In Proceedings of the 38th International Conference on Machine Learning , pages 8844-8856. PMLR, July 2021. URL https://proceedings.mlr.press/v139/rao21a.html . ISSN: 2640-3498.

Damiano Sgarbossa, Umberto Lupo, and Anne-Florence Bitbol. Generative power of a protein language model trained on multiple sequence alignments. eLife , 12:e79854, February 2023. ISSN 2050-084X. doi: 10.7554/eLife.79854.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers) , pages 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v1/N19-1423. URL https://aclanthology.org/N19-1423 .

Roshan Rao, Nicholas Bhattacharya, Neil Thomas, Yan Duan, Xi Chen, John Canny, Pieter Abbeel, and Yun S. Song. Evaluating Protein Transfer Learning with TAPE, June 2019. URL http://arxiv.org/abs/1906. 08230 . arXiv:1906.08230 [cs, q-bio, stat].

Ahmed Elnaggar, Hazem Essam, Wafaa Salah-Eldin, Walid Moustafa, Mohamed Elkerdawy, Charlotte Rochereau, and Burkhard Rost. Ankh: Optimized Protein Language Model Unlocks General-Purpose Modelling, January 2023b. URL http://arxiv.org/abs/2301.06568 . arXiv:2301.06568 [cs, q-bio].

Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In International Conference on Machine Learning , 2015. URL https://api. semanticscholar.org/CorpusID:160705 .

Donatas Repecka, Vykintas Jauniskis, Laurynas Karpus, Elzbieta Rembeza, Irmantas Rokaitis, Jan Zrimec, Simona Poviloniene, Audrius Laurynenas, Sandra Viknander, Wissam Abuajwa, Otto Savolainen, Rolandas Meskys, Martin K. M. Engqvist, and Aleksej Zelezniak. Expanding functional protein sequence spaces using generative adversarial networks. Nature Machine Intelligence , 3(4):324-333, April 2021. ISSN 2522-5839. doi: 10.1038/s42256-021-00310-5. URL https://www.nature.com/articles/s42256-021-00310-5 . Number: 4 Publisher: Nature Publishing Group.

Ali Madani, Ben Krause, Eric R. Greene, Subu Subramanian, Benjamin P. Mohr, James M. Holton, Jose Luis Olmos, Caiming Xiong, Zachary Z. Sun, Richard Socher, James S. Fraser, and Nikhil Naik. Large language models generate functional protein sequences across diverse families. Nature Biotechnology , pages 1-8, January 2023. ISSN 1546-1696. doi: 10.1038/s41587-022-01618-2.

Karen S. Sarkisyan, Dmitry Bolotin, Margarita V. Meer, Dinara R. Usmanova, Alexander S. Mishin, George V. Sharonov, Dmitry N. Ivankov, Nina G. Bozhanova, Mikhail S. Baranov, Onuralp Soylemez, Natalya S. Bogatyreva, Peter K. Vlasov, Evgeny S. Egorov, Maria D. Logacheva, Alexey S. Kondrashov, Dmitry M. Chudakov, Ekaterina V. Putintseva, Ilgar Z. Mamedov, Dan S. Tawfik, Konstantin A. Lukyanov, and Fyodor A. Kondrashov. Local fitness landscape of the green fluorescent protein. Nature , 533:397 - 401, 2016. URL https://api.semanticscholar.org/CorpusID:4472875 .

Samuel Thompson, Yang Zhang, Christine Ingle, Kimberly A. Reynolds, and Tanja Kortemme. Altered expression of a quality control protease in e. coli reshapes the in vivo mutational landscape of a model enzyme. eLife , 9, 2020a. URL https://api.semanticscholar.org/CorpusID:220717580 .

Louisa Gonzalez Somermeyer, Aubin Fleiss, Alexander S. Mishin, Nina G. Bozhanova, Anna A. Igolkina, Jens Meiler, Maria-Elisenda Alaball Pujol, Ekaterina V. Putintseva, Karen S. Sarkisyan, and Fyodor A. Kondrashov. Heterogeneity of the GFP fitness landscape and data-driven protein design. eLife , 11:e75842, May 2022. ISSN 2050-084X. doi: 10.7554/eLife.75842.

Nicholas C Wu, Lei Dai, C Anders Olson, James O Lloyd-Smith, and Ren Sun. Adaptation in protein fitness landscapes is facilitated by indirect paths. eLife , 5:e16965, July 2016a. ISSN 2050-084X. doi: 10.7554/eLife.16965.

Daniel D. Brauer, Celine B. Santiago, Zoe N. Merz, Esther McCarthy, Danielle Tullman-Ercek, and Matthew B. Francis. Comprehensive Fitness Landscape of a Multi-Geometry Protein Capsid Informs Machine Learning Models of Assembly, December 2021.

David Ding, Anna G. Green, Boyuan Wang, Thuy-Lan Vo Lite, Eli N. Weinstein, Debora S. Marks, and Michael T. Laub. Co-evolution of interacting proteins through non-contacting and non-specific mutations. Nature Ecology & Evolution , 6(5):590-603, May 2022. ISSN 2397-334X. doi: 10.1038/s41559-022-01688-0.

Jeffrey M. Spencer and Xiaoliu Zhang. Deep mutational scanning of S. pyogenes Cas9 reveals important functional domains. Scientific Reports , 7(1):16836, December 2017. ISSN 2045-2322. doi: 10.1038/ s41598-017-17081-y.

Philipp Koch, Steven Schmitt, Alexander Heynisch, Anja Gumpinger, Irene Wüthrich, Marina Gysin, Dimitri Shcherbakov, Sven N. Hobbie, Sven Panke, and Martin Held. Optimization of the antimicrobial peptide Bac7 by deep mutational scanning. BMC Biology , 20(1):114, May 2022. ISSN 1741-7007. doi: 10.1186/ s12915-022-01304-4.

Tobias Stadelmann, Daniel Heid, Michael Jendrusch, Jan Mathony, Stéphane Rosset, Bruno E. Correia, and Dominik Niopek. A deep mutational scanning platform to characterize the fitness landscape of anti-CRISPR proteins, August 2021.

Sarah Gersing, Matteo Cagiada, Marinella Gebbia, Anette P. Gjesing, Atina G. Coté, Gireesh Seesankar, Roujia Li, Daniel Tabet, Amelie Stein, Anna L. Gloyn, Torben Hansen, Frederick P. Roth, Kresten Lindorff-Larsen, and Rasmus Hartmann-Petersen. A comprehensive map of human glucokinase variant activity, June 2022.

Helen T. Hobbs, Neel H. Shah, Sophie R. Shoemaker, Jeanine F. Amacher, Susan Marqusee, and John Kuriyan. Saturation mutagenesis of a predicted ancestral Syk-family kinase. Protein Science: A Publication of the Protein Society , 31(10):e4411, October 2022. ISSN 1469-896X. doi: 10.1002/pro.4411.

Christopher D. Aakre, Julien Herrou, Tuyen N. Phung, Barrett S. Perchuk, Seán Crosson, and Michael T. Laub. Evolving new protein-protein interaction specificity through promiscuous intermediates. Cell , 163:594-606, 2015. URL https://api.semanticscholar.org/CorpusID:16389199 .

Nicholas C. Wu, Lei Dai, Anders Olson, James O. Lloyd-Smith, and Ren Sun. Adaptation in protein fitness landscapes is facilitated by indirect paths. eLife , 5, 2016b. URL https://api.semanticscholar.org/ CorpusID:16039694 .

Zihang Jiang, Weihao Yu, Daquan Zhou, Yunpeng Chen, Jiashi Feng, and Shuicheng Yan. ConvBERT: Improving BERTwith Span-based Dynamic Convolution, February 2021. URL http://arxiv.org/abs/2008.02496 . arXiv:2008.02496 [cs].

Jacquier, A. Birgy, H. Le Nagard, Y. Mechulam, E. Schmitt, J. Glodt, B. Bercot, E. Petit, J. Poulain, G. Barnaud, P.-A. Gros, and O. Tenaillon. Capturing the mutational landscape of the beta-lactamase TEM-1. Proceedings of the National Academy of Sciences , 110(32):13067-13072, August 2013. ISSN 0027-8424, 1091-6490. doi: 10.1073/pnas.1215206110. URL http://www.pnas.org/cgi/doi/10.1073/pnas. 1215206110 .

Jochen Weile, Song Sun, Atina G Cote, Jennifer Knapp, Marta Verby, Joseph C Mellor, Yingzhou Wu, Carles Pons, Cassandra Wong, Natascha Lieshout, Fan Yang, Murat Tasan, Guihong Tan, Shan Yang, Douglas M Fowler, Robert Nussbaum, Jesse D Bloom, Marc Vidal, David E Hill, Patrick Aloy, and Frederick P Roth. A framework for exhaustively mapping functional missense variants. Molecular Systems Biology , 13(12):957, December 2017. ISSN 1744-4292, 1744-4292. doi: 10.15252/msb.20177908. URL https: //onlinelibrary.wiley.com/doi/10.15252/msb.20177908 .

Samuel Thompson, Yang Zhang, Christine Ingle, Kimberly A Reynolds, and Tanja Kortemme. Altered expression of a quality control protease in E. coli reshapes the in vivo mutational landscape of a model enzyme. eLife , 9: e53476, July 2020b. ISSN 2050-084X. doi: 10.7554/eLife.53476. URL https://elifesciences.org/ articles/53476 .

Richard N. McLaughlin Jr, Frank J. Poelwijk, Arjun Raman, Walraj S. Gosal, and Rama Ranganathan. The spatial architecture of protein function and adaptation. Nature , 491(7422):138-142, November 2012. ISSN 00280836, 1476-4687. doi: 10.1038/nature11500. URL http://www.nature.com/articles/nature11500 .

Andrew O. Giacomelli, Xiaoping Yang, Robert E. Lintner, James M. McFarland, Marc Duby, Jaegil Kim, Thomas P. Howard, David Y. Takeda, Seav Huong Ly, Eejung Kim, Hugh S. Gannon, Brian Hurhula, Ted Sharpe, Amy Goodale, Briana Fritchman, Scott Steelman, Francisca Vazquez, Aviad Tsherniak, Andrew J. Aguirre, John G. Doench, Federica Piccioni, Charles W. M. Roberts, Matthew Meyerson, Gad Getz, Cory M. Johannessen, David E. Root, and William C. Hahn. Mutational processes shape the landscape of TP53 mutations in human cancer. Nature Genetics , 50(10):1381-1387, October 2018. ISSN 1061-4036, 1546-1718. doi: 10.1038/s41588-018-0204-y. URL http://www.nature.com/articles/s41588-018-0204-y .

Jason D. Fernandes, Tyler B. Faust, Nicolas B. Strauli, Cynthia Smith, David C. Crosby, Robert L. Nakamura, Ryan D. Hernandez, and Alan D. Frankel. Functional Segregation of Overlapping Genes in HIV. Cell , 167(7):1762-1773.e12, December 2016. ISSN 00928674. doi: 10.1016/j.cell.2016.11.031. URL https: //linkinghub.elsevier.com/retrieve/pii/S0092867416316038 .

Benjamin P. Roscoe, Kelly M. Thayer, Konstantin B. Zeldovich, David Fushman, and Daniel N.A. Bolon. Analyses of the Effects of All Ubiquitin Point Mutants on Yeast Growth Rate. Journal of Molecular Biology , 425(8):1363-1377, April 2013. ISSN 00222836. doi: 10.1016/j.jmb.2013.01.032. URL https: //linkinghub.elsevier.com/retrieve/pii/S0022283613000636 .

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In ICLR , 2019.

Jannis Born and Matteo Manica. Regression Transformer enables concurrent sequence regression and generation for molecular language modelling. Nature Machine Intelligence , 5(4):432-444, April 2023. ISSN 2522-5839. doi: 10.1038/s42256-023-00639-z. URL https://www.nature.com/articles/s42256-023-00639-z . Number: 4 Publisher: Nature Publishing Group.

Ahmed Elnaggar, Michael Heinzinger, Christian Dallago, Ghalia Rihawi, Yu Wang, Llion Jones, Tom Gibbs, Tamas Feher, Christoph Angerer, Martin Steinegger, et al. Prottrans: towards cracking the language of life's code through self-supervised deep learning and high performance computing. arXiv preprint arXiv:2007.06225 , 2020.

Martin Steinegger, Markus Meier, Milot Mirdita, Harald Vöhringer, Stephan J. Haunsberger, and Johannes Söding. Hh-suite3 for fast remote homology detection and deep protein annotation. BMC Bioinformatics , 20, 2019.

Thomas A Hopf, John B Ingraham, Frank J Poelwijk, Charlotta PI Schärfe, Michael Springer, Chris Sander, and Debora S Marks. Mutation effects predicted from sequence co-variation. Nature biotechnology , 35(2): 128-135, 2017b.

Appendix

A ProteinGym extensions

A.1 Curation of additional DMS assays

In this work we extensively leverage the ProteinGym benchmark introduced in Notin et al. [2022a]. We further expand the available pool of DMS assays by incorporating 13 additional assays, resulting in a total of 100 DMS assays. We refer to this broader pool in the text as 'the extended benchmark'. The selection of these assays was guided by various criteria, including recency (we included recent DMS experiments published in 2022 and 2023, not available when ProteinGym was initially released) and quality (e.g., high number of mutants, high correlation between replicates). We also favored the inclusion of assays that contain multiple mutants, as it was one of the shortcomings of the first ProteinGym benchmarks. Lastly, we focused on DMS assays that involved proteins relevant to protein design, such as GFP, capsid proteins, antitoxins, antimicrobial agents, CRISPR/Cas9, and kinases. We give further details on all assays that belong to the extended benchmark in the reference file included in the code repository. The 13 new assays that we include are shown in Table 3.

For consistency, we pre-processed the raw DMS data following the approach from Notin et al. [2022a]. As such, we exclude silent mutations from our benchmark, which are nucleotide substitutions that do not affect the resulting protein sequence. Duplicate mutants, either arising from repeated nucleotide substitutions or indels leading to identical protein sequences, are removed by calculating the average of all DMS measurements across duplicates. Finally, mutants with missing assay measurements are excluded.

A.2 Creation of cross-validation schemes

As described in the text, we designed 3 types of cross-validation schemes: Random, Contiguous and Modulo. For the Random scheme, we randomly assign each mutation to one of 5 folds, which may result in mutations at the same position to be assigned to different folds. The Contiguous scheme is obtained by splitting the sequence in contiguous segments along its length. We ensure that the segments are comprised of the same number of positions by first counting the total number of mutated position in the sequence. If this total is not divisible by the total number of fold, the first few folds will include one more position than the rest. We only consider positions mutated, which may not span the entire length of the protein sequence. 'Contiguous' therefore means contiguous only with respect to the mutated range. Lastly, the Modulo scheme is obtained by assigning positions to each fold using the modulo of the residue position. Therefore, for a 5 fold cross-validation, position 1 is

assigned to fold 1, position 2 to fold 2, ..., position 6 to fold 1, etc. Once again, we make sure to only consider mutated positions. If position 12 is assigned to fold 2, and position 13 and 14 are not mutated, position 15 will be assigned to fold 3. We operate a five fold cross-validation for all assays except for assays F7YBW8 MESOWAakre et al. [2015] and SPG1 STRSG Wu et al. [2016b], as these contain only 4 mutated positions. We therefore do a 4-fold cross-validation for these two proteins (the Contiguous and Modulo schemes are therefore identical for these two proteins). Note that the Contiguous and Modulo cross-validation schemes are only used when assessing the effects of single mutants. Multiple mutants generally involve several positions that may not be easily separated into the independent folds as discussed in the Contiguous and Modulo schemes above. Thus, all analyses about multiple mutants conducted in this work (for example in § 4.4) are based on the Random cross-validation scheme only.

B ProteinNPT details

B.1 Model architecture

Instead of learning a parametric model that makes protein property predictions based on the features of a single input protein sequence, our model makes predictions both based on these features but also based on similarity with a mini-batch of already labeled instances via self-attention. Using similarity between protein sequences is a cornerstone of computational biology and protein modeling via the use of Multiple Sequence Alignments (MSAs). The MSA Transformer showed how self-attention across the sequences of an MSA is a very effective way to learn compact protein representations in an unsupervised setting, and the idea of performing self-attention across labeled instances has first been proposed by Kossen et al. [2022] - we discuss how ProteinNPT relates to these two works in more details in Appendix B.5.

We summarize our architecture in Fig. 1. We embed the input mini-batch of protein sequences with a pre-trained protein language models. We experimented with various models in our ablations (Appendix C.2), obtained superior performance with the MSA Transformer and used this model in all subsequent analyses discussed in this work. Assuming the input sequence is of length L , the resulting embedded mini-bath of B sequences is of size B.L.D , where D is the dimension for the embedding of a single amino acid token 2 . To prevent overfitting when training a ProteinNPT model on a relatively small labeled dataset, the embedding dimension d used in subsequent self-attention layers of the network is smaller than D (e.g., D = 768 for the MSA Transformer, and d = 200 in our final architecture based on ablations). We apply linear projections to reduce dimensionality accordingly.

In what follows, we assume the input target labels and auxiliary labels are both unidimensional continuous inputs, but everything seamlessly extends to continuous or categorical inputs of arbitrary dimensions (our codebase handles both input types). In our final ProteinNPT architecture the auxiliary labels are zero-shot fitness predictions obtained with the MSA Transformer, following the procedure described in Meier et al. [2022]. To disambiguate masked entries from true input zeros, we append to continuous vectors a binary flag indicating whether the corresponding label is masked (for categorical input, a separate category corresponds to masked targets). After applying standard scaling on the input targets, we use linear layers (a separate one for each target and each auxiliary label) to project them into a d -dimensional vector. We then concatenate the protein sequence embeddings with the input target embeddings and auxiliary label embeddings.

Wealso experiment with various transforms of the resulting embeddings: Convolutional layer (CNN), ConvBERT layer [Jiang et al., 2021], or no transform. We find in our ablations that applying a CNN layer delivers the best performance on downstream tasks performance.

The transformed embeddings are subsequently fed into 5 consecutive ProteinNPT layers (chosen based on ablations). Each ProteinNPT layer applies successively row-attention, column-attention and a feedforward layer. Each of these transforms is preceded by a LayerNorm operator and we add residual connections to the output of each step. As proposed in the MSA Transformer, we leverage tied row-wise attention, ie., the attention maps in row attention are averaged across rows, resulting in lower memory footprint. This is a sensible design for the applications studied in our work, in which labeled sequences are relatively similar to one another and share a common three-dimensional contact structure.

Finally, we make predictions for the targets of interests by feeding the corresponding target embeddings into a L2-penalized linear projector (a separate linear map is learned for each target). We summarize all architecture details from our final architecture in Table 4.

B.2 Ablations

We carried out thorough ablations to develop ProteinNPT. Following [Meier et al., 2022, Notin et al., 2022a], to support theses various ablations, we set aside a relatively small set of DMS assays (8 assays out of the 100 assays in the extended ProteinGym substitution benchmark), which are used to decide between different model architectures while not overfitting these decisions to the benchmark:

BLAT ECOLX [Jacquier et al., 2013]

CALM1 HUMAN [Weile et al., 2017]

DYR ECOLI [Thompson et al., 2020b]

DLG4 RAT [McLaughlin Jr et al., 2012]

P53 HUMAN [Giacomelli et al., 2018]

REV HV1H2 [Fernandes et al., 2016]

RL401 YEAST [Roscoe et al., 2013]

TAT HV1BR [Fernandes et al., 2016]

Together, these 8 assays cover diverse protein families in terms of taxa, mutation depths, sequence lengths and MSA depth.

We report results for several of these ablations in this subsection, in particular varying the model used to obtain protein sequence embeddings (Table 7) and the use of auxiliary labels (Table 6). Based on these analyses, we use for our final model architecture the MSA Transformer [Rao et al., 2021] to obtain protein sequence embeddings, use auxiliary labels (ie. zero-shot fitness predictions with MSA Transformer), and use batches with 1k training instances to make predictions at inference.

B.3 Model training

At train time, we feed as input to our model a batch of 425 protein sequences and the corresponding labels. Since we mask at random 15% of labels, this equates to having on average 64 masked training instances - which we want to accurately predict the values of, and 361 unmasked instances - which are used to support the predictions of masked instances. We also mask at random 15% of input protein sequence tokens. As discussed in § 3, our training loss is the weighted average of two different losses: the Cross-Entropy (CE) loss over masked tokens predictions (akin to a standard masked-language modeling objective [Devlin et al., 2019]) and the Mean Squared Error (MSE) over the masked targets predictions. Throughout training, we vary the weight used to balance out these two losses, starting with equal weights at the beginning of training, and annealing progressively the token prediction objective with a cosine schedule over the course of training.

We train our model with the AdamW optimizer [Loshchilov and Hutter, 2019]. We experimented with various early stopping criteria to train our model by setting aside a validation set, monitoring the validation MSE or Spearman's rank correlation between prediction and targets, varying the patience parameter (e.g., 3, 5 and 10) but obtained systematically better performance on the downstream tasks from our 'validation benchmark' when instead training for a fixed number of training steps (10k gradient steps). All training hyperparameters are summarized in Table 5.

Lastly, before training a ProteinNPT model on a given assay, we compute and persist to disk the sequence embeddings for all mutated proteins in that assay. During training, we load from disk the embeddings corresponding to each mini batch. Since the parameters of the model used to extract sequence embeddings are frozen, this procedure is strictly equivalent to computing embeddings on the fly. However, training is much faster as embeddings are computed only once and loading from disk is relatively fast. Additionally, we can work with larger training batches as memory requirements are lower - computing embeddings with large protein language models can necessitate substantial amounts of GPU memory.

B.4 Inference

During the inference process, our objective is to make predictions for unlabeled protein sequences. These sequences are handled in our model using the same approach as we employed for sequences with masked targets during training. The input mini-batch is also comprised of labeled sequences from the training set, which are used to support the predictions of unlabeled instances. Ideally we would fit the entire training set when making predictions. However, this approach is often hindered by GPU memory bottlenecks in practice when the training set is sufficiently large. To that end, we set an upper bound M on the maximum number of training sequences we can leverage at inference time: if the training dataset is smaller than M , we use all training sequences with no filtering; if the training dataset is larger, we sample at random M sequences with no replacement. As seen in ablations above, we choose a value of M = 1 k in practice as larger values only yield marginal gains on downstream tasks performance. We also experimented with a sampling scheme with replacement in earlier ablations, but obtained systematically comparable performance with the scheme described above. The Regression Transformer [Born and Manica, 2023] can be seen as a special case of our architecture in which column-wise attention is ablated. In our experiments, this however results in a significant performance drop (Table 2).

B.5 Difference with axial transformer, MSA Transformer and Non-parametric transformers

Our work builds on prior work from the self-attention literature, in particular the Axial Transformer [Ho et al., 2019], MSA Transformer [Rao et al., 2021] and Non-Parametric Transformer [Kossen et al., 2022].

Focusing on image and video modeling tasks, the Axial Transformer introduced the concept of axial attention, alternating in a given self-attention layer between row-wise attention and column-wise attention. We use this concept throughout our architecture to operate self-attention across residues (in embedding layers), across

residues and labels (in ProteinNPT layers), across sequences of a retrieved MSA (in embeddings layers) and across labeled observations (in ProteinNPT layers).

The MSA Transformer leveraged axial attention to learn protein sequence representations by training across a large number of Multiple Sequence Alignments (MSAs). Axial attention is used to alternatively learn dependencies across residues in the protein sequences, and across MSA sequences at a given position. The MSA Transformer is a fully unsupervised architecture and does not leverage available labels. Authors also show that a regression trained on attention maps of the MSA transformer achieve strong performance in predicting the residues that are in contact in the tertiary structure of the protein. It is an integral part of our ProteinNPT architecture, allowing us to obtain high-quality embeddings for each residue, which are then leveraged in the subsequent ProteinNPT layers.

Focusing on tabular datasets, Non-Parametric transformers (NPTs) leveraged axial attention to learn dependencies across features and labels of the tabular input, and across labeled instances. ProteinNPT differs from and extends NPTs in the following ways:

In lieu of tabular inputs, we focus on a different data modality (protein sequences);

Given the scarcity of labels involved with respect to the complexities of the underlying objects and their relations with the targets of interest, our architecture is semi-supervised and we leverage protein embeddings pre-trained on large quantities of unlabeled natural sequences;

In our final architecture, we operate tri-axial self-attention as described above (instead of bi-axial self-attention in MSA Transformer and in standard NPTs);

We introduce the concept of auxiliary labels which is critical for us to achieve strong performance, but it also broadly applicable beyond the protein-specific use cases we focus on in this work (§ 3);

As discussed in Appendix C.1, we leverage tied row attention ;

We explore the multi-task learning setting (§ 4.5 and Appendix F);

We investigate uncertainty quantification to support the Bayesian optimization tasks (Appendix G.1).

We demonstrate the ability of our model to support conditional sampling of new sequences - which is critically important to prioritize sensible points to subsequently score in silico (§ 5.2);

C Baselines details

C.1 Model architecture

As discussed in § 4.2, our first baseline is the top performing model that was introduced in Hsu et al. [2022], which learns a ridge regression on top of One-Hot-Encoded (OHE) representations of input protein sequences and unsupervised fitness predictions from DeepSequence [Riesselman et al., 2018]. We then sought to improve upon this initial baseline in two ways: 1) superior zero-shot fitness predictors 2) improved protein representations.

For the former, we ran ablations in which we replaced the zero-shot fitness predictions from DeepSequence with more recent alternatives: ESM-1v [Meier et al., 2022], MSA Transformer [Rao et al., 2021], Tranception [Notin et al., 2022a] and TranceptEVE [Notin et al., 2022b](see Appendix C.2). We generally see that gains in zero-shot fitness predictions translate into increased performance of the corresponding semi-supervised 'OHE - augmented' models.

For the latter, we build on the extensive literature [Heinzinger et al., 2019, Alley et al., 2019, Dallago et al., 2021, 2022, Elnaggar et al., 2023b, Stärk et al., 2021] of semi-supervised models for protein property prediction. These works typically first extract embeddings for labeled protein sequences with a large pre-trained protein language model. Given the dimensionality of the resulting full sequence embeddings - relatively large in practice compared with the size of the training data, a dimensionality reduction technique is typically used. The most commonly used approach consists in applying a mean pooling [Alley et al., 2019] operator across all the tokens of the input sequences. After pooling, a non-linear transform (e.g., MLP, CNN) is used to learn complex dependencies between the corresponding protein representations and the targets of interest. Albeit simple, this approach has lead to strong performance across various benchmarks [Dallago et al., 2022]. We then propose to combine these ideas with the zero-shot fitness prediction augmentations from Hsu et al. [2022] to obtain superior baselines. The corresponding baselines are referred to as 'Embeddings - Augmented'.

Wecarry out extensive ablations to optimize this new baseline further, and present a summary of the main variants tested in the next subsection (Appendix C.2). Our final top-performing baseline ('Embeddings - Augmented (MSA Transformer)') is based on embeddings obtained with the MSA Transformer, fed into a CNN layer with kernel size 9, followed by a ReLU non-linear activation, mean-pooled across sequence length, concatenated with zero-shot fitness predictions, and then passed to a linear projection for the final prediction. Following Hsu et al. [2022], the L2-penalty on the zero-shot fitness prediction parameter is set to a much lower value ( 10 -8 ) compared with other parameters ( 5 . 10 -2 ). We summarize details of the final architecture in Table 9.

As seen in Table 1, our improved baseline ('Embeddings - Augmented (MSA Transformer)') outperforms all OHE-based variants across all cross-validation schemes, including the ones augmented with the current best-in-class zero-shot fitness predictors. The lift is particularly important on the 'Contiguous' and 'Modulo' cross-validation schemes, for which the non-augmented OHE baseline performs just as poorly as random guessing, and the augmented equivalents perform on par with zero-shot fitness predictors.

C.2 Ablations

We present in this section the main ablations conducted to give rise to our various baselines. We use the same subset of 8 DMS assays as described in Appendix B.2.

The first analysis (Table 11) looks at the effect of the model used for the zero-shot fitness predictions in OHE baselines. We find that the better the underlying zero-shot fitness predictor, the higher the performance of the corresponding OHE augmented model.

The second analysis (Table 12) investigates various pre-trained protein language model, and concludes that representations obtained with the MSA Transformer are the most robust across the various cross-validation schemes.

The third analysis (Table 13) compares several of the non-linear transforms that have been suggested by the prior literature to predict protein properties based on the sequence embeddings. More specifically, we investigated L2-penalized regression [Dallago et al., 2022], shallow Multi-Layer Perceptron (MLP) [Elnaggar et al., 2020], Convolutional Neural Network (CNN) [Dallago et al., 2022], ConvBERT [Elnaggar et al., 2023b] and Light Attention [Stärk et al., 2021]. We obtain our best results with a CNN layer with a kernel of size 9.

Lastly, we investigate the effect of the size of the training data on the model performance to assess the ability of the ProteinNPT architecture to generalize to shallow training settings (Table 8). We do not observe a drop in performance on datasets with fewer labeled instances relative to assays with more labels available during training, confirming the ability of the architecture to mitigate overfitting risks via its hybrid learning objective.

C.3 Model training

For all baselines ('OHE' and 'Embeddings'), the loss we optimize is the MSE loss between property predictions and true labels. Similarly to what we discussed in Appendix B.3, we train our various baselines for 10k steps, using the AdamW optimizer with a cosine learning rate schedule. All training hyperparameters are summarized in Table 10. At train time we find it helpful to pre-compute the MSA Transformer embeddings of all labeled examples to prevent GPU memory bottlenecks and speed up training. Following Rao et al. [2021], the embeddings for each DMS sequence are computed by first filtering sequences in the corresponding MSA with HHFilter [Steinegger et al., 2019] to ensure minimum coverage of 75% and maximum sequence identity of 90%, and then performing a weighted sampling of the MSA to extract 384 natural sequences. The sampling weights are based on the procedure described in Hopf et al. [2017b].

D Property prediction experiments for single mutants

D.1 Experiment details

In this analysis we seek to assess the ability of various models to accurately predict the properties of single amino acid substitutions. We performed this analysis on all single mutants from the 100 assays from the extended ProteinGym substitution benchmark, and considered the 3 cross-validation schemes described in Appendix A.2.

We compared the performance of the following models:

Zero-shot (MSA Transformer) : zero-shot predictions obtained with the MSA Transformer (ensembling over 5 weighted samples from the corresponding MSA, as described in Meier et al. [2022]);

OHE : ridge regression on trained one-hot-encodings of the input protein sequences;

OHE - Augmented (DeepSequence) : ridge regression on trained one-hot-encodings of the input protein sequences and zero-shot predictions from DeepSequence [Riesselman et al., 2018];

OHE - Augmented (MSA Transformer) : ridge regression on trained one-hot-encodings of the input protein sequences and zero-shot predictions from the MSA Transformer [Rao et al., 2021];

Embeddings - Augmented (MSA Transformer) : ridge regression on the concatenation of meanpooled embeddings from the MSA Transformer and zero-shot predictions from the MSA Transformer (as described in Appendix C);

ProteinNPT : final ProteinNPT architecture as described in § 3.

For brevity, we use the following shorthands in the Figures and Tables throughout this manuscript:

DS stands for DeepSequence

MSAT stands for MSA Transformer

PNPT stands for ProteinNPT

Embed. stands for Embeddings

Aug. stands for Augmented

D.2 Detailed performance results

We report the DMS-level and average performance for each cross-validation scheme in Fig 6 (Random crossvalidation scheme), Fig. 7 (Modulo cross-validation scheme) and Fig. 8 (Contiguous cross-validation scheme). We also report the standard error for the Spearman's rank performance metric in Table 14, which we compute via non-parametric bootstrap of the difference between the Spearman performance of a given model and that of

the best overall model (i.e., ProteinNPT), computed over 10k bootstrap samples from the set of proteins in the extended ProteinGym benchmark.

We observe that ProteinNPT outperforms all baselines across the different cross-validation schemes. The largest performance lift is observed with the Random cross-validation as the model fully benefits from its column attention modules by learning from mutants that occurred at the same positions in the training set.

E Property prediction experiments for multiple mutants

E.1 Experiment details

In this analysis we investigate the ability of the different models listed in Appendix D.1 to predict the effects of multiple mutations. We carry out this analysis on the subset of 17 DMS assays with multiple mutants from the extended ProteinGym substitution benchmark. As discussed in § 4.4, since multiple mutants occur across several positions simultaneously, we only conduct this analysis based on the Random cross-validation scheme.

E.2 Detailed performance results

We report aggregated performance and corresponding standard error for all models in Table 15. Performance is reported in terms of Spearman's rank correlation and Mean Squared Error (MSE) between model predictions and experimental measurements. In the zero-shot setting, a fitness predictor is only able to rank order mutants based on relative fitness, but is unable to predict the actual value of the experimental phenotype. Thus, we do not report MSE performance in the zero-shot setting and only provide Spearman performance. We also report DMS-level performance (Spearman) in Fig. 2 (left).

We find that ProteinNPT markedly outperforms all other baselines both in terms of Spearman's rank correlation and Mean Squared Error. As shown in Fig. 2 (right), this is the case both for aggregate performance and across all mutation depths.

Spearman

Models

Zero-shot (MSAT)

OHE - Aug. (MSAT)

OHE - Aug. (DeepSequence)

OHE

Embeddings - Aug. (MSAT)

ProteinNPT

CCDB_ECOL_Adkar_2012

BLAT_ECOLX_Firnberg_2014

A4GRB6_PSEAI_Chen_2020

SPIKE_SARS2_Starr_expr_2020 -

P84126_THETH_Chan_2017

P53_HUMAN_Giacomelli_WT_Nutlin_2018

RL401_YEAST_Roscoe_2013

GRB2_HUMAN_Faure_2021

CP2C9_HUMAN_Amorosi_activity_2021

SPG1_STRSG_Wu_2016

IF1_ECOLI_Kelsic_2016

RASH_HUMAN_Bandaru_2017 -

AMIE_PSEAE_Wrenbeck_2017 -

BLAT_ECOLX_Deng_2012

NUD15_HUMAN_Suiter_2020 -

D7PM05_CLYGR_Somermeyer_2022

TRPC_SACS2_Chan_2017

TRPC_THEMA_Chan_2017 -

R1AB_SARS2_Flynn_growth_2022-

POLG_HCVJF_Qi_2014

VKOR1_HUMAN_Chiasson_abundance_2020-

DLG4_RAT_McLaughlin_2012

SRC_HUMAN_Ahler_CD_2019-

NRAM_133A0_Jiang_standard_2016

RL401_YEAST_Roscoe_2014 -

KKA2_KLEPN_Melnikov_2014

SPIKE_SARS2_Starr_bind_2020 -

CP2C9_HUMAN_Amorosi_abundance_2021

B3VI55_LIPST_Klesmith_2015 -

YAP1_HUMAN_Araya_2012

ANCSZ_Hobbs_2022

TADBP_HUMAN_Bolognesi_2019 -

DLG4_HUMAN_Faure_2021 -

DYR_ECOL_Thompson_plusLon_2019

AOA192B1T2_9HIV1_Haddox_2018

A4_HUMAN_Seuma_2021

C6KNH7_9INFA_Lee_2018 -

Q59976_STRSQ_Romero_2015

Q6WV13_9MAX1_Somermeyer_2022

P53_HUMAN_Giacomelli_NULL_Etoposide_2018

RL401_YEAST_Mavor_2016 -

P53_HUMAN_Kotler_2018 -

MTH3_HAEAE_Rockah-Shmuel_2015 -

BLAT_ECOLX_acquier_2013

A0A2Z5U3Z0_9INFA_Doud_2016

UBC9_HUMAN_Weile_2017

I6TAH8_168A0_Doud_2015

P53_HUMAN_Giacomelli_NULL_Nutlin_2018

ESTA_BACSU_Nutschel_2020 -

CCDB_ECOL_Tripathi_2016 -

HSP82_YEAST_Mishra_2016 -

GAL4_YEAST_Kitzman_2015

PTEN_HUMAN_Matreyek_2021

SC6A4_HUMAN_Young_2021

BAC7_Held_2022

Q8WTC7_9CNID_Somermeyer_2022

SYUA_HUMAN_Newberry_2020 -

HXK4_HUMAN_Gersing_2022

TPOR_HUMAN_Bridgford_S505N_2020-

CAPSD_BPMS2_Brauer_2022

Q2N0S5_9HIV1_Haddox_2018 -

A0A247D711_LISMN_Stadelmann_2021

POLG_CXB3N_Mattenberger_2021

PTEN_HUMAN_Mighell_2018

AACC1_pSEAI_Dandage_2018

NCAP_134A1_Doud_2015

GFP_AEQVI_Sarkisyan_2016

TAT_HV1BR_Fernandes_2016 -

ADRB2_HUMAN_ones_2020

SUMO1_HUMAN_Weile_2017 -

TPMT_HUMAN_Matreyek_2018 -

MSH2_HUMAN_Jia_2020

ENV_HV1B9_DuenasDecamp_2016

MK01_HUMAN_Brenan_2016

F7YBW8_MESOW_Aakre_2015

BRCA1_HUMAN_Findlay_2018 -

HSP82_YEAST_Flynn_2019

A0A140D2T1_ZIKV_Sourisseau_growth_2019-

A4D664_9INFA_Soh_CCL141_2019

AOA2Z5U3Z0_9INFA_Wu_2014

TPK1_HUMAN_Weile_2017 -

ENV_HV1BR_Haddox_2016 -

UBE4B_MOUSE_Starita_2013

VKOR1_HUMAN_Chiasson_activity_2020 -

PA_134A1_Wu_2015 -

PARE_ECOLX_Ding_2022 -

CALM1_HUMAN_Weile_2017 -

SPCAS9_Spencer_positive_2022

AOA220GHA5_9CAUD_Stadelmann_2021

AOA119GEU1_NEIME_Kennouche_2019

GCN4_YEAST_Staller_induction_2018

SCN5A_HUMAN_Glazer_2019 -

SPCAS9_Spencer_negative_2022-

HIS7_YEAST_Pokusaeva_2019

Average

0.4

0.2

0.0

0.2

0.4

0.6

0.8

1.0

Spearman

Models

Zero-shot (MSAT)

OHE - Aug. (MSAT)

OHE - Aug. (DeepSequence)

OHE

Embeddings - Aug. (MSAT)

ProteinNPT

P53_HUMAN_Giacomelli_WT_Nutlin_2018

BLAT_ECOLX_Firnberg_2014

A4GRB6_PSEAI_Chen_2020 -

TADBP_HUMAN_Bolognesi_2019

CP2C9_HUMAN_Amorosi_activity_2021

NRAM_133A0_iang_standard_2016

PABP_YEAST_Melamed_2013 -

NUD15_HUMAN_Suiter_2020 -

CP2C9_HUMAN_Amorosi_abundance_2021

Q59976_STRSQ_Romero_2015

AMIE_PSEAE_Wrenbeck_2017

CAPSD_AAV2S_Sinai_substitutions_2021

D7PM05_CLYGR_Somermeyer_2022-

B3VI55_LIPST_Klesmith_2015 -

P84126_THETH_Chan_2017 -

BLAT_ECOLX_acquier_2013 -

KKA2_KLEPN_Melnikov_2014

P53_HUMAN_Giacomelli_NULL_Etoposide_2018 -

RL401_YEAST_Rosc0e_2013

POLG_HCVJF_Qi_2014

MTH3_HAEAE_Rockah-Shmuel_2015

CCDB_ECOL_Adkar_2012

SC6A4_HUMAN_Young_2021-

P53_HUMAN_Giacomelli_NULL_Nutlin_2018

TRPC_SACS2_Chan_2017

GRB2_HUMAN_Faure_2021

CCDB_ECOLI_Tripathi_2016-

A4_HUMAN_Seuma_2021

A0A192B1T2_9HIV1_Haddox_2018 -

C6KNH7_9INFA_Lee_2018 -

HXK4_HUMAN_Gersing_2022 -

VKOR1_HUMAN_Chiasson_abundance_2020 -

GAL4_YEAST_Kitzman_2015

RASH_HUMAN_Bandaru_2017 -

SPIKE_SARS2_Starr_expr_2020

PTEN_HUMAN_Matreyek_2021

Q2N0S5_9HIV1_Haddox_2018 -

P53_HUMAN_Kotler_2018

HSP82_YEAST_Mishra_2016

ANCSZ_Hobbs_2022

Q6WV13_9MAX_Somermeyer_2022

TRPC_THEMA_Chan_2017 -

POLG_CXB3N_Mattenberger_2021

SYUA_HUMAN_Newberry_2020 -

DLG4_RAT_McLaughlin_2012

GFP_AEQVI_Sarkisyan_2016

PTEN_HUMAN_Mighel_2018 -

MSH2_HUMAN_Jia_2020 -

ADRB2_HUMAN_ones_2020 -

UBC9_HUMAN_Weile_2017

TPMT_HUMAN_Matreyek_2018 -

DYR_ECOLI_Thompson_plusLon_2019

TPOR_HUMAN_Bridgford_S505N_2020 -

IF1_ECOLI_Kelsic_2016

DLG4_HUMAN_Faure_2021

RL401_YEAST_Mavor_2016

ESTA_BACSU_Nutschel_2020 -

BLAT_ECOLX_Deng_2012

BAC7_Held_2022

SRC_HUMAN_Ahler_CD_2019 -

SUMO1_HUMAN_Weile_2017 -

BRCA1_HUMAN_Findlay_2018

HSP82_YEAST_Flynn_2019

YAP1_HUMAN_Araya_2012

A0A140D2T1_ZIKV_Sourisseau_growth_2019

TAT_HV1BR_Fernandes_2016

AACC1_PSEAI_Dandage_2018

A4D664_9INFA_Soh_CCL141_2019

MKO1_HUMAN_Brenan_2016

Q8WTC7_9CNID_Somermeyer_2022

RL401_YEAST_Roscoe_2014

NCAP_134A1_Doud_2015

I6TAH8_168A0_Doud_2015 -

R1AB_SARS2_Flynn_growth_2022 -

KCNH2_HUMAN_Kozek_2020 -

GCN4_YEAST_Staller_induction_2018

AOA247D711_LISMN_Stadelmann_2021

VKOR1_HUMAN_Chiasson_activity_2020

A0A2Z5U3Z0_9INFA_Wu_2014

SPIKE_SARS2_Starr_bind_2020

F7YBW8_MESOW_Aakre_2015

ENV_HV1BR_Haddox_2016 -

UBE4B_MOUSE_Starita_2013-

TPK1_HUMAN_Weile_2017

ENV_HV1B9_DuenasDecamp_2016

PA_134A1_Wu_2015

CAPSD_BPMS2_Brauer_2022

SPG1_STRSG_Olson_2014

CALM1_HUMAN_Weile_2017 -

PARE_ECOLX_Ding_2022 -

SPCAS9_Spencer_positive_2022

SCN5A_HUMAN_Glazer_2019

AOA220GHA5_9CAUD_Stadelmann_2021

AOA119GEU1_NEIME_Kennouche_2019

SPCAS9_Spencer_negative_2022

HIS7_YEAST_Pokusaeva_2019 -

SPG1_STRSG_Wu_2016

Average

0.6

0.4

0.2

0.0

0.2

0.4

0.6

0.8

1.0

Spearman

Models

Zero-shot (MSAT)

OHE

OHE - Aug. (MSAT)

Embeddings - Aug. (MSAT)

OHE - Aug. (DeepSequence)

ProteinNPT

CCDB_ECOL_Tripathi_2016

BLAT_ECOLX_Firnberg_2014 -

TADBP_HUMAN_Bolognesi_2019 -

A4GRB6_PSEA1_Chen_2020

P53_HUMAN_Giacomelli_WT_Nutlin_2018

CP2C9_HUMAN_Amorosi_activity_2021

NRAM_133A0_iang_standard_2016

D7PM05_CLYGR_Somermeyer_2022

Q59976_STRSQ_Romero_2015

MTH3_HAEAE_Rockah-Shmuel_2015

BLAT_ECOLX_acquier_2013

NUD15_HUMAN_Suiter_2020 -

KKA2_KLEPN_Melnikov_2014

CAPSD_AAV2S_Sinai_substitutions_2021 -

TRPC_SACS2_Chan_2017 -

P84126_THETH_Chan_2017 -

CP2C9_HUMAN_Amorosi_abundance_2021

HSP82_YEAST_Mishra_2016

P53_HUMAN_Kotler_2018

SC6A4_HUMAN_Young_2021

P53_HUMAN_Giacomelli_NULL_Etoposide_2018

HXK4_HUMAN_Gersing_2022

VKOR1_HUMAN_Chiasson_abundance_2020

B3V155_LIPST_Klesmith_2015 -

AMIE_PSEAE_Wrenbeck_2017

GRB2_HUMAN_Faure_2021

RASH_HUMAN_Bandaru_2017

SYUA_HUMAN_Newberry_2020 -

Q2N0S5_9HIV1_Haddox_2018 -

TPMT_HUMAN_Matreyek_2018 -

ANCSZ_Hobbs_2022

AOA192B1T2_9HIV1_Haddox_2018

DYR_ECOL_Thompson_plusLon_2019

GAL4_YEAST_Kitzman_2015

P53_HUMAN_Giacomelli_NULL_Nutlin_2018

UBC9_HUMAN_Weile_2017 -

TAT_HV1BR_Fernandes_2016

A4_HUMAN_Seuma_2021

Q6WV13_9MAXI_Somermeyer_2022-

POLG_CXB3N_Mattenberger_2021

TPOR_HUMAN_Bridgford_S505N_2020

DLG4_RAT_McLaughlin_2012

BLAT_ECOLX_Deng_2012

DLG4_HUMAN_Faure_2021

POLG_HCVJF_Qi_2014

RL401_YEAST_Roscoe_2014

GFP_AEQVI_Sarkisyan_2016 -

ADRB2_HUMAN_ones_2020

BRCA1_HUMAN_Findlay_2018

PTEN_HUMAN_Mighell_2018

ESTA_BACSU_Nutschel_2020 -

SUMO1_HUMAN_Weile_2017-

CCDB_ECOLI_Adkar_2012

RL401_YEAST_Mavor_2016

AACC1_PSEAI_Dandage_2018

RL401_YEAST_Roscoe_2013

A0A247D711_LISMN_Stadelmann_2021

TRPC_THEMA_Chan_2017 -

Q8WTC7_9CNID_Somermeyer_2022-

PTEN_HUMAN_Matreyek_2021

KCNH2_HUMAN_Kozek_2020

SPIKE_SARS2_Starr_expr_2020

IF1_ECOLI_Kelsic_2016

A0A140D2T1_ZIKV_Sourisseau_growth_2019

A4D664_9INFA_Soh_CCL141_2019

HSP82_YEAST_Flynn_2019 -

GCN4_YEAST_Staller_induction_2018

BAC7_Held_2022

A0A2Z5U3Z0_9INFA_Doud_2016

YAP1_HUMAN_Araya_2012

ENV_HV1B9_DuenasDecamp_2016

VKOR1_HUMAN_Chiasson_activity_2020 -

MSH2_HUMAN_Jia_2020

SRC_HUMAN_Ahler_CD_2019 -

ENV_HV1BR_Haddox_2016 -

AOA2Z5U3Z0_9INFA_Wu_2014

F7YBW8_MESOW_Aakre_2015

NCAP_134A1_Doud_2015

MK01_HUMAN_Brenan_2016 -

R1AB_SARS2_Flynn_growth_2022 ~

UBE4B_MOUSE_Starita_2013

TPK1_HUMAN_Weile_2017

SPG1_STRSG_Olson_2014 -

CALM1_HUMAN_Weile_2017

SPIKE_SARS2_Starr_bind_2020 -

I6TAH8_168A0_Doud_2015

PA_134A1_Wu_2015

CAPSD_BPMS2_Brauer_2022 -

PARE_ECOLX_Ding_2022

SPCAS9_Spencer_positive_2022

REV_HV1H2_Fernandes_2016

AOA119GEU1_NEIME_Kennouche_2019

SCN5A_HUMAN_Glazer_2019

AOA220GHA5_9CAUD_Stadelmann_2021

SPCAS9_Spencer_negative_2022

HIS7_YEAST_Pokusaeva_2019 -

SPG1_STRSG_Wu_2016

Average

0.6

0.4

0.2

0.0

0.2

0.4

0.6

0.8

1.0

F Property prediction experiments for multiple properties

F.1 Experiment details

We seek to evaluate the ability of the models to predict several properties simultaneously for the same mutated sequences. This is often highly important in practice as the majority of real-life protein design use cases call for the joint optimization of several properties, or the optimization of one property under constraints over other properties.

To that end, we extract from the extended ProteinGym substitution benchmark a subset of proteins for which several phenotypes have been measured. The seven proteins with at least two distinct measurements we considered were as follows:

The three proteins with at least three measurements were as follows:

RL401 YEAST

BLAT ECOLX

P53 HUMAN

We focus on the Random cross-validation scheme, and consider all possible mutants across the different assays for the same protein, which sometimes result in certain target values being missing for a subset of mutants.

For the various baselines we compare against, a separate linear head is used to predict each target based on the concatenation of the mean-pooled embeddings and zero-shot predictions. The loss that we optimize is the sum of the MSE losses for each target. When a subset of target values are missing for a given observation, we ignore the corresponding missing values in the loss computation and only leverage the loss for non-missing target values for these training instances.

In ProteinNPT, we simply add as input as many target columns as required, using the same masking procedure as described in Appendix B.3. Each target is then predicted by inputting the corresponding last-layer target embedding into a target-specific linear projection. The loss we optimize is the same composite loss which combines cross-entropy loss for the masked token prediction and the sum of MSE losses for each target prediction. When a subset of target values are missing for a given observation, we simply mask the corresponding input values. This leads to an implicit imputation scheme internally which contributes to further boosting the predictive performance.

F.2 Detailed performance results

We report the performance of ProteinNPT and other baselines for the simultaneous predictions of two protein properties in Table 16, and three properties in Table 17. Consistent with the results from the above analyses, ProteinNPT outperforms all baselines across the different settings, both in terms of Spearman's rank correlation and MSE.

Model

ProteinNPT

Embeddings - Aug. (MSAT)

OHE - Aug. (MSAT)

OHE

0.7

0.6

0.5

Spearman

0.4

0.3

0.2

0.1

0.0

Assay 1

Assay 2

0.8

0.7

0.6

Spearman

0.5

0.4

0.3

0.2

0.1

0.0

Assay 1

Assay 2

Assay 3

G Bayesian optimization experiments

G.1 Uncertainty quantification

We developed and compared three uncertainty quantification schemes:

MC dropout: We use a fixed inference batch at inference, apply Monte Carlo dropout Gal and Ghahramani [2015] and use the standard deviation of the prediction across multiple forward passes as our uncertainty metric;

Batch resampling: For a given point that we want to quantify the uncertainty of, we perform several forward passes by completing the input batches with a different sample (with replacement) of training

0.85

0.80

Hybrid-Random

Hybrid-Contiguous

Hybrid -Modulo

0.75

MCdropout-Random

MCdropout-Contiguous

MC dropout-Modulo

M

0.70

Batchresampling-Random

Batchresampling-Contiguous

Batchresampling-Modulo

0.65

0.60

0.55

80

60

40

20

0

%of points removed based on uncertainty cutoff

points for each pass. No dropout is applied and we use the standard deviation across forward passes as our uncertainty metric;

Hybrid scheme: We combine the first two schemes together (ie, turning on dropout in the batch resampling sampling).

The hybrid scheme delivers superior performance in calibration curve analyses (Fig 10), for a fixed computational budget.

Uncertainty for all baseline models is calculated using MC dropout based on 25 forward passes through the corresponding models. The dropout rate is the same as in training, and is set to 0.1. Uncertainty for ProteinNPT models is obtained with the hybrid scheme based on 5 MC dropout samples for 5 distinct resampled batches, with both an attention dropout and an activation dropout set to 0.1. For each batch of sequences, we draw at random a set number (set to M = 1 k , or to the total number of training points if less than that for Bayesian optimization experiments) of datapoints from the training set, which are required for non-parametric inference.

G.2 Experiment details

We conduct the iterative protein redesign experiments on all DMS assays in the ProteinGym extended benchmark.

In this analysis we only consider single mutants, and exclude the multiple mutants that may be present in the original assay. We split assays in three groups based on the total number of datapoints they include:

For assays with less than 1250 points, we start with 50 labels and acquire N = 50 new labels at each iteration;

For assays with more than 1250 points and less than 2500, we start with 100 labels and acquire N = 100 new labels at each iteration;

For assays with more than 2500 points, we start with 50 labels and acquire N = 200 new labels at each iteration.

Note that at iteration 0, we therefore start with a few labeled examples, which explains why performance does not start at 0 on the performance plots. At each round, the available pool is constituted of all sequences minus the ones already selected. As such, we do not disregard any sequence, and all available sequences are scored using the upper confidence bound acquisition function. For robustness, we run each experiments three times, with

a different starting pool. We set the α parameter, controlling the exploration trade-off to 0.1, after optimising performance on a subset of 3 assays.

G.3 Detailed performance results

We plot the recall rate of high fitness points (top 3 deciles) as a function of the number of batch acquisition cycles for all DMS assays in the extended benchmark. We find that ProteinNPT performs better than baselines on almost all assays. Moreover, when ProteinNPT does not perform well, it is usually because the assay is 'hard' to learn, because the property is intrinsically difficult to predict or experimental measurements are noisy (or both) - on these assays the performance of all baselines is close to that of a random acquisition strategy. This is for example the case for TAT HV1BR, CALM1 HUMAN and REV HV1H2 experiments. Lastly, we observe that ProteinNPT performs better for assays with a larger number of mutants. Since real-world design scenarios are less restricted than the in silico settings we consider here, it is therefore encouraging to see that ProteinNPT performs better when larger mutant pools are available.

H Codebase and License

The entire codebase to train and evaluate ProteinNPT and the various baselines considered in this work is available on our GitHub repository ( https://github.com/OATML-Markslab/ProteinNPT ) under the MIT license.

I Compute

All experiments carried out in this work were conducted in Pytorch, on A100 GPUs with either 40GB or 80GB of GPU memory.

J Limits

As a semi-supervised pseudo-generative model, ProteinNPT is dependent on the number and quality of available labels used for training. The diversity of the 100 DMS assays from the extended ProteinGym benchmark provide a sense for the sensitivity of performance to various characteristics of the underlying protein family and available labels in the corresponding assay. While compute cost is typically a secondary consideration in protein design endeavours as it is dwarfed by the cost of wet lab experiments, one should note that ProteinNPT is more computationally expensive than the other baselines discussed in this work. This overhead directly depends on the protein sequence length and the number of training instances per mini-batch since both impact the compute and memory complexities of the row and column attention in ProteinNPT layers.

K Ethics

The models introduced in this work can be used to support the design of novel proteins of any kind. As such we acknowledge the dual potential for both positive and negative applications. We commit to exercising caution in their deployment (we will not be releasing model parameters for instance) and we are dedicated to actively promoting their responsible use.

A0A140D2T1_ZIKV_Sourisseau_growth_2019

A0A192B1T2_9HIV1_Haddox_2018

A0A119GEU1_NEIME_Kennouche_2019

0.6

0.5

0.7

0.5

0.4

0.4

0.3

0.5+

0.3

0.2

0.2

0.1

0.1

0.0

10

0.0

2

8

10

10

A0A2Z5U3Z0_9INFA_Doud_2016

A0A2Z5U3Z0_9INFA_Wu_2014

A4_HUMAN_Seuma_2021

0.8

0.5

0.7

0.4

0.6

.5

0.5

0.3

0.4

0.4

0.3

.2

0.1

.

0.1

0.0

0.0

0.0

4

6

10

4

6

10

6

10

A4D664_9INFA_Soh_CCL141_2019

A4GRB6_PSEAI_Chen_2020

AACC1_PSEAI_Dandage_2018

1.0

0.4

0.8

0.8

0.3

0.6

0.6

0.4

0.4

0.2

0.2

0.0

ADRB2_HUMAN_Jones_2020

10

0.0 L

AMIE_PSEAE_Wrenbeck_2017

10

B3VI55_LIPST_Klesmith_2015

6

10

0.5

0.8

0.8

0.7

0.7

0.4

0.6

0.5

0.3

0.4

0.4

0.2

0.3

0.3

0.2

.21

0.1

0.1

4

6

10

0.0

6

8

10

0.0

2

6

10

BLAT_ECOLX_Deng_2012

BLAT_ECOLX_Firnberg_2014

BLAT_ECOLX_acquier_2013

1.0

1.0

0.8

8'0

0.8

0.6

0.6

0.6

0.4

0.4

0.2

0.2

0.2

0.0

0.0

6

10

10

4

6

10

BLAT_ECOLX_Stifler_2015

BRCA1_HUMAN_Findlay_2018

0.8

0.8

0.7

0.5

0.6

0.6

0.5

0.4

0.3

0.1

10

0.0

Q

6

10

00

10

CALM1_HUMAN_Weile_2017

CAPSD_AAV2S_Sinai_substitutions_2021

CCDB_ECOL_Adkar_2012

0.7

0.35

1.0

0.6

0E'0

0.8

0.25

0.20

0.6

0.15

0.4

0.10

0.2

0.05

0.2

0.1

0.00

0.0

4

6

8

10

2

4

6

8

10

2

6

8

10

CCDB_ECOL_Tripathi_2016

CP2C9_HUMAN_Amorosi_abundance_2021

CP2C9_HUMAN_Amorosi_activity_2021

0.7

0.7

.7

0.8

0.6

0.6

0.5

0.5

0.4

0.4

0.4

0.3

0.3

0.2

0.1

0.1

0.0

0.0

10

10

10

DLG4_HUMAN_Faure_2021

DLG4_RAT_McLaughlin_2012

DYR_ECOL_Thompson_plusLon_2019

1.0

0.7

0.8

0.6

0.8

0.5

0.6

0.6

0.4

0.3

0.4

0.4

0.2

0.1

0.2

0.2

00

0.0

0

4

10

10

0

8

10

ENV_HV1B9_DuenasDecamp_2016

ENV_HV1BR_Haddox_2016

ESTA_BACSU_Nutschel_2020

0.5

0.8

0.8

0.4

0.6

0.3

0.6

0.4

0.2

0.4

0.2

0.1 -

0.2

0.0

0.0

10

4

6

10

0

8

10

F7YBW8_MESOW_Aakre_2015

GAL4_YEAST_Kitzman_2015

GCN4_YEAST_Staller_induction_2018

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.4

0.4 -

0.4

0.2

0.2

0.2

0.0

0.0 L

2

4

6

10

2

10

0

4

10

GFP_AEQVI_Sarkisyan_2016

GRB2_HUMAN_Faure_2021

HIS7_YEAST_Pokusaeva_2019

0.175 -

1.0

0.014

0.150

0.8

0.012

0.125

0.010

0.100

0.6

8000

0.075

0.4 -

9000

0.050

000

0.025

0.2

200°0

000°0

4

6

10

0.0

10

0000

4

6

8

10

HSP82_YEAST_Flynn_2019

HSP82_YEAST_Mishra_2016

I6TAH8_168A0_Doud_2015

0.8

0.7

0.6

0.20

0.6

0.5

0.15

0.5

0.4

0.4

0.3

0.10

0.3

0.2

0.2

0.05

0.1

0.1

0.0

4

2

6

10

0

4

8

10

IF1_ECOLI_Kelsic_2016

KCNH2_HUMAN_Kozek_2020

KKA2_KLEPN_Melnikov_2014

1.0

1.0

0.8

0.8

0.8 -

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

10

2

10

0.0

0

6

8

10

MK01_HUMAN_Brenan_2016

MSH2_HUMAN_Jia_2020

MTH3_HAEAE_Rockah-Shmuel_2015

0.7

0E'0

1.0

0.6

0.25

0.8

0.5

0.20

0.4

0.6

0.15

0.4

0.10

0.2

0.1

0.05

0.2

0.0

0

2

10

0.00

2

4

6

10

0

6

10

NCAP_134A1_Doud_2015

NRAM_33A0_iang_standard_2016

NUD15_HUMAN_Suiter_2020

1.0

0.5

8'0

0.4

0.8

0.6

0.3

0.6

0.4

0.2

0.4

0.1

0.2

00

0'0

10

10

P53_HUMAN_Giacomelli_NULL_Etoposide_2018

P53_HUMAN_Giacomelli_NULL_Nutlin_2018

P53_HUMAN_Giacomelli_WT_Nutlin_2018

0.7

0.8

0.6

0.5

0.6

0.4

0.3

0.4

0.4

0.2

0.1

0°0

2

10

10

P53_HUMAN_Kotler_2018

P84126_THETH_Chan_2017

PA_134A1_Wu_2015

0.8

8°0

0.8

0.6

0.6

0.4

0.2

0.0

4

10

4

10

10

POLG_CXB3N_Mattenberger_2021

POLG_HCVJF_Qi_2014

1.0

0.5

0.8

0.6

0.3

0.4

0.1

0.2

10

6

10

6

10

PTEN_HUMAN_Matreyek_2021

PTEN_HUMAN_Mighel_2018

Q2N0S5_9HIV1_Haddox_2018

0.7

0.6

0.5

0.5

0.4

0.4

0.3

0.3

0.2

0.1

0.1

0.1

0.0

10

10

6

10

Q59976_STRSQ_Romero_2015

R1AB_SARS2_Flynn_growth_2022

RASH_HUMAN_Bandaru_2017

0.6

0.8

S°0

0.8

0.6

0.4

0.6

0.3

0.4

0.2

0.1

0.2

6

10

10

6

10

REV_HV1H2_Femandes_2016

RL401_YEAST_Mavor_2016

RL401_YEAST_Roscoe_2013

0.6

1.0

0.8

0.8

0.4

0.6

0.3

0.4

0.4

0.1

0.2

0.2

10

10

RL401_YEAST_Roscoe_2014

SC6A4_HUMAN_Young_2021

SCN5A_HUMAN_Glazer_2019

1.0

S°0

0.8

0.4

0.3

0.4

0.2

Z'C

0.1

0.2

2

4

8

10

0

4

6

10

2

6

10

SPG1_STRSG_Olson_2014

SPIKE_SARS2_Starr_bind_2020

SPIKE_SARS2_Star_expr_2020

0.05

0.8

0.8

0.04

0.6

0.6

0.03

0.02

0.4

0.4

0.01

20

0.2

0.00

SRC_HUMAN_Ahler_CD_2019

SUMO1_HUMAN_Weile_2017

SYUA_HUMAN_Newberry_2020

0.8

0.8

0.7

0.8

0.6

0.6

0.6

0.5

0.4

0.4

0.4

0.3

0.2

0.2

10

2

o

0.0

4

8

10

TADBP_HUMAN_Bolognesi_2019

TAT_HV1BR_Femandes_2016

TPK1_HUMAN_Weile_2017

1.0

0.8

0.8

0.7

0.7

0.6

0.6

0.6

0.5

0.5

0.4

0.4

E'0

0.3

0.2

0.2

0.2

0.1

0.0

b

10

10

4

10

TPMT_HUMAN_Matreyek_2018

TPOR_HUMAN_Bridgford_S505N_2020

TRPC_SACS2_Chan_2017

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

4

10

0.0

2

10

0

10

TRPC_THEMA_Chan_2017

UBC9_HUMAN_Weile_2017

UBE4B_MOUSE_Starita_2013

1.0

0.8

0.8

0.7

0.7

0.8

0.6

0.6

0.6

0.5 -

0.5

0.4

0.4

0.4

0.3

0.3

0.2

0.2

0.2

0.1 -

0.1

Z

4

6

8

10

0.0

2

4

6

10

10

VKOR1_HUMAN_Chiasson_abundance_2020

VKOR1_HUMAN_Chiasson_activity_2020

YAP1_HUMAN_Araya_2012

8'0

0.7

0.7

0.8

0.6

0.6

0.5

0.5

0.6

0.4

0.4

0.3

0.4

0.3

0.2

0.2

0.1

0.2

0.1

0.0

2

4

10

6

10

0.0

6

8

Q8WTC7_9CNID_Somermeyer_2022

D7PM05_CLYGR_Somermeyer_2022

Q6WV13_9MAX1_Somermeyer_2022

0.5

0.8

0.6

0.4

0.5

0.6

0.3

0.4

0.3

0.2

0.4

0.2

0.1

0.2

0.1

0.0

4

8

0.0

2

8

10

0.0

0

8

10

SPG1_STRSG_Wu_2016

CAPSD_BPMS2_Brauer_2022

PARE_ECOLX_Ding_2022

0.175

0.8

0.6

0.150

0.5

0.125

0.6

0.100

0.4

0.075

0.4

0.3

0.050

0.2

0.025

0.2

0.1

0.000

0

6

8

10

0.0

2

6

10

0

2

6

8

10

SPCAS9_Spencer_positive_2022

SPCAS9_Spencer_negative_2022

BAC7_Held_2022

0.30

0.25

0200

0.25

0.20

0.20

0.015

0.15

0.15 -

0.010

0.10

0.10

S000

0.05

0.05

000'0

10

10

ProteinNPT

Embeddings - Aug. (MSAT)

==- OHE - Aug. (MSAT)

==- OHE - Aug. (DeepSequence)

=-OHE

Zero-shot (MSAT)

==- Random

A0A247D711_LISMN_Stadelmann_2021

A0A220GHA5_9CAUD_Stadelmann_2021

0.8

0.8

0.6

0.6

0.4

0.4

0.2

0.2

0

2

4

6

8

10

0

2

4

6

8

10

HXK4_HUMAN_Gersing_2022

ANCSZ_Hobbs_2022

0.6

0.8

0.5

0.7

0.6

0.4

0.5

0.3

0.4

0.2

0.3

0.2

0.1

0.1

0.0

2

10

0.0

0

8

10

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduced ProteinNPT, which is a transformer based semi-supervised learning model. By combining MSA transformer and Non-parametric transformer(NPT), it outperformed their baselines on property and mutation effect predictions. Then they tested their model on the protein design task and showed its potential to help on designing novel protein sequences with desired labels. In addition, its contributions also include the new extended benchmark dataset, ProteinGym, as part of their test set to evaluate their model. This paper leveraged large quantities of unlabelled natural sequences to pretrain the model to have a more informative embedding. Moreover, the tri-axial attention helps the model to learn relationships not only within the row, but also across the data points in a batch. All these build up together to boost the performance of the model on multiple downstream tasks.\n\nSTRENGTHS:\nThe paper addresses an important topic which is the subject of much research. The flow of this paper is clear. It introduces the motivation first and then its solution to the problem fits the motivation nicely. In addition, the experimental design supports their claims by showing their model outperforming the baseline models in different downstrain tasks. Furthermore, they summarize their contributions and novelties clearly\nThe authors nicely demonstrate improvements in the quality of representation learning by comparing ProteinNPT with baselines methods on property prediction and mutation effect prediction tasks.  \nTo enhance the claim of advantages of tri-axial attention and auxiliary labels, the paper further shows that it can succeed in the protein fitness prediction problem. In this task, the authors provide some examples to illustrate their good performance while their evaluation metrics show that their model outperforms the other models in general. \nUsing ablation studies, the authors demonstrate that their embedding can help in different cross validation schemes compared to other methods.\nThe paper well built an experiment with ProteinNPT + Bayesian optimization methods and showed its ability to redesign protein with desired properties in an iterative manner.\n\nWEAKNESSES:\nThe paper did not illustrate the model design very clearly. The supplied figure is just too simplistic. The paper’s annotations and formulas give the audience a basic understanding of what the authors aim to do. However, it is hard to follow how the MSA transformer is introduced to the model. They should have more explicit illustration about how it is trained(i.e with what data) and where it is used as part of the ProteinNPT. In addition to that, it is vague about the training data for the MSA transformer in the baselines. It will be helpful for the audience to understand the experiment design if the paper can provide more details about how MSA sequences are generated(if used) and if the ProteinNPT used the same MSA transformer as the baseline.\n\nThe authors mention their novelties in different aspects. In fact, the dedicate a whole subsection for it, which helps detail the novelty. However, beyond the additional data and testing that come with it, most of the novelties appear to be subtle changes or directly borrowed ideas from MSA transformer or NPT. For example, the authors claim one of their novelties is a semi-supervised architecture but it is actually mentioned in NPT paper. \n\nIn the Figure 5 and 6, their model actually has similar performance or even lose to the DeepSequence and MSAT models in a lot of datasets. However, there is no analysis of the reason.\n\nMinor comments:\nLine 152, 167, 268: There is no “Fig. 4”.\nTable 6: the second “100” should be “1000”\nDefine \"MSA Transformer\" line 91,94 then switch to \"MSAT\" line 100\n\"labelled\" line 294\n\"a a sets\" line 94\n\"is that is\" line 178\n\"faction of test\" line 197\nDescription line 291-296 unclear\n\nReference in line 400 does not appear correctly\n\nQUESTIONS:\nsee above.\nAlso: Line 180: The idea of using informative other labels as features make sense. In reality they use MSAT predictions which makes their model like an ensamble of sorts. What happens if they take these labels out? How is performance affected?\n\nLIMITATIONS:\nsee above", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

comments: SUMMARY:
The paper researches for the purpose of iterative protein design. It first collects 13 additional fitness metrics to the ProteinGym benchmark. Then, it proposes ProteinNPT, an MSA-Transformer-based model for fitness prediction, and achieved state-of-the-art on ProteinGym. Finally, it applies the model to iterative protein redesign via Bayesian Optimization.

STRENGTHS:
The experimental results are promising.

WEAKNESSES:
1. Novelty: The way adapts Transformer to property prediction is straightforward. The novelty is limited.
2. Transferability: Additional training is required for different fitness settings. At the same time, the architecture seems to be difficult to even generalize to closely correlated fitness settings, for the labels are plugged into the input. The transferability is limited.
3. Real-world application: Labeled data is required for training, which would bring wet-lab costs. Additionally, the method is based on MSAs consisting of mutated sequences, which tend to be gathered at a certain point in the protein space rather than cover much space, so the method seems unsuitable to be applied in the discovery stage. As for the optimization stage, I doubt whether the precision of the proposed model satisfies the requirements from biology researchers. Overall, I think the space for real-world application is limited.
4. Other suggestions: Exchange "left" and "right" in the text description of Fig. 1, and declare “OHE” as One-Hot-Encoded in the main paper.

QUESTIONS:
N/A

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nhis paper introduces a non-parametric transformer variant called ProteinNPT, which is utilized for protein property prediction and design tasks. The study compares ProteinNPT with re-implemented top-performing baselines on the extended ProteinGym benchmark. \n\nThe objective of the paper is to address challenges faced in protein engineering, including the vast design space, sparse functional regions, and limited availability of labels. The results demonstrate that the proposed ProteinNPT outperforms all the compared methods in various protein property prediction tasks and also shows potentials in the protein design task.\n\nSTRENGTHS:\nThis paper introduces ProteinNPT, a non-parametric variant of transformers, which is employed for protein property prediction and design tasks. The authors extend the ProteinGym benchmark and re-implement several top-performing baselines for comparison with ProteinNPT. The primary focus of ProteinNPT is to tackle the challenge of limited availability of labels, and it can be effectively combined with state-of-the-art protein language models.\n\nThe problems investigated in this paper hold significant importance and appeal from a biological perspective. The authors have expanded the ProteinGym benchmark and conducted a comprehensive experimental analysis, with very detailed and sufficient results provided in both regular paper and Appendix. In conclusion, the proposed ProteinNPT offers substantial value and can contribute to the advancement of related research.\n\nWEAKNESSES:\nThe paper's presentation style may pose challenges for readers in terms of comprehension. It is important to note that this comment does not reflect the quality of the writing itself, but rather the accessibility and reader-friendliness of the content. In order to fully comprehend the paper, readers are required to go through the entirety of the appendix. In summary, it is good for the paper to be as self-contained as possible. For instance, \"the difference with axial transformer, MSA Transformer and Non-parametric transformers\" would be better suited within the main body of the paper rather than relegating it to the Appendix. In general, it is crucial to explicitly state the differences between the proposed method and other comparable approaches when the leaving space of paper is enough.\n\nFurthermore, it would be beneficial to include additional mathematical formulations pertaining to ProteinNPT in Section 4, rather than solely relying on the architecture diagram for explanation.\n\nQUESTIONS:\nI am interested in the part of model training (Section 4.2 and Appendix B.3). It would be beneficial if the authors could provide a detailed description of the multi-task optimization process. For instance, could you explain how the annealing process is implemented for the token prediction objective? Additionally, could you elaborate on the use of a cosine schedule throughout the training process?\n\nMinor Miscellaneous Suggestions\n\n- Line 152 and 268: 'Fig 2', not 'Fig 4'?\n- Table-1: The full name of abbreviation \"OHE\" needs to be added.\n\nLIMITATIONS:\nNo obvious limitations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose to use non parametric transformers to model the ‘fitness’ landscape and the protein primary sequences with a Bert-like model. Both tokens and continuous attributes can be masked in the modeling. Additionally, the authors introduce several novel datasets to test and benchmark their model. Finally, a directed-evolution like experiment is conducted to evaluate the capability of the proposed model to be used to assist the design.\n\nSTRENGTHS:\n1- The authors’ problem is well motivated and clearly presented. The method is sound and particularly adapted to the tackled problem hence necessitating the conducted investigation.\n\n2- I found the discussion and the different splits for training the model very interesting and provides additional insight on how the model behaves given novel mutations.\n\n3- Adding novel datasets to the field in a standardized way is also in my opinion a great contribution to the bio-ML community. Nonetheless, given that there is room for additional details, I am disappointed about the lack of description of the novel datasets in the main part of the paper.\n\n4 - The results seem rather conclusive that the method is well suited to adress the learning of the joint distribution despite several more investigations could be conducted for confirmation.\n\nWEAKNESSES:\n1 -  A notable exception in my opinion to the clarity of the paper is section 6 that i struggled to understand.\n\n2 - There is in my opinion little discussion on how the inputs are constructed. For instance, how many proteins are included in the “alignement”. Are there any learning tradeoffs there ? For example: number of doable gradient steps before overfitting  versus  available able information in the alignment ? \n\n3 - I found it hard to guess which baseline was what in Table 1.  Moreover, the results in table 1 does not present any variance in the experiments.\n\n4 - Since I found the presented work interesting, I was interested in what problem could be addressed by such a model and the expected \nlimits (minimum number of datapoints, minimum number of datapoints in the alignments, gradient steps and so on) ?\n\n5 - A lot of paper in the ML for protein design literature uses datasets such as GFP. Can the authors comment on their choice of not using this standard dataset.\n\nQUESTIONS:\n1-  How would the model behave if at inference, we specify the desired target and mask (randomly or any other fashion) some tokens ? \n\n2 - To follow up the previous question, since your model learns a joint probability distribution, can you use this property for protein design ? I don’t believe that this is what is done in your section 6, but devising to an experiment like this would prove the ability to generate tokens conditionally to the output value. \n\n3 - What is the influence of the number of rows in the input sequence ? \n\n4 - Given that the current version does not exceed the 9 page limit, I strongly encourage the authors to provide either some hyper-parameterisation or some explanation on the dataset.\n\n5 - I also found unclear in the main paper whether the authors were responsible for the experiments that led to the new datasets. \n\n6 - How does the model scale with the number of data points ? \n\n\nNote that i am willing to change my evlulation if my questions are addressed.\n\nLIMITATIONS:\nSee Weaknesses & Questions\n\nMinor typo: line 96", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-08"}