PaperlayerView exact revision on arXiv

Source-linked AI summary

Conformal Prediction Under Feedback Covariate Shift for Biomolecular Design

Clara Fannjiang, Stephen Bates, Anastasios N. Angelopoulos, Jennifer Listgarten, Michael I. Jordan

arXiv:2202.03613v5cs.LGq-bio.QMstat.ME

TL;DR

When a trained model influences which data are collected next, uncertainty about predictions on those inputs is difficult to quantify. The paper constructs conformal confidence sets for feedback covariate shift and demonstrates their use in protein design. In the reported design experiments, coverage is close to 0.9, while uncertainty measures can rise as inverse temperature increases.

  • Problem

    In feedback settings, test inputs depend on training data, making it unclear whether model predictions on those inputs can be trusted.

    Source: §1 · p. 1

  • Method

    The paper uses conformal prediction to construct confidence sets under feedback covariate shift, including settings where a trained model induces the test input distribution.

    Source: §2.1 · p. 4, §2.2 · p. 5

  • Results

    Empirical coverage stays very close to 1 − α = 0.9 across considered inverse temperatures, while infinite-size sets and mean non-infinite set size increase with inverse temperature.

    Source: §3.2.3 · p. 16

  • Takeaways & Limitations

    In protein design, the confidence sets can help practitioners choose hyperparameters that trade off predicted fitness against predictive uncertainty.

    Source: §3.2.3 · p. 16, §4 · p. 16

    Without further assumptions on the training and test input distributions, the method provides no corresponding upper bound on coverage probability.

    Source: §2.3 · p. 6

Abstract

from arXiv · show

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.

1 Uncertainty quantification under feedback loops

When test data are chosen using a trained model, feedback loops can make it unclear whether predictions on designed inputs are trustworthy. The paper formalizes feedback covariate shift and extends conformal prediction to quantify uncertainty in this setting.

1 Uncertainty quantification under feedback loops

  • In protein design and other feedback-loop deployments, test inputs can depend on training data even while the conditional label distribution remains unchanged.The paper calls this setting feedback covariate shift (FCS).Source: §1 · pp. 1–2

1.1 Quantifying uncertainty with valid confidence sets

  • The authors generalize conformal prediction to produce finite-sample valid confidence sets for any regression model under FCS.Coverage is marginal over the randomness in training and test data, not conditional on a particular test input.Source: §1.1 · pp. 2–3

1.2 Our contributions

  • The paper formalizes FCS, introduces valid and randomized exact-coverage confidence sets, and demonstrates uncertainty quantification for designed proteins using real data sets.It recommends using the method to select design-algorithm hyperparameters that balance high predictions with low predictive uncertainty.Source: §1.2 · p. 3

1.3 Prior work

  • Unlike standard covariate shift, FCS includes dependence between training and test data; the authors adapt conformal machinery for covariate shift to address it.They note that common model-based uncertainty estimates do not come with formal guarantees, while conformal prediction can be applied to any modeling strategy.Source: §1.3 · pp. 3–4

2 Conformal prediction under feedback covariate shift

The section formalizes feedback covariate shift and adapts conformal prediction to produce confidence sets with coverage guarantees when test inputs depend on the training data.

2.1 Feedback covariate shift

  • FCS allows i.i.d. training data to induce an order-invariant, otherwise unrestricted test-input distribution, while keeping the conditional label distribution fixed.Source: §2.1 · p. 4

2.2 Conformal prediction for exchangeable data

  • For exchangeable data, conformal prediction includes candidate labels whose test scores are sufficiently small relative to a quantile of training-point scores.Source: §2.2 · p. 5

2.3 Conformal prediction under FCS

  • Under FCS, likelihood-ratio weights account for dependence and differing input distributions, while randomizing the quantile yields exact coverage.Source: §2.3 · pp. 5–6
  • For FCS, full conformal sets achieve marginal coverage at least 1−α for every α, assuming each induced test-input distribution is absolutely continuous with respect to training inputs.Source: §2.3 · p. 6

