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.

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.

Figure 9:. Multiple target prediction

Avg. Spearman’s rank correlation between model predictions and experimental measurements, for proteins with 2 or 3 distinct experimental measurements predicted simultaneously (respectively top and bottom plots).

Figure 10:. Uncertainty calibration curves for each cross-validation scheme.

Uncertainty calibration curves plot a performance metric of interest, here MSE (y-axis), as a function of the proportion of points set aside based on their uncertainty (x-axis) (from right to left, we set aside an increasing fraction of the most uncertain points). To be properly calibrated, an uncertainty quantification metric should monotonically improve as we set aside an increasing proportion of the most uncertain points. We experiment with three different uncertainty quantification schemes: MC dropout, Batch resampling, and a hybrid scheme. For a fixed compute budget, the hybrid scheme delivers optimal performance across our three cross-validation schemes.

Figure 11:. DMS-level performance for iterative protein redesign experiments (Assays 1–24).

We plot the recall rate of high fitness points (top 3 deciles) as the function of the number of batch acquisition cycles. The shaded regions represent the standard errors of each method.

Figure 12:. DMS-level performance for iterative protein redesign experiments (Assays 25–48).

We plot the recall rate of high fitness points (top 3 deciles) as the function of the number of batch acquisition cycles. The shaded regions represent the standard errors of each method.

Figure 13:. DMS-level performance for iterative protein redesign experiments (Assays 49–72).

We plot the recall rate of high fitness points (top 3 deciles) as the function of the number of batch acquisition cycles. The shaded regions represent the standard errors of each method.

Figure 14:. DMS-level performance for iterative protein redesign experiments (Assays 73–96).

We plot the recall rate of high fitness points (top 3 deciles) as the function of the number of batch acquisition cycles. The shaded regions represent the standard errors of each method.

Figure 15:. DMS-level performance for iterative protein redesign experiments (Assays 97–100).

We plot the recall rate of high fitness points (top 3 deciles) as the function of the number of batch acquisition cycles. The shaded regions represent the standard errors of each method.

Figure 1:. ProteinNPT architecture.

(Left) The model takes as input the primary structure of a batch of proteins of length $L_{seq}$ along with the corresponding $L_t$ labels and, optionally, $L_a$ auxiliary labels (for simplicity we consider $L_t=L_a=1$ here). Each input is embedded separately, then all resulting embeddings are concatenated into a single tensor. Several ProteinNPT layers are subsequently applied to learn a representation of the entire batch, which is ultimately used to predict both masked tokens and targets (depicted by question marks). (Right) A ProteinNPT layer alternates between tied row and column attention to learn rich embeddings of the labeled batch.

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.

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.

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.