SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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.
${}^{2}$  Depending on the underlying pre-trained protein language models used,additional beginning of sequence (BOS) and end of sequence (EOS) tokens are also added in practice. For instance,in the case of the MSA Transformer,a BOS token is pretended to the beginning of the input sequence. The size of the embedded sequence is thus  $B.\left( {L + 1}\right) .D$
Table 4: ProteinNPT - Architecture details
Hyperparameter	Value
Nb. ProteinNPT layers	5
Embedding dimension (d)	200
Feedforward embedding dimension	400
Nb. attention heads	4
CNN kernel size (post embedding)	7
Weight decay	5.10-3
Dropout	0.0
Table 5: ProteinNPT - Training details
Hyperparameter	Value
Training steps	10k
Learning rate warmup steps	100
Peak learning rate	3.10-4
Optimizer	AdamW
Gradient clipping norm	1.0
Learning rate schedule	Cosine
Training batch (masked)	64
Training batch (unmasked)	361
Table 6: ProteinNPT ablations - Use of auxiliary labels Spearman's rank correlation between model predictions and experimental measurements. The auxiliary label used here is the zero-shot fitness prediction with the MSA Transformer.
CV scheme	No aux. label	With aux. label
Random	0.683	0.684
Modulo	0.527	0.531
Contiguous	0.439	0.501
Average	0.549	0.572
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.
Table 7: ProteinNPT ablations - Protein sequence embeddings Spearman's rank correlation between model predictions and experimental measurements, when leveraging different protein language models (ESM-1v, Tranception or MSA Transformer) to obtain sequence embeddings and using as auxiliary labels the zero-shot predictions from the corresponding protein language model. 'Learned' corresponds to the case of learned token embeddings from scratch and no auxiliary labels.
CV scheme	Learned	Trancection	ESM1v	MSAT
Random	0.422	0.665	0.659	0.684
Modulo	0.098	0.433	0.468	0.531
Contiguous	-0.048	0.375	0.380	0.501
Average	0.141	0.491	0.502	0.572
Table 8: ProteinNPT ablations - Impact of the size of the training set Average Spearman's rank correlation between predictions and DMS labels. Assays are grouped based on 'Small size' (≤2k labels), 'Medium size' (2-8k labels) and 'Large size' (>8k labels).
CV scheme	Small	Medium	Large
Random	0.633	0.738	0.573
Modulo	0.492	0.591	0.421
Contiguous	0.481	0.546	0.381
Average	0.535	0.625	0.458
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  $\S 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 = 1k$  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], Trancection [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, Stark 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'.
We carry 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.
Table 9: Embeddings Augmented - Architecture details
Hyperparameter	Value
Embedding dimension	768
CNN kernel size (post embedding)	9
Weight decay (fitness pred.)	10-8
Weight decay (others)	5.10-2
Dropout	0.1
Max context length	1024
Table 10: Embeddings Augmented - Training details
Hyperparameter	Value
Training steps	10k
Learning rate warmup steps	100
Peak learning rate	3.10-4
Optimizer	AdamW
Gradient clipping norm	1.0
Learning rate schedule	Cosine
Training batch	64
Table 11: Baselines ablations - Augmented OHE Spearman's rank correlation between model predictions and experimental measurements.
CV scheme	No Aug.	DS	ESM1v	MSAT	Trancection	TrancaptEVE
Random	0.572	0.510	0.508	0.494	0.499	0.521
Modulo	0.033	0.403	0.405	0.396	0.390	0.423
Contiguous	0.105	0.411	0.418	0.402	0.396	0.434
Average	0.237	0.442	0.444	0.431	0.428	0.459
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.
Table 12: Baselines ablations - Protein sequence embeddings Spearman's rank correlation between model predictions and experimental measurements.
CV scheme	Trancection	ESM1v	MSAT
Random	0.646	0.633	0.624
Modulo	0.498	0.534	0.553
Contiguous	0.413	0.487	0.509
Average	0.519	0.551	0.562
Table 13: Baselines ablations - Non-linear transform post embedding Spearman's rank correlation between model predictions and experimental measurements. We perform this ablation on the Random CV scheme only.
Transform	Spearman
Light attention	0.405
Linear	0.480
MLP	0.509
ConvBERT	0.617
CNN	0.624
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 mean-pooled 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 cross-validation 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
Table 14: Fitness prediction performance. We compute the Spearman's rank correlation between model predictions and experimental measurements, averaged across the 100 DMS assays from the extended ProteinGym benchmark. DS and MSAT are shorthand for DeepSequence and MSA Transformer respectively. The Spearman for the MSA Transformer in the zero-shot setting is 0.41 (independent of CV split).
Model name	Contiguous	Modulo	Random	Average
OHE	0.08 (0.02)	0.02 (0.02)	0.54 (0.01)	0.21 (0.02)
OHE - Augmented (DS)	0.40 (0.01)	0.39 (0.01)	0.48 (0.01)	0.42 (0.01)
OHE - Augmented (MSAT)	0.41 (0.01)	0.40 (0.01)	0.50 (0.01)	0.44 (0.01)
Embeddings - Augmented (MSAT)	0.47 (0.01)	0.49 (0.01)	0.57 (0.01)	0.51 (0.01)
ProteinNPT	0.48 (0.00)	0.51 (0.00)	0.65 (0.00)	0.54 (0.00)
Table 15: Multiple mutants fitness prediction - Performance summary. We report the aggregated performance of ProteinNPT and other baselines listed in Appendix D.1, measured by the Spearman's rank correlation  $\rho$  and Mean Squared Error (MSE) between model scores and experimental measurements. The standard error of the average performance is reported in parentheses.
Model name	Spearman ↑	MSE ↓
Zero-shot (MSAT)	0.367 (0.002)	-
OHE	0.717 (0.002)	0.541 (0.008)
OHE - Augmented (DeepSequence)	0.719 (0.001)	0.658 (0.106)
OHE - Augmented (MSAT)	0.722 (0.001)	0.516 (0.008)
Embeddings - Augmented (MSAT)	0.695 (0.002)	0.479 (0.008)
ProteinNPT	0.788 (0.002)	0.259 (0.001)
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.
Figure 6: Single mutants fitness prediction - Random cross-validation scheme We report the DMS-level performance (measured by the Spearman's rank correlation  $\rho$  between model scores and experimental measurements) of ProteinNPT and other baselines listed in Appendix D.1.
Figure 7: Single mutants fitness prediction - Modulo cross-validation scheme We report the DMS-level performance (measured by the Spearman's rank correlation  $\rho$  between model scores and experimental measurements) of ProteinNPT and other baselines listed in Appendix D.1.
Figure 8: Single mutants fitness prediction - Contiguous cross-validation scheme We report the DMS-level performance (measured by the Spearman's rank correlation  $\rho$  between model scores and experimental measurements) of ProteinNPT and other baselines listed in Appendix D.1.