2.4 Efficient computation of confidence sets under feedback covariate shift

  • When candidate-label predictions are computationally simple, such as linear in y, scores and weights require n+1 model fits rather than (n+1)|Y| fits.Source: §2.4 · p. 8

Ridge regression.

  • For ridge regression, candidate-label predictions are linear in y, so stored slopes and intercepts enable computing candidate-grid scores and weights after n+1 fits.Source: Ridge regression · p. 9

Gaussian process regression.

  • For Gaussian processes, candidate-label mean predictions are linear in y for linear and arbitrary kernels, enabling analogous efficient score and weight computation.Source: Gaussian process regression · p. 9

2.5 Data splitting

  • When data are abundant or score computation is inefficient, splitting training and calibration data restores independent standard covariate shift conditional on training; randomized split conformal achieves exact coverage.Source: §2.5 · p. 9

3 Experiments with protein design

The experiments evaluate confidence sets for designed proteins using combinatorially complete fluorescence data and AAV capsid packaging data. They examine coverage and uncertainty, and use these measures to assess the trade-off between predicted fitness and predictive certainty.

3.1.1 Protocol for design experiments

  • Fluorescence experiments used combinatorially complete data to simulate design with labeled test sequences, fitting ridge regression to interaction features and sampling designs according to predicted fitness.The inverse temperature λ controls the trade-off between predicted fitness and distance from training data.Source: §3.1 · p. 10, §3.1.1 · p. 11

3.1.2 Results

  • Empirical coverage reached at most 0.924 for n = 96, λ = 6, close to the 0.9 lower bound.The authors interpret the small excess over the lower bound as suggesting intervals are nearly as small as possible while achieving coverage.Source: §3.1.2 · p. 12
  • For any λ, intervals tend to be smaller with more training data and larger as λ increases; some sparse-data, high-λ settings yield intervals spanning the full fitness range.The authors attribute wider intervals at higher λ to designed sequences having higher predicted fitness and greater model uncertainty.Source: §3.1.2 · p. 12
  • Standard covariate-shift confidence sets tend to be more conservative than feedback covariate-shift sets, especially with less training data or higher λ.The methods become more similar when the mapping from data to the test-input distribution is less sensitive to changes in the data.Source: §3.1.2 · pp. 12–13
  • For red fluorescence, the probability of exceeding the wild-type fitness decreases above λ=2; for blue fluorescence, it increases as interval width grows more slowly.The red–blue trade-off differs across fitness functions, so the appropriate λ can differ as well.Source: §3.1.2 · pp. 13–14

3.2.1 Protein design problem: AAV capsid proteins with improved packaging ability

  • Most AAV capsid sequences fail to assemble into a genome-packaging capsid, a prerequisite for gene-therapy delivery.The AAV design problem concerns improving packaging ability.Source: §3.2.1 · p. 15

3.2.2 Protocol for design experiments

  • For AAV design, a neural network parameterized nucleotide distributions to seek higher predicted fitness, and randomized data splitting produced confidence sets for rejection-sampled designs.The study used held-out labeled sequences for proposal sampling and calibration.Source: §3.2.2 · pp. 15–16

3.2.3 Results

  • At λ=7, more than 17% of AAV designs received infinite confidence sets; λ∈{4, 5} was judged an acceptable balance of higher predicted fitness and low uncertainty.The fraction of infinite sets and the mean size of finite sets both increase with λ.Source: §3.2.3 · p. 16

4 Discussion

The paper presents a way to assess learned-model predictions when those predictions influence future data. Its protein-design examples illustrate a trade-off between desirable predictions and predictive certainty, while the FCS framework also covers other feedback-loop settings.

4 Discussion

  • The methodology addresses whether learned-model predictions can be trusted when they influence the data encountered later.The discussion frames this as relevant to consequential decisions involving feedback.Source: §4 · p. 16
  • Protein-design examples demonstrate how the approach can navigate the trade-off between desirable predictions and predictive certainty.Source: §4 · p. 16
  • FCS captures settings including active learning, adaptive experimental design, and Bayesian optimization, with applications and extensions identified as future directions.Source: §4 · p. 16

