Our first goal was to evaluate the calibration and quality of a variety of UQ methods. We implemented seven uncertainty methods for this benchmark: linear Bayesian ridge regression (BRR) [16, 17], Gaussian processes (GPs) [18], and five methods using variations on a convolutional neural network (CNN) architecture. The CNN implementation from FLIP [3] provided the core architecture used by our dropout [19], ensemble [20], evidential [21], mean-variance estimation (MVE) [22], and last-layer stochastic variational inference (SVI) [23] methods. Additional model details are provided in the Methods section.

The landscapes used in this work were taken from the Fitness Landscape Inference for Proteins (FLIP) benchmark [3]. These include the binding domain of an immunoglobulin binding protein (GB1), adeno-associated virus stability (AAV), and thermostability (Meltome) data landscapes, 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 generalization 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—random sampling with no domain shift (AAV/Random, Meltome/Random, and GB1/Random); the highest (and most relevant) domain-shift regimes (AAV/Random vs. Designed and GB1/1 vs. Rest); and less aggressive domain shifts (AAV/7 vs. Rest, GB1/2 vs. Rest, and GB1/3 vs. Rest). The Datasets section of the Methods provides notes on the nomenclature used for these tasks.

We trained the seven models on each of the eight tasks described above and evaluated their performance on the test set using the metrics described in the Evaluation Metrics section. We compare model calibration and accuracy in Fig 2 and the percent coverage versus average width relative to range in Fig 3. These figures illustrate the results for models trained on the embeddings from a pretrained ESM language model [15]; the corresponding results using one-hot encodings are shown in Figs A and B in S1 Appendix.

As expected, the splits with the least required domain extrapolation tend to have more accurate models (lower RMSE; Fig 2). However, the relationship between miscalibration area and extrapolation is less clear; some models are highly calibrated on the most difficult (highest domain shift) splits, while others are 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, ensembling is often one of the highest accuracy CNN models, but also one of the most poorly calibrated. Additionally, GP and BRR models are often better calibrated than CNN models. For the AAV and GB1 landscapes (Fig 2a and 2c), model miscalibration area usually increases slightly while RMSE increases more substantially with increasing domain shift.

In addition to accuracy and calibration, we assessed each method in terms of the coverage and width of its uncertainty estimates. A good uncertainty method results in high coverage (ideally, the true value falls within the 95% confidence region 95% of the time) while still maintaining a small average width. The latter is necessary because predicting a very large and uniform value of uncertainty for every point would result in good coverage, so coverage alone is not sufficient. Fig 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. Similarly to Fig 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 GB1 splits (Fig 3b) than for the AAV splits. Most models trained on the AAV landscape (Fig 3a) have a similar average width/range ratio for all splits, but for the GB1 landscape (Fig 3c), this ratio typically increases as the domain shift increases. The locations of the sets of points for each model type shared some similarities across landscapes. CNN SVI often has low coverage and low width, CNN MVE often has moderate coverage and moderate width, and CNN Evidential and BRR often have high coverage and high width. These trends across landscapes could point to a general problem of under- or over-confidence with some model types, and indicates that post-hoc calibration may be necessary. The results for all prediction and uncertainty metrics (along with their standard deviations across 5 different initialization seeds) are shown in Tables A to AR in S1 Appendix.

We next assessed how target predictions and uncertainty estimates depended on the degree of domain shift. Across datasets and splits, we compared the ranking performance of each method in terms of predictions relative to true values and uncertainty estimates relative to true errors (ESM in Fig 4 and one-hot encodings (OHE) in Fig C in S1 Appendix). The splits are ordered according to domain shift within their respective landscapes (lowest to highest shift from left to right). We observe that the rank correlation of the predictions to the true labels generally decreases moving from less to more domain shift within a landscape, consistent with expectation, with the exception of AAV/Random vs. Designed models performing better than AAV/7 vs. Rest models (Fig 4a). Most methods exhibit similar performance in Spearman rank correlations of predictions to targets (ρ) within the same task. For many tasks, GP and BRR models perform as well or better than CNN models. Performance on Spearman rank correlations of uncertainties to prediction residuals (ρ_unc) is generally much worse than that on ρ, with some results showing negative correlation (Fig 4b). MVE and evidential uncertainty methods are most performant in ρ_unc for most cases of low to moderate domain shift. Most methods have ρ_unc near zero for the most challenging splits. Despite the relatively good performance of MVE on tasks with low to moderate domain shift, it performs poorly in cases of high domain shift, which is consistent with its intended use as an estimator of aleatoric (data-dependent) uncertainty.

