B ENCHMARKING U NCERTAINTY Q UANTIFICATION FOR P ROTEIN E NGINEERING
A BSTRACT
Machine learning sequence-function models for proteins could enable significant advances in protein engineering, especially when paired with state-of-theart methods to select new sequences for property optimization and/or model improvement.
Such methods (Bayesian optimization and active learning) require calibrated estimations of model uncertainty.
While studies have benchmarked a variety of deep learning uncertainty quantification (UQ) methods on standard and molecular machine-learning datasets, it is not clear how well these results extend to protein datasets.
In this work, we implement a panel of deep learning UQ methods on the Fitness Landscape Inference for Proteins (FLIP) benchmark regression tasks.
We compare results across different degrees of distributional shift using metrics that assess each UQ method’s accuracy, calibration, coverage, width and rank correlation to provide recommendations for the effective design of biological sequences.
I NTRODUCTION
Machine learning (ML) has already begun to accelerate the field of protein engineering by providing low-cost predictions of phenomena that require time- and resource-intensive labeling by experiments or physics-based simulations (Yang et al., 2019).
It is often necessary to have an estimate of model uncertainty in addition to the property prediction, as the performance of an ML model can be highly dependent on the domain shift between its training and testing data (Kendall & Gal, 2017).
Because protein engineering data is often collected in a biased manner (Dallago et al., 2021), tailored ML methods are required to guide the selection of new experiments from a protein landscape.
Uncertainty quantification (UQ) can inform the selection of experiments in order to improve a ML model or optimize protein function through active learning or Bayesian optimization.
In chemistry and materials science, several studies have benchmarked common UQ methods against one another on standard datasets and have used or developed appropriate metrics to quantify the quality of these uncertainty estimates (Scalia et al., 2020; Tran et al., 2020; Hirschfeld et al., 2020;
Nigam et al., 2021; Soleimany et al., 2021).
This work has illustrated that the best choice of UQ method can be dependent on the dataset and other considerations such as scaling.
While some protein engineering work has leveraged uncertainty estimates, these studies have been limited to single UQ methods such as ensembles (Mariet et al., 2020) or Gaussian processes (GPs)
(Hie et al., 2020).
In this work, we use a group of standardized, public protein datasets to benchmark a panel of UQ methods.
Our chosen datasets include splits with varied degrees of domain extrapolation, enabling us to evaluate the methods in a setting similar to what might be experienced while collecting new experimental data.
We assess each model using a variety of metrics that capture different aspects of desired performance, including accuracy, calibration, coverage, width, and rank correlation.
M ETHODS
This work uses several dataset splits to evaluate the performance of a panel of uncertainty methods across various evaluation metrics (Figure 1). Uncertainty Function prediction [Figure]
one
[Figure] hot encoding [Figure] 2. D ATASETS AND T ASKS
The landscapes used in this work were taken from the Fitness Landscape Inference for Proteins (FLIP) benchmark (Dallago et al., 2021).
These included the binding domain of an immunoglobulin binding protein (GB1), adeno-associated virus stability (AAV), and thermostability (Meltome), which cover a large sequence space and a broad range of protein families.
The FLIP benchmark includes several train-test splits, or tasks, for each landscape; most of these tasks are designed to mimic common, real-world data collection scenarios and are thus a more realistic assessment of generalizability than random train-test splits.
However, random splits are also included as a point of reference.
We chose 8 of the 15 FLIP tasks to benchmark the panel of uncertainty methods.
We selected these tasks to be representative of several regimes of domain shift; they include random sampling with no domain shift (AAV/sampled, Meltome/mixed, and GB1/sampled); the highest (and most relevant) domain-shift regimes (AAV/Mut-Des, AAV/7 vs. Many, and GB1/1 vs. Rest); and less aggressive domain shifts (GB1/2 vs. Rest and GB1/3 vs. Rest).
2. E VALUATION M ETRICS
To give a comprehensive report of model accuracy, we compute the following metrics on the test sets: root mean square error (RMSE), mean absolute error (MAE), coefficient of determination ( R 2 ), and Spearman rank correlation ( ρ ).
RMSE is more sensitive to outliers than MAE, so while both are informative independently, the combination of the two gives additional information about the distribution of errors.
R 2 and ρ are both unitless and are thus more easily interpreted and compared across datasets.
We evaluated the quality of the uncertainty estimates using four metrics.
First, ρ unc is the the
Spearman rank correlation between uncertainty and absolute prediction error.
Following Kompa et al.
(2021), we measure the coverage as the percentage of true values that fall within the 95% confidence interval ( ± 2 σ ) of each prediction.
Kompa et al.
(2021) define the width as the size of the 95% confidence region ( 4 σ ), but we calculate the width relative to the range of the training set as 4 σ/R to make these values more interpretable across datasets.
Finally, the miscalibration area (also called the area under the calibration error curve or AUCE) quantifies the absolute difference between the calibration plot and perfect calibration in a single number (Gustafsson et al., 2020).
2. U NCERTAINTY M ETHODS
We implemented eight uncertainty methods for this benchmark: linear Bayesian ridge regression (BRR), Gaussian processes (GPs), and six methods using variations on a convolutional neural network (CNN) architecture.
We used a scikit-learn implementation of BRR (Pedregosa et al., 2011) and a GPyTorch implementation of continuous-kernel GPs (Gardner et al., 2018).
The CNN implementation from FLIP provided the core architecture used by our dropout (Gal & Ghahramani, 2016), ensemble (Lakshminarayanan et al., 2017), evidential (Amini et al., 2020), mean-variance estimation (MVE) (Nix & Weigend, 1994), and last-layer stochastic variational inference (SVI) (Hoffman et al., 2013) models.
For each dropout model, we tested dropout fractions of 0.1, 0.2, 0.3, 0.4, and 0.5 and reported the model with the lowest miscalibration area.
The output of the evidential uncertainty can be divided into an epistemic and aleatoric component following the analysis of Amini et al.
, so we report these uncertainties separately in Tables 1-4 along with their sum as the total uncertainty in Figures 2-4.
For all models and landscapes, the sequences were featurized as one-hot encodings.
R ESULTS AND D ISCUSSION
We trained each of the seven models described in Section 2.3 on each of the eight tasks described in Section 2.1 and evaluated their performance on the test set using the metrics described in Section 2.2.
We compare model calibration and accuracy in Figure 2 and the percent coverage vs. average width relative to range in Figure 3.
As expected, the splits with the least required domain extrapolation tend to have more accurate models (lower RMSE).
However, the relationship between miscalibration area and extrapolation is less clear; some models were highly calibrated on the most rigorous (highest domain shift) splits, while others were poorly calibrated even on random splits.
There is no single method that performs consistently well across splits and landscapes, but some trends can be observed.
For example, the CNN ensemble is often one of the highest accuracy models, but also one of the most poorly calibrated.
Figure 3 illustrates that many methods perform relatively well in either coverage or width (corresponding to the the top and left limits of the plot, respectively), but few methods perform well in both.
Among those that do, the GP model is the only method that is among the best across all landscapes and most splits.
Similarly to Figure 2, there is some observable trend that more challenging splits are further from the optimal part (upper left) of the plot; this trend is more clear for the GB splits than for the AAV splits.
Figure 4 compares the ranking performance of each method in terms of predictions relative to true values and uncertainty estimates relative to true errors.
The splits are ordered according to domain shift within their respective landscapes, and the rank correlation of the predictions to the true labels generally decreases moving from left to right within a landscape.
Ensembling was often among the best performance in ρ , and the GP and BRR models performed better on ρ on the splits with the highest domain shift.
Performance on ρ unc is generally much worse than that on ρ , with a large number of results showing negative correlation.
Dropout has one of the highest and most consistent ρ unc across splits.
All methods have ρ unc near zero for the most challenging splits.
MVE performs particularly poorly in regimes of high domain shift, which is consistent with its intended use as an estimator of aleatoric (data-dependent) uncertainty.
C ONCLUSION
Calibrated uncertainty estimations are necessary for effective property optimization or model improvement using Bayesian optimization or active learning.
In this work, we have benchmarked a panel of uncertainty quantification (UQ) methods on protein datasets, including on train-test splits that are representative of real-world data collection practices.
After evaluating each method based on accuracy, calibration, coverage, width, and rank correlation, there is no method that performs consistently well across all metrics or all landscapes and splits.
However, this study only examines models using one-hot-encoding representations of sequences; further research is needed to understand how more informative and generalizable representations such as the pretrained embeddings from (Dallago et al., 2021) may impact performance, particularly on splits with higher domain shifts.
Additionally, while the evaluation metrics used herein provide valuable information, they are ultimately only a proxy for expected performance in Bayesian optimization and active learning.
Retrospective studies using holdouts of the training sets would be useful for evaluating this performance more directly.
A more thorough understanding of how to best apply UQ to sequence-function models will ultimately enable more effective protein engineering. [Figure]
A UTHOR C ONTRIBUTIONS
[Figure] P.S. and K.K.Y. conceived the project.
K. P. G. wrote the computer code, analyzed the data, and wrote the first manuscript draft.
A.P.S. and K.K.Y. supervised the research and edited the manuscript.
A CKNOWLEDGMENTS
K.P.G. was supported by a Microsoft Research micro-internship and by the National Science Foundation Graduate Research Fellowship Program under Grant No. 1745302.
We acknowledge the
MIT Lincoln Laboratory Supercloud cluster (Reuther et al., 2018) at the Massachusetts Green High
[Figure] Performance Computing Center (MGHPCC) for providing high-performance computing resources to train our machine learning models. [Figure]
A A PPENDIX
[Table] [Table]
AAV (7-vs-many) 0.
AAV (mut-des) 0.
Meltome (mixed) 0.
GB1 (sampled) 0.
GB1 (3-vs-rest) 0. GB1 (2-vs-rest) 0.
GB1 (1-vs-rest) 0.
[Table] [Table]