A1.1 Proof of Theorem 1

The proof generalizes conformal prediction to pseudo-exchangeable data, then applies that result to data generated under feedback covariate shift. The argument relies on characterizing score distributions conditional on the unordered data values.

A1.1 Proof of Theorem 1

  • Pseudo-exchangeability represents the joint density through factor functions and a core function that is invariant to input ordering.Source: A1.1 Proof of Theorem 1 · p. 22
  • Conditioning on score values without their ordering enables an exact description of the scores’ distribution for the proof.Source: A1.1 Proof of Theorem 1 · p. 22
  • Theorem 2 generalizes conformal prediction to pseudo-exchangeable random variables and provides a confidence set with a coverage guarantee.Source: A1.1 Proof of Theorem 1 · pp. 23–24
  • The proof identifies FCS data as pseudo-exchangeable, allowing Theorem 2 to yield Theorem 1 as a corollary.Source: A1.1 Proof of Theorem 1 · p. 24

A1.2 A randomized confidence set achieves exact coverage

The section randomizes the quantile used to construct a full conformal confidence set. Under FCS and an absolute-continuity condition, the resulting set has exact coverage; the result also applies to standard covariate shift.

A1.2 A randomized confidence set achieves exact coverage

A1.3 Data splitting

Data splitting conditions on the training data, reducing FCS to standard covariate shift for calibration and test data. The approach uses density-ratio weights to construct split conformal confidence sets, trading data efficiency for fitting only one model.

A1.3 Data splitting

  • Split conformal partitions labeled data, fits one regression model on the training portion, and uses calibration scores to form confidence sets.Source: A1.3 Data splitting · p. 26
  • Data splitting reduces FCS to standard covariate shift by conditioning on the training data and treating the fitted model as fixed.The calibration and test data are then independent and share the same conditional label distribution.Source: A1.3 Data splitting · p. 27
  • Density-ratio weights compare test-input and training-input densities when constructing split conformal confidence sets.The score may use a heuristic uncertainty measure, and setting it to 1 recovers the residual score.Source: A1.3 Data splitting · p. 27
  • The split conformal set has a coverage guarantee under standard covariate shift, but uses labeled data less efficiently than full conformal prediction.The efficiency trade-off may matter when labeled data are limited.Source: A1.3 Data splitting · pp. 26–27

A1.4 Data splitting with randomization achieves exact coverage

Randomizing the quantile gives split conformal confidence sets with exact coverage under stated conditions. A staircase construction exploits the piecewise-constant inclusion probability over candidate labels, avoiding a separate calculation for every label.

A1.4 Data splitting with randomization achieves exact coverage

  • Under standard covariate shift, the randomized split conformal set has exact coverage for score functions of the specified form, any miscoverage level, and an absolutely continuous test-input distribution.The result applies because standard covariate shift is a special case of feedback covariate shift.Source: A1.4 Data splitting with randomization… · p. 28
  • The staircase construction sorts m scores and requires O(m) additional floating-point operations.The exact-coverage property follows from independently drawing a randomized quantile for each candidate label.Source: A1.4 Data splitting with randomization… · p. 28
  • Candidate-label inclusion probability is piecewise constant, with deterministic inclusion or exclusion outside the quantile-boundary case and randomized inclusion at the boundary.The probability is highest at the predicted mean and decreases as labels move farther away.Source: A1.4 Data splitting with randomization… · p. 28
  • The randomized staircase set computes inclusion probabilities over intervals between sorted calibration scores instead of testing every candidate label.For the residual score, each interval maps to two equal-length intervals around the prediction.Source: A1.4 Data splitting with randomization… · pp. 28–29
  • For any continuous score function, exact coverage extends to efficiently computable preimages of the score intervals; the algorithm assumes distinct scores almost surely when handling endpoints.Computational feasibility depends on being able to compute those preimages efficiently.Source: A1.4 Data splitting with randomization… · pp. 29–30

