Conformal Prediction Under Feedback Covariate Shift for Biomolecular Design

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.

Code at https://github.com/clarafy/conformal-for-design. Updated title to match published version

References in corpus (12)

• Practical Bayesian Optimization of Machine Learning Algorithms
• Protein sequence design with deep generative models
• Generating and designing DNA with deep generative models
• Testing for Outliers with Conformal p-values
• Robust Validation: Confident Predictions Even When Distributions Shift
• AdaLead: A simple and robust adaptive greedy search algorithm for sequence design
• Adaptive Conformal Inference Under Distribution Shift
• Learn then Test: Calibrating Predictive Algorithms to Achieve Risk Control
• Autofocused oracles for model-based design
• Distribution-free uncertainty quantification for classification under label shift
• PAC Confidence Predictions for Deep Neural Network Classifiers
• Tracking the risk of a deployed model and detecting harmful distribution shifts

Cited by in corpus (3)

• Is novelty predictable?
• Non-exchangeable Conformal Prediction for Temporal Graph Neural Networks
• Non-Exchangeable Conformal Risk Control