We find that the models trained on ESM embeddings outperform those trained on one-hot encodings in 21 out of 51 cases for rank correlation of test set predictions, and 29 out of 51 cases for rank correlation of test set uncertainties. The relative performance of the two representations on prediction and uncertainty rank correlation is shown in Fig D in S1 Appendix. In terms of predictions, ESM embeddings often yield substantially better performance for tasks with high domain shift (e.g. GB1/1 vs. Rest and Meltome/Random), while OHE performs slightly better on tasks with lower domain shift (e.g. AAV/Random and GB1/3 vs. Rest). The relative uncertainty rank correlation performance, on the other hand, does not have a clear relationship to domain shift.

Since there is no single best UQ method across datasets, splits, and metrics, it is prudent for practitioners to quantify the performance of uncertainty estimates on each new task and to prioritize metrics according to the situation (e.g. prioritize high coverage over low width in high-risk or safety-critical situations).

In protein engineering, the purpose of uncertainty estimation is typically to intelligently prioritize sample acquisition to facilitate downstream experimentation. One such use case is in active learning, where uncertainty estimates are used to inform sampling with the goal of improving model predictions overall (i.e., to achieve an accurate model with less training data; Fig 5a). Having assessed the calibration and accuracy of the panel of UQ methods above, we next evaluated whether uncertainty-based active learning could make the learning process more sample-efficient. Across all datasets and splits using the pretrained ESM embeddings, data acquisition was simulated as iterative selection from the data library according to a given sampling strategy (acquisition function; see Methods for details). The results are summarized in Fig 5 for Spearman rank correlation (ρ) on three methods and one split per landscape, and additional results are shown in the Figs H to BH in S1 Appendix for other metrics, uncertainty methods, and splits. Across most models, the performance difference between the start of active learning (10% of training data) and end of active learning (100% of training data) is relatively small, and many models begin to plateau in performance before reaching 100% of training data. In addition to the active learning experiments run with 10% of the training data in the initial sample, we also show results starting with 1% and 5% of the training data in Figs F and G in S1 Appendix, respectively. These results showed worse initial performance given the smaller initial training data, but otherwise similar trends to the trials starting at 10%.

The “explorative greedy” and “explorative sample” acquisition functions (which sample based on uncertainty alone or sample randomly weighted by uncertainty, respectively) sometimes outperform random sampling, but this is not true across all methods and landscapes (Fig 5b–5d). In some cases, the performance of the uncertainty-based sampling strategies also varies depending on the fraction of the total training data available to the model. For example, for the Meltome/Random split and CNN evidential model (Fig 5c), explorative greedy sampling results in a decrease in model performance after the first round of active learning while the explorative sample strategy increases performance. By the fourth round of active learning for this task, the two explorative strategies outperform random sampling. This indicates that in the early stages of active learning when a model’s uncertainty estimates are poorly calibrated, it may be advantageous to sample with at least some randomness included in an uncertainty-based acquisition function. We also analyzed how the mean test set uncertainty changed as more data was acquired during active learning. Fig E in S1 Appendix illustrates that in some cases, the mean test set uncertainty decreased with increasing training data, while in other cases, it increased. While one might expect adding more data from uncertain sequences would always cause a decrease in the mean uncertainty, this assumes that (1) the added training data comes from the same distribution as the test data, and (2) the uncertainty estimates are well-calibrated at the beginning and remain so after retraining with more data. Overall, the results indicate that uncertainty-informed active learning can outperform random sampling and thus lead to more accurate machine learning models with fewer training points needing to be measured (Fig 5b–5d).