A2.1 Ridge regression

For ridge-regression feedback covariate shift, the likelihood-based scores and weights can be computed more efficiently than by fitting a model for every candidate label.

A2.1 Ridge regression

  • When test-input likelihood depends on ridge-regression predictions, scores and weights require n + 1 model fits plus O(n·p·|Y|) operations, rather than (n + 1)×|Y| fits.The method computes candidate-label scores and likelihood ratios using the fitted ridge model.Source: A2.1 Ridge regression · p. 30
  • Algorithm 3 takes training data, a test input, candidate labels, and train- and test-input likelihood routines as inputs, and returns scores and likelihood ratios.For each candidate label, it computes leave-one-out residual scores and corresponding likelihood ratios.Source: A2.1 Ridge regression · p. 30

A2.2 Gaussian process regression

For Gaussian-process regression, the predictive mean's dependence on the candidate label permits an efficient computation of feedback-shift confidence-set scores and weights. The fluorescent-protein details describe the sequence features and simulated measurement noise used in the experiments.

A2.2 Gaussian process regression

  • Gaussian-process regression scores and weights can be computed efficiently when test-input likelihood depends on the model’s predictive mean and variance.The method uses the mean and variance at the input in each augmented leave-one-out data set.Source: A2.2 Gaussian process regression · p. 31
  • Because the predictive mean is linear in the candidate label and predictive variance is constant over that label, computation needs n + 1 model fits rather than (n + 1)×|Y|.This allows the ridge-regression computation strategy to be reused.Source: A2.2 Gaussian process regression · p. 31

A3 Additional details and results on designing fluorescent proteins

  • Fluorescent-protein sequences were represented by thirteen signed bits and their pairwise interaction terms.The bits encode which of two wild-type parents each site’s amino acid matches.Source: A3 Additional details and results on… · p. 31
  • The fluorescent-protein experiments added zero-mean Gaussian measurement noise, with sequence-specific standard deviations estimated from residuals of a fitted model with terms up to seventh order.Noise was added each time a sequence was sampled for training or design.Source: A3 Additional details and results on… · p. 31

A4 Additional details on AAV experiments

The supplied AAV-experiment details describe the NNK training distribution, simulated noise and neural-network setup, and figures that assess uncertainty across design settings. The figure captions specify comparisons and encodings but do not report the plotted outcomes.

A4 Additional details on AAV experiments

  • The NNK distribution samples the first two nucleotides of each codon uniformly from four bases and the last from T or G.The specified probabilities aim for amino-acid diversity while avoiding stop codons.Source: A4 Additional details on AAV experiments · p. 31
  • The AAV neural network used two 100-unit tanh hidden layers and was trained for 10 epochs on 7,552,729 data points with variance-based weights.Sequences were one-hot encoded; training used Adam with batch size 64.Source: A4 Additional details on AAV experiments · p. 32
  • Figure A2 compares red-fluorescence label distributions across inverse temperatures and training sizes, then reports coverage, interval widths, and FCS-versus-SCS Jaccard distances.The caption attributes labels beyond the observed fitness range to additional simulated measurement noise.Source: Figure A2 · p. 34
  • Figures A3 and A4 vary feature interaction order or ridge regularization strength and compare FCS with SCS on Jaccard distance, coverage, and interval width for blue and red fluorescence.Both captions specify n = 48 and λ = 6; A3 varies interaction order, while A4 varies γ.Source: Figure A3 · p. 35, Figure A4 · p. 36
  • Figures A5 and A6 contrast SCS's n + 1 weights with FCS's (n + 1) × |Y| weight matrix, whose entries vary linearly with candidate label y.A5 varies inverse temperature; A6 varies ridge regularization strength.Source: Figure A5 · p. 37, Figure A6 · p. 38

Related papers

Loading 2202.03613v5…