Conformal Prediction Under Feedback Covariate Shift for Biomolecular Design

TL;DR

This work extends conformal prediction to the challenging setting of feedback covariate shift (FCS), where test inputs depend on the training data through iterative design or optimization procedures. By introducing data-dependent likelihood-ratio weights, the authors construct full conformal confidence sets under FCS that maintain finite-sample coverage for any regression model, and they offer randomized versions that achieve exact coverage. The methodology is instantiated with ridge regression and Gaussian process regression, with practical computational strategies to avoid exponential model retraining. Through protein design experiments—including combinatorially complete fluorescence datasets and AAV capsid data—the approach quantifies predictive uncertainty for designed sequences and guides design-hyperparameter choice. Overall, the framework provides a principled, model-agnostic tool for uncertainty-aware design in iterative ML pipelines and broadens conformal prediction to dependent-training/test regimes.

Abstract

Many applications of machine learning methods involve an iterative protocol in which data are collected, a model is trained, and then outputs of that model are used to choose what data to consider next. For example, one data-driven approach for designing proteins is to train a regression model to predict the fitness of protein sequences, then use it to propose new sequences believed to exhibit greater fitness than observed in the training data. Since validating designed sequences in the wet lab is typically costly, it is important to quantify the uncertainty in the model's predictions. This is challenging because of a characteristic type of distribution shift between the training and test data in the design setting -- one in which the training and test data are statistically dependent, as the latter is chosen based on the former. Consequently, the model's error on the test data -- that is, the designed sequences -- has an unknown and possibly complex relationship with its error on the training data. We introduce a method to quantify predictive uncertainty in such settings. We do so by constructing confidence sets for predictions that account for the dependence between the training and test data. The confidence sets we construct have finite-sample guarantees that hold for any prediction algorithm, even when a trained model chooses the test-time input distribution. As a motivating use case, we demonstrate with several real data sets how our method quantifies uncertainty for the predicted fitness of designed proteins, and can therefore be used to select design algorithms that achieve acceptable trade-offs between high predicted fitness and low predictive uncertainty.

Figures (11)

• Figure 1: Illustration of feedback covariate shift. In the left graph, the blue distribution represents the training input distribution, $P_X$. The dark gray line sandwiched by lighter gray lines represents the mean $\pm$ the standard deviation of $P_{Y \mid X}$, the conditional distribution of the label given the input, which does not change between the training and test data distributions (left and right graphs, respectively). The blue dots represent training data, $Z_{1:n} = \{Z_1, \ldots, Z_n\}$, where $Z_i = (X_i, Y_i)$, which is used to fit a regression model (middle). Algorithms that use that trained model to make decisions, such as in design problems, active learning, and Bayesian optimization give rise to a new test-time input distribution, $P_{X; Z_{1:n}}$ (right graph, green distribution). The green dots represent test data.
• Figure 2: Illustration of single-shot protein design. The gray distribution represents the distribution of fitnesses under the training sequence distribution. The blue circles represent the fitnesses of three training sequences, and the goal is to propose a sequence with even higher fitness. To that end, we fit a regression model to the training sequences labeled with experimental measurements of their fitnesses, then deploy some design procedure that uses that trained model to propose a new sequence believed to have a higher fitness (green circle).
• Figure 3: Quantifying predictive uncertainty for designed proteins, using the blue fluorescence data set. (a) Distributions of labels of designed proteins, for different values of the inverse temperature, $\lambda$, and different amounts of training data, $n$. Labels surpass the fitness range observed in the combinatorially complete data set, $[0.091, 1.608]$, due to additional simulated measurement noise. (b) Empirical coverage, compared to the theoretical lower bound of $1 - \alpha = 0.9$ (dashed gray line), and (c) distributions of confidence interval widths achieved by full conformal prediction for feedback covariate shift (our method) over $T = 2000$ trials. In (a), and (c), the whiskers signify the minimum and maximum observed values. (d) Distributions of Jaccard distances between the confidence intervals produced by full conformal prediction for feedback covariate shift and standard covariate shift tibshirani2019conformal. (e, f) Same as (b, c) but using full conformal prediction for standard covariate shift.
• Figure 4: Comparison of trade-off between predicted fitness and predictive certainty on the red and blue fluorescence data sets. (a) Trade-off between mean confidence interval width and mean predicted fitness for different values of the inverse temperature, $\lambda$, and $n = 384$ training data points. (b) Empirical probability that the smallest fitness value in the confidence intervals of designed proteins exceeds the true fitness of one of the wild-type parent sequences, mKate2. (c) For $n = 384$ and $\lambda = 6$, the distributions of both confidence interval width and predicted fitnesses of designed proteins.
• Figure 5: Quantifying uncertainty for predicted fitnesses of designed adeno-associated virus (AAV) capsid proteins. (a) Mean true fitness of designed sequences resulting from different values of the inverse temperature, $\lambda \in \{1, 2, \ldots, 7\}$. The dashed black line is the mean true fitness of sequences drawn from the NNK sequence distribution (i.e., the training distribution). (b) Top: empirical coverage of randomized staircase confidence sets (Section \ref{['app:rand-split']}) constructed for designed sequences. The dashed black line is the expected empirical coverage of $1 - \alpha = 0.9$. Bottom: fraction of confidence sets with infinite size (dashed gray line) and mean size of non-infinite confidence sets (solid gray line). The set size is reported as a fraction of the range of fitnesses in all the labeled data, $[-7.53, 8.80]$. (c) Trade-off between mean predicted fitness and mean confidence set size for $\lambda \in \{1, 2, \ldots, 7\}$. The dashed black line is the mean predicted fitness for sequences from the training distribution.

Theorems & Definitions (14)

• Definition 1
• Theorem 1
• Definition 2
• Proposition 1
• proof
• Lemma 1
• proof
• Theorem 2
• proof
• Theorem 3