Uncertainty estimates can also be leveraged to identify top-performing sequences with as few samples as possible. The true objective can be approximated with a surrogate model, and we can use the predictions of this model as well as the uncertainties in these model predictions to guide a search toward higher or lower values of the true objective. This approach, referred to as Bayesian optimization, can also be represented by Fig 5a. Bayesian optimization methods use an acquisition function computed from the predicted mean and uncertainty values to trade off exploration and exploitation when choosing new sequences to sample, in contrast to the acquisition functions used in active learning that maximize exploration.

We compared two popular acquisition functions (upper-confidence bound (UCB) and Thompson sampling (TS)) against random and greedy baselines (see Methods for details). UCB and TS are intended to accelerate identification of top-performing instances over random and greedy baselines by taking into account both the predicted objective and the uncertainty in that prediction. Fig 6 shows the % of top-100 scores found versus fraction of training data seen for three UQ methods and one split per landscape. Across these cases, the uncertainty-based methods almost always perform better than the random baseline but never outperform greedily sampling the sequences with the highest predicted values. In most cases, UCB sampling performs about the same as greedy sampling, with the notable exception of evidential uncertainty on the AAV sampled dataset (Fig 6a), for which it performed worse than random sampling. The performance of TS (a probabalistic method) was typically intermediate between greedy/UCB (deterministic methods) and random sampling. Overall, the uncertainty-based methods do not outperform greedy baselines, suggesting that better UQ methods are needed for protein engineering or that the landscapes studied here are simple enough that they can be optimized by pure exploitation.

All tasks studied in this work are regression problems, in which we attempt to fit a model to a dataset with D data points (x_i, y_i). x_i is a protein sequence representation (either a one-hot encoding or an embedding vector from an ESM language model), and y i ∈ R is a scalar-valued target property from the protein landscapes described in the Datasets section.

The landscapes and splits in this work are taken from the FLIP benchmark [3]. GB1 is a landscape commonly used for investigating epistasis (interactions between mutations) using the binding domain of protein G, an immunoglobulin binding protein in Streptococcal bacteria. These splits are designed primarily to test generalization from few- to many-mutation sequences. The AAV landscape is based on data collected for the Adeno-associated virus capsid protein, which help the virus integrate a DNA payload into a target cell. The mutations in this landscape are restricted to a subset of positions within a much longer sequence. The Meltome landscape includes data from proteins across 13 different species for a non-protein-specific property (thermostability), so it includes both local and global variations. The total number of data points in the GB1, AAV, and Meltome sets are 8,733, 284,009, and 27,951, respectively. In the AAV set, 82,583 are sampled (mutations) and 201,426 are designed. For AAV, only the 82,583 sampled sequences are used for the Random and 7 vs. Rest tasks, while all 284,009 are used for the Sampled vs. Designed task.

The names of several of the tasks were changed slightly from the original FLIP nomenclature for clarity: GB1/Random was originally called GB1/Sampled, AAV/Random was originally called AAV/Sampled, AAV/7 vs. Rest was originally called AAV/7 vs. Many, AAV/Sampled vs. Designed was originally called AAV/Mut-Des, and Meltome/Random was originally called Meltome/Mixed.

We used the pretrained, 650M-parameter ESM-1b model (esm1b_t33_650M_UR50S) from [15] to generate embeddings of the protein sequences in this study and to compare these embeddings to one-hot encoding representations. Sequence embeddings from the final representation layer (layer 33) were mean pooled per amino acid over the length of each protein sequence, which resulted in a fixed embedding size of 1280 for each sequence. In other words, the output of the ESM-1b model is a tensor of size L × 1280, and we averaged over each sequence to obtain a representation vector of size 1280 for each sample.

The base architecture of all CNN models in this work was taken from the CNNs in the FLIP benchmark [3], which took the architecture from previous work [36]. For the one-hot encoding inputs (with a vocabulary of 22 tokens), this was comprised of a convolution with 1024 output channels and kernel width 5, a ReLU non-linear activation function, a linear mapping to 2048 dimensions, a max pool over the sequence, and a linear mapping to 1 dimension. For ESM embedding inputs (of size 1280), the architecture was the same except with 1280 input channels rather than 1024, and a linear mapping to 2560 dimensions rather than 2048.

