SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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.
Model name	Spearman (↑)	MSE (↓)
Contig.	Mod.	Rand.	Avg.	Contig.	Mod.	Rand.	Avg.
OHE	0.08	0.02	0.54	0.21	1.17	1.11	0.92	1.06
OHE - Aug. (DS)	0.40	0.39	0.48	0.42	0.98	0.93	0.78	0.90
OHE - Aug. (MSAT)	0.41	0.40	0.50	0.44	0.97	0.92	0.77	0.89
Embed. - Aug. (MSAT)	0.47	0.49	0.57	0.51	0.93	0.85	0.67	0.82
ProteinNPT	0.48	0.51	0.65	0.54	0.93	0.83	0.53	0.77
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.
Figure 2: Multiples mutants performance. (Left) Spearman's rank correlation between model predictions and experimental measurements, for each assay in ProteinGym with multiple mutants (see Appendix A.1). (Right) Average Spearman's rank correlation, overall and by 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:
$$
a (x; \lambda) = \mu (x) + \lambda . \sigma (x) \tag {6}
$$
where  $\mu$  is the predicted fitness value,  $\sigma$  quantifies uncertainty in the prediction, and  $\lambda$  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.
Algorithm 1 Iterative protein redesign
1:	Input: Initial labeled data DL; Initial unlabeled data DU;
Batch size B; Batch set S = ∅; Acquisition function
α(x;λ)
2:	for t ∈ 1,2,...,10 do
3:	Train model model on DL
4:	for t ∈ 1,2,...,B do
5:	xnew = arg maxx∈DL α(x;λ)
6:	Obtain label ynew for xnew
7:	S ← S ∪ {(xnew, ynew)}
8:	DU ← DUS{xnew}
9:	end for
10:	DL ← DL ∪ S, S ← ∅
11:	end for
12:	Output: DL, Trained model on DL
Figure 3: In silico protein redesign. (Left) Iterative redesign algorithm (Right) Recall rate of high fitness points (top 3 deciles) vs acquisition cycle, averaged across all DMS assays in ProteinGym.
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.
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.
Figure 4: Conditional sampling. ProteinNPT is used to sample novel sequences for the GFP protein conditioned on high fitness values, leading to sequences with high predicted fitness relative to controls.
Figure 5: Row-wise attention map. Row-wise attention between residues & fitness, mapped on the DHFR enzyme structure. The high intensity residue (red) corresponds to a substrate binding site.
Table 2: Impact of column-wise attention on performance. (Left) Impact of ablating the column-wise attention in ProteinNPT layers. (Right) Impact of the number of labeled sequences used at inference for a model trained with column-wise attention. Performance is measured by the Spearman's rank correlation between model predictions and experimental measurements.
CV scheme	Column attention	Nb. of labeled sequences at inference
Without	With	0	100	200	500	1000	2000
Random	0.669	0.684	0.398	0.677	0.678	0.679	0.684	0.685
Modulo	0.530	0.531	0.299	0.533	0.531	0.531	0.531	0.531
Contiguous	0.425	0.501	0.254	0.496	0.504	0.502	0.501	0.500
Average	0.542	0.572	0.317	0.569	0.571	0.571	0.572	0.572
Column-wise attention is critical to reach top performance. We test out the impact of column-attention 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).
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.
P.N. is supported by GSK and the UK Engineering and Physical Sciences Research Council (ESPRC ICASE award no.18000077). R.W. is supported by the UK Engineering and Physical Sciences Research Council, as part of the CDT in Health Data Science. D.S.M. is supported by a Chan Zuckerberg Initiative Award (Neurodegeneration Challenge Network, CZI2018-191853) and a NIH Transformational Research Award (TR01 1R01CA260415). Y.G. holds a Turing AI Fellowship (Phase 1) at the Alan Turing Institute, which is supported by EPSRC grant reference V030302/1. We thank Jannik Kossen for helpful feedback and discussion.
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.
Table 3: DMS assays added to ProteinGym to for the extended benchmark
Uniprot ID	Reference	Selection assay	Function grouping
Q8WTC7 9CNID	Gonzalez et al. [2022]	Green fluorescence intensity	Activity
D7PM05 CLYGR	Gonzalez et al. [2022]	Green fluorescence intensity	Activity
Q6WV13 9MAXI	Gonzalez et al. [2022]	Green fluorescence intensity	Activity
SPG1 STRSG	Wu et al. [2016a]	Binding affinity to IgG	Binding
CAPSD BPMS2	Brauer et al. [2021]	Quantitative capsid self-assembly	Activity
PARE ECOLX	Ding et al. [2022]	Growth under toxin	Activity
CAS9 STRP1	Spencer and Zhang [2017]	Expression of Cas9 and correct cleavage	Activity
CAS9 STRP1	Spencer and Zhang [2017]	Expression of Cas9 and non-cleavage of off target sequence	Activity
CTHL3 BOVIN	Koch et al. [2022]	Antimicrobial activity in minimal inhibitory concentration	Activity
A0A247D711 LISMN	Stadelmann et al. [2021]	Activity against SpyCas9 inducing an RFP reporter	Activity
A0A220GHA5 9CAUD	Stadelmann et al. [2021]	Activity against SpyCas9 inducing an RFP reporter	Activity
HXK4 HUMAN	Gersing et al. [2022]	Functional complementation to reduced growth on glucose medium	Organismal Fitness
ANCSZ	Hobbs et al. [2022]	Successful phosphorylation of bait peptide	Activity
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 MESOW Aakre 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.
We also 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.