Efficient and Differentiable Conformal Prediction with General Function Classes

1 INTRODUCTION . Modern machine learning models can yield highly accurate predictions in many applications . As these predictions are often used in critical decision making , it is increasingly important to accompany them with an uncertainty quantification of how much the true label may deviate from the prediction . A common approach to quantifying the uncertainty in the data is to learn a prediction set—a setvalued analogue of usual ( point ) predictions—which outputs a subset of candidate labels instead of a single predicted label . For example , this could be a prediction interval for regression , or a discrete label set for multi-class classification . A common requirement for learned prediction sets is that it should achieve valid coverage , i.e . the set should cover the true label with high probability ( such as 90 % ) on a new test example ( Lawless & Fredette , 2005 ) . In addition to coverage , the prediction set is often desired to have a good efficiency , such as a low length or small cardinality ( Lei et al. , 2018 ; Sadinle et al. , 2019 ) , in order for it to be informative . Note that coverage and efficiency typically come as a trade-off , as it is in general more likely to achieve a better coverage using a larger set . This paper is concerned with the problem of finding the most efficient prediction set with valid coverage . Our approach builds on conformal prediction ( Vovk et al. , 2005 ) , a powerful framework for generating prediction sets from ( trained ) base predictors with finite-sample coverage guarantees . Conformal prediction has been used for learning prediction sets in a variety of tasks in regres- sion ( Lei & Wasserman , 2014 ; Lei et al. , 2018 ; Romano et al. , 2019 ) , classification ( Cauchois et al. , 2020b ; Romano et al. , 2020 ; Angelopoulos et al. , 2020 ) , structured prediction ( Bates et al. , 2021 ) , and so on . However , the conformalization step in conformal prediction by default does not offer the flexibility for optimizing additional efficiency metrics , as the efficiency is already determined by the associated score function and the target coverage level . As a concrete example , the Conformalized Quantile Regression algorithm learns a single width adjustment parameter that turns a two-sided quantile predictor into a prediction interval of valid coverage ( Romano et al. , 2019 ) ; however , it does not offer a way of further optimizing its length ( cf . Figure 1 Left ) . For certain efficiency metrics and prediction tasks , several approaches have been proposed , for example by designing a better score function ( Angelopoulos et al. , 2020 ) , using base predictors of a specific form ( Izbicki et al. , 2019 ; 2020 ; Sadinle et al. , 2019 ) , or selecting a best training hyperparameter ( Yang & Kuchibhotla , 2021 ) . However , optimizing the efficiency for more general tasks or efficiency metrics still largely requires “ manual ” efforts by the researcher , as it ( 1 ) often relies on specific domain knowledge about the task at hand ; ( 2 ) is often done in conjunction with conformal prediction in multiple rounds of trial-and-error ; ( 3 ) is often done by reasoning about high-level properties of the efficiency loss and coverage constraints ( e.g . what makes the length short ) , but not by directly optimizing the efficiency-coverage trade-off in a data-dependent way . To the best of our knowledge , there is a lack of a more principled and unified approach for optimizing any efficiency metric subject to valid coverage over any class of prediction sets . In this paper , we cast the above task as a constrained empirical risk minimization ( ERM ) problem of optimizing the efficiency subject to the coverage constraint , over any general function class of prediction sets with potentially multiple learnable parameters . This is motivated by a simple observation that vanilla conformal prediction is already equivalent to solving such a constrained ERM with one learnable parameter ( Section 2.1 ) . Overall , our algorithm can be viewed as an automatic and data-dependent approach for optimizing the efficiency simulatneously with conformal prediction . Our contributions are summarized as follows . • We propose CP-Gen ( Conformal Prediction with General Function Class ) , a generalization of conformal prediction to learning multiple parameters . CP-Gen selects within an arbitrary class of prediction sets by solving the constrained ERM problem of best efficiency subject to valid empirical coverage ( Section 3.1 ) , and is a systematic extension of existing algorithms . • We show theoretically that CP-Gen achieves approximately valid coverage and near-optimal efficiency within class , whenever the class is low-capacity with respect to both the coverage and the efficiency loss ( Section 3.2 , with concrete examples in Appendix C ) . We also provide a practical variant CP-Gen-Recal using data splitting and reconformalization , which achieves exact coverage , as well as good efficiency under additional assumptions ( Section 3.3 ) . • To address the issue that CP-Gen and CP-Gen-Recal involve a non-differentiable coverage constraint , we develop a differentiable approximation using surrogate losses and Lagrangians ( Section 4 ) . This allows us to solve the constrained ERM problem over higher-dimensional continuous parameter spaces via gradient-based optimization , and is more flexible than existing algorithms that require discretization and brute-force search . • We empirically demonstrate that CP-Gen-Recal with our gradient-based implementation can learn prediction sets with valid coverage and significantly improved efficiency on three real-data tasks : prediction intervals for regression with improved length , minimum-volume prediction sets for multi-output regression , and label prediction sets for ImageNet ( Section 5 & Appendix F ) . We illustrate our main insight via the coverage-vs-efficiency trade-off plots in Figure 1 : While vanilla conformal prediction only learns a single parameter ( within its conformalization step ) by a simple thresholding rule over a coverage-efficiency curve , our CP-Gen is able to further improve the efficiency by thresholding a region formed by a larger function class . 1.1 RELATED WORK . Learning prediction sets via conformal prediction The framework of conformal prediction for learning prediction sets is originated in the early works of ( Vovk et al. , 1999 ; 2005 ; Shafer & Vovk , 2008 ) . The main advantage of conformal prediction is that it yields ( marginal ) coverage guarantees regardless of the data distribution ( i.e . distribution-free ) . More recently , conformal prediction has been applied to a variety of uncertainty quantification tasks , such as prediction intervals for regression ( Papadopoulos , 2008 ; Vovk , 2012 ; 2015 ; Lei & Wasserman , 2014 ; Vovk et al. , 2018 ; Lei et al. , 2018 ; Romano et al. , 2019 ; Izbicki et al. , 2019 ; Guan , 2019 ; Gupta et al. , 2019 ; Kivaranovic et al. , 2020 ; Barber et al. , 2021 ; Foygel Barber et al. , 2021 ) , label prediction sets for classification problems ( Lei et al. , 2013 ; Sadinle et al. , 2019 ; Romano et al. , 2020 ; Cauchois et al. , 2020b ; a ; Angelopoulos et al. , 2020 ) , and prediction sets for structured output ( Bates et al. , 2021 ) . Optimizing efficiency in addition to valid coverage The problem of finding a prediction set with ( approximate ) valid coverage and small size has been considered , e.g . in Pearce et al . ( 2018 ) ; Chen et al . ( 2021 ) for regression and Park et al . ( 2019 ) for classification ; however , these approaches do not use conformal prediction . Yang & Kuchibhotla ( 2021 ) propose to minimize the length of the conformal interval over either a finite class or a linear aggregation of base predictors , and provides coverage and efficiency guarantees . All above works formulate this task as a risk minimization problem , yet are restricted to considering either finite classes or specific efficiency loss functions . Our work is inspired by ( Yang & Kuchibhotla , 2021 ) and generalizes the above works by allowing any function class and efficiency loss , along with providing a differentiable approximate implementation . The problem of optimizing the efficiency can also be done by utilizing structures of the particular efficiency loss to choose a specific base predictor and an associated prediction set ( Lei & Wasserman , 2014 ; Sadinle et al. , 2019 ; Izbicki et al. , 2019 ; 2020 ) . By contrast , our approach does not require either the efficiency loss or the base predictor to possess any structure , and is thus complementary . Other algorithms and theory An alternative line of work constructs prediction intervals / prediction sets by aggregating the prediction over multiple base predictors through Bayesian neural network ( Mackay , 1992 ; Gal & Ghahramani , 2016 ; Kendall & Gal , 2017 ; Malinin & Gales , 2018 ; Maddox et al. , 2019 ) or ensemble methods ( Lakshminarayanan et al. , 2016 ; Ovadia et al. , 2019 ; Huang et al. , 2017 ; Malinin et al. , 2019 ) . However , these methods do not typically come with ( frequentist ) coverage guarantees . The recent work of Hoff ( 2021 ) studies ways of enhancing Bayes-optimal prediction with frequentist coverage . Prediction intervals can also be obtained by parameter estimation using a parametric model for the data ( Cox , 1975 ; Bjornstad , 1990 ; Beran , 1990 ; Barndorff-Nielsen & Cox , 1996 ; Hall et al. , 1999 ; Lawless & Fredette , 2005 ) ; see ( Tian et al. , 2020 ) for a review . However , the coverage of such prediction intervals relies heavily on the parametric model being correct ( well-specified ) , and can even fail in certain high-dimensional regimes where the model is indeed correct ( Bai et al. , 2021 ) . 2 PRELIMINARIES . Uncertainty quantification via prediction sets We consider standard learning problems in which we observe a dataset D of examples ( xi , yi ) ∈ X × Y from some data distribution , and wish to predict the label y from the input x . A prediction set is a set-valued function C : X → 2Y where C ( x ) is a subset of Y . Two prevalent examples are regression ( Y = R ) in which we can choose C ( x ) ⊂ R as a prediction interval , and ( multi-class ) classification ( Y = [ L ] : = { 1 , . . . , L } ) in which we can choose C ( x ) ⊂ [ L ] as a ( discrete ) label prediction set . Coverage and efficiency The ( marginal ) coverage probability ( henceforth coverage ) of a prediction set C is defined as Coverage ( C ) : = P ( Y ∈ C ( X ) ) where ( X , Y ) is a test example from the same data distribution . We also define the ( mis ) -coverage loss Lcoverage ( C ) : = 1− Coverage ( C ) = P ( Y /∈ C ( X ) ) . A learned prediction set is often desired to achieve valid coverage in the sense that Coverage ( C ) ≥ 1− α for some α ∈ ( 0 , 1 ) . Here 1− α is a pre-determined target coverage level ; typical choices are e.g . 1 − α ∈ { 90 % , 95 % } , which corresponds to picking α ∈ { 0.1 , 0.05 } . In addition to valid coverage , it is often desired that the prediction set has a good efficiency ( such as small size ) . This is motivated by the fact that valid coverage can be achieved trivially if we do not care about the size , e.g . by always outputting C = Y , which is not informative . Throughout this paper we will use ` eff to denote the particular efficiency loss we care about , where ` eff ( C ; ( x , y ) ) measures the efficiency loss of C on an example ( x , y ) , such as the length ( Lebesgue measure ) of prediction intervals , or the size ( cardinality ) of label prediction sets . Nested set framework We adopt the nested set framework of ( Gupta et al. , 2019 ) for convenience for our presentation and analysis . A family { Ct } t∈T ⊂R is said to be a ( family of ) nested sets if t ≤ t′ implies that Ct ( x ) ⊂ Ct′ ( x ) for all x ∈ X . Throughout this paper out notation Ct or Cθ , t are assumed to be nested sets with respect to t. We assume that our efficiency loss ` eff is non-decreasing w.r.t . its ( set-valued ) argument , i.e . ` eff ( C ; ( x , y ) ) ≤ ` eff ( C ′ ; ( x , y ) ) if C ⊆ C ′ . Therefore , for nested sets the loss t 7→ ` eff ( Ct ; ( x , y ) ) is non-decreasing in t. As the coverage loss L ( Ct ) = P ( Y /∈ Ct ( X ) ) ( and its empirical version ) is instead non-increasing in t , the efficiency loss and the coverage loss always comes as a trade-off .

Summary: This paper generalizes the standard conformal prediction calibration setup to a constrained empirical risk minimization problem. Specifically, this work seeks to optimize some efficiency loss, while satisfying coverage constraints. This formulation allows for the introduction of multiple, learnable parameters which can help find a better set-based predictor. The paper explains the implications of this approach by analyzing the generalization error that may occur when transferring the solution learned by constrained ERM to a test population. It also explains practical ways of learning this problem via differentiable surrogate losses and Lagrangians. Contributions-wise, the paper contributes validating theoretical analysis that proves that this method can achieve approximate coverage and near optimal efficiency for certain set-function classes. It also empirically shows that the proposed method can improve over baselines that are not directly optimized for efficiency.