To train our CNN models, we used a batch size of 256 (GB1, AAV) or 30 (Meltome). Adam [37] was used for optimization with the following learning rates: 0.001 for the convolution weights, 0.00005 for the first linear mapping, and 0.000005 for the second linear mapping. Weight decay was set to 0.05 for both the first and second linear mappings. CNNs were trained with early stopping using a patience of 20 epochs. Each model was trained on an NVIDIA Volta V100 GPU. Reported metrics in Figs 2–4 are the average of training 5 models per split with different seeds for initialization of the CNN parameters and batching / stochastic gradient descent. Code, data, and instructions needed to reproduce results can be found at https://github.com/microsoft/protein-uq.

For all models and landscapes, the sequences were featurized using either one-hot encodings or embeddings from a pretrained language model (see the ESM Embeddings section).

We used the scikit-learn [38] implementation of Bayesian ridge regression (BRR) with default hyperparameters. BRR for one-hot encodings of the Meltome/Random split was not feasible because the required work array was too large to perform the computation with standard 32-bit LAPACK in scipy.

For Gaussian processes (GPs), we used the GPyTorch [39] implementation with the constant mean module, scaled rational quadratic (RQ) kernel covariance module, and Gaussian likelihood. Some GP models (for AAV one-hot encodings and ESM embeddings, and Meltome one-hot encodings) were not feasible to train due to GPU-memory requirements for exact GP models, so these are omitted from the results.

For our uncertainty methods that rely on sampling (dropout, ensemble, and SVI), the final model prediction is defined as the mean of the set of inference samples, and the uncertainty is the standard deviation of these samples. In other words, for a set of predictions E = { G 1 ( x ) , G 2 ( x ) , … , G n ( x ) } (each coming from an individual model G_i), the final prediction is defined as G ^ ( x ) = ∑ G ∈ E G ( x ) n (1) and the uncertainty U(x) is defined as U ( x ) = ∑ G ∈ E ( G ^ ( x ) - G ( x ) ) 2 n (2)

The uncertainty is sometimes defined as the variance U², but using the standard deviation puts the uncertainty in the same units as the predictions.

For dropout uncertainty [19], a single model G was trained normally. At inference time, we applied n = 10 random dropout masks with dropout probability p to obtain the set of predictions E for each input x_i. We tested dropout rates of p ∈ {0.1, 0.2, 0.3, 0.4, 0.5} and reported the model with the lowest miscalibration area.

Similarly for last-layer stochastic variational inference (SVI) [23], we obtained E using n = 10 samples from a set of models where each G_i has the weight and bias terms of its last layer themselves sampled from a distribution q(θ) that has been trained to approximate the true posterior p ( θ | D ).

Traditional model ensembling calculated E using n = 5 models trained using different random seeds for initialization of the CNN parameters and batching / stochastic gradient descent. The computational cost of this approach is 5 times that of a standard CNN model since the cost scales linearly with the size of the ensemble.

In mean-variance estimation (MVE) models, we adapt the base CNN architecture to produce 2 outputs (θ = {μ, σ²}) for each data point (x_i, y_i) in the last layer rather than 1, and we train using the negative log-likelihood loss: L ( θ ) = 1 N ∑ i = 1 N ( y i - μ ( x i ) ) 2 2 σ 2 ( x i ) + 1 2 log ( 2 π σ 2 ( x i ) ) (3)

In practice, the variance (σ²) is clamped to a minimum value of 10⁻⁶ to prevent division by 0.

Evidential deep learning modifies the loss function of the traditional CNN to jointly maximize the model’s fit to data while also minimizing its evidence on errors (increasing uncertainty on unreliable predictions) [21]: L ( x ) = L N L L ( x ) + λ L R ( x ) (4) where L N L L ( x ) is the negative log-likelihood loss defined above, L R ( x ) is the evidence regularizer as defined in Amini et al. [21], and λ controls the trade-off between these two terms. In this study, we use λ = 1 for all evidential models. In these models, the last layer of the model produces 4 outputs m = {γ, ν, α, β} that parameterize the Normal-Inverse-Gamma distribution. This distribution assumes that targets y_i are drawn i.i.d. from a Gaussian distribution with unknown mean and variance θ = {μ, σ²}, where the mean is drawn from a Gaussian and the variance is drawn from an Inverse-Gamma distribution. The output of the evidential model can be divided into the prediction and the epistemic (model) and aleatoric (data) uncertainty components following the analysis of Amini et al. [21]: E [ μ ] = γ ︸ prediction , E [ σ 2 ] = β α - 1 ︸ aleatoric , Var [ μ ] = β ν ( α - 1 ) ︸ epistemic (5)

We report the sum of the aleatoric and epistemic uncertainties as the total uncertainty.

To give a comprehensive report of model accuracy, we computed the following metrics on the test sets: root mean square error (RMSE), mean absolute error (MAE), coefficient of determination (R²), 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² 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 Spearman rank correlation between uncertainty and absolute prediction error. This metric may be particularly relevant in an active learning context, where one wants to acquire labels for the most uncertain points hoping that these are also the highest-error points. This application does not require the uncertainties to be well-calibrated.

Second, 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 [40]. Good calibration may be more important in safety-critical applications.

Following Kompa et al. [41], we measured the coverage as the percentage of true values that fall within the 95% confidence interval (±2σ) of each prediction. This is a indication of the reliability of the uncertainty estimates. A model with high coverage is appropriately cautious in its predictions, which may be most important in applications where safety is a major consideration.

Kompa et al. [41] also define another metric, the width, as the size of the 95% confidence region (4σ). We normalized this width relative to the range (R) of the training set as 4σ/R to make these values unitless and thus more interpretable across datasets. The width is a measure of precision in the uncertainty. In practical applications, narrower intervals (lower width) can help in making more precise and cost-effective decisions. Ideal uncertainties have high coverage and low width, but in some cases, there may be trade-offs between the two. For example, wider widths can help to detect distribution shift, but these wider intervals may not be reliable if coverage is low. The coverage and width metrics may also be more easily interpretable than other calibration metrics [41].

Each active learning run began with a random sample of 1%, 5%, or 10% of the full training data, which was taken from the random splits of the three landscapes. We evaluated several alternatives for adding to this initial dataset using different sampling strategies (acquisition functions): explorative greedy, explorative sample, and random. “Explorative greedy” sampled the sequences with the highest uncertainty; “explorative sample” sampled the data according to the probability of sampling a sequence equal to the ratio of its uncertainty to the sum of all uncertainties in the dataset (i.e. random sampling weighted by uncertainty); and “random” sampled the data uniformly from all unobserved sequences. We employed these sampling strategies 5 times in each active learning run, with the 5 training set sizes equally spaced on a log scale. We repeated this process using 3 folds (different random seeds for sampling initial dataset and “explorative sample” probabilities) and calculated the mean and standard deviation across these folds.

Similarly to active learning, each Bayesian optimization run began with a random sample of 10% of the full training data, which was taken from the random splits of the three landscapes. We used the following acquisition functions: greedy, upper-confidence bound (UCB) [42], Thompson sampling (TS) [43], and random. In greedy sampling, the sequence with the best predicted value was selected. The UCB strategy added the predicted uncertainties to the predicted values and selected the sequence with the largest sum. For TS, we added each predicted value to a number sampled randomly from a Gaussian distribution with a mean of 0 and a standard deviation of the corresponding predicted uncertainty, and again selected the largest sum. The “random” strategy used in Bayesian optimization was the same as that used in active learning (new points were sampled with uniform probability). As with active learning, we used these strategies 5 times in each run, with the 5 training set sizes equally spaced on a log scale. We report the mean and standard deviation across 3 folds (different random seeds for sampling the initial dataset).