E FFICIENT AND D IFFERENTIABLE C ONFORMAL P REDICTION WITH G ENERAL F UNCTION C LASSES
Yu Bai
Salesforce Research yu.bai salesforce
Huan Wang & Yingbo Zhou & Caiming Xiong
Salesforce Research { huan.wang, yingbo.zhou, cxiong salesforce com }
A BSTRACT
Quantifying the data uncertainty in learning tasks is often done by learning a prediction interval or prediction set of the label given the input.
Two commonly desired properties for learned prediction sets are valid coverage and good efﬁciency (such as low length or low cardinality).
Conformal prediction is a powerful technique for learning prediction sets with valid coverage, yet by default its conformalization step only learns a single parameter, and does not optimize the efﬁciency over more expressive function classes.
In this paper, we propose a generalization of conformal prediction to multiple learnable parameters, by considering the constrained empirical risk minimization (ERM) problem of ﬁnding the most efﬁcient prediction set subject to valid empirical coverage.
This meta-algorithm generalizes existing conformal prediction algorithms, and we show that it achieves approximate valid population coverage and near-optimal efﬁciency within class, whenever the function class in the conformalization step is low-capacity in a certain sense.
Next, this ERM problem is challenging to optimize as it involves a non-differentiable coverage constraint.
We develop a gradient-based algorithm for it by approximating the original constrained ERM using differentiable surrogate losses and Lagrangians.
Experiments show that our algorithm is able to learn valid prediction sets and improve the efﬁciency signiﬁcantly over existing approaches in several applications such as prediction intervals with improved length, minimum-volume prediction sets for multi-output regression, and label prediction sets for image classiﬁcation.
I NTRODUCTION
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 quantiﬁcation 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 classiﬁcation.
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 efﬁciency , 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 efﬁciency 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 ﬁnding the most efﬁcient 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 ﬁnite-sample coverage guarantees.
Conformal prediction has been used for learning prediction sets in a variety of tasks in regres-
[Figure]
sion (Lei & Wasserman, 2014; Lei et al., 2018; Romano et al., 2019), classiﬁcation (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 ﬂexibility for optimizing additional efﬁciency metrics, as the efﬁciency 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 efﬁciency 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 speciﬁc form (Izbicki et al., 2019; 2020; Sadinle et al., 2019), or selecting a best training hyperparameter (Yang & Kuchibhotla, 2021).
However, optimizing the efﬁciency for more general tasks or efﬁciency metrics still largely requires “manual” efforts by the researcher, as it (1) often relies on speciﬁc 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 efﬁciency loss and coverage constraints (e.g. what makes the length short), but not by directly optimizing the efﬁciency-coverage trade-off in a data-dependent way.
To the best of our knowledge, there is a lack of a more principled and uniﬁed approach for optimizing any efﬁciency 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 efﬁciency 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 efﬁciency 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 efﬁciency 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 efﬁciency within class, whenever the class is low-capacity with respect to both the coverage and the efﬁciency 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 efﬁciency 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 ﬂexible 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 signiﬁcantly improved efﬁciency 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-efﬁciency 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-efﬁciency curve , our CP-Gen is able to further improve the efﬁciency by thresholding a region formed by a larger function class.
1. R ELATED 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 quantiﬁcation 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 classiﬁcation 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 efﬁciency in addition to valid coverage The problem of ﬁnding a prediction set with (approximate) valid coverage and small size has been considered, e.g. in Pearce et al.
; Chen et al.
(2021) for regression and Park et al.
for classiﬁcation; however, these approaches do not use conformal prediction.
Yang & Kuchibhotla (2021) propose to minimize the length of the conformal interval over either a ﬁnite class or a linear aggregation of base predictors, and provides coverage and efﬁciency guarantees.
All above works formulate this task as a risk minimization problem, yet are restricted to considering either ﬁnite classes or speciﬁc efﬁciency loss functions.
Our work is inspired by (Yang & Kuchibhotla, 2021) and generalizes the above works by allowing any function class and efﬁciency loss, along with providing a differentiable approximate implementation.
The problem of optimizing the efﬁciency can also be done by utilizing structures of the particular efﬁciency loss to choose a speciﬁc 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 efﬁciency 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-speciﬁed), and can even fail in certain high-dimensional regimes where the model is indeed correct (Bai et al., 2021).
P RELIMINARIES
Uncertainty quantiﬁcation via prediction sets We consider standard learning problems in which we observe a dataset D of examples ( x i , y i ) ∈
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 → 2 Y 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) classiﬁcation ( Y =
[ L ] := { 1 , . . .
, L } ) in which we can choose C ( x ) ⊂
[ L ] as a (discrete) label prediction set .
Coverage and efﬁciency The (marginal) coverage probability (henceforth coverage ) of a prediction set C is deﬁned as
[Equation]
where ( X, Y ) is a test example from the same data distribution.
We also deﬁne the (mis)-coverage loss L coverage ( 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 efﬁciency (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 (cid:96) eﬀ to denote the particular efﬁciency loss we care about, where (cid:96) eﬀ ( C ; ( x, y )) measures the efﬁciency 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 { C t } t ∈T ⊂ R is said to be a (family of) nested sets if t ≤ t (cid:48) implies that C t ( x ) ⊂ C t (cid:48)
( x ) for all x ∈ X .
Throughout this paper out notation C t or
C θ,t are assumed to be nested sets with respect to t .
We assume that our efﬁciency loss (cid:96) eﬀ is non-decreasing w.r.t.
its (set-valued) argument, i.e. (cid:96) eﬀ ( C ; ( x, y )) ≤ (cid:96) eﬀ ( C (cid:48) ; ( x, y ))
if C ⊆ C (cid:48) .
Therefore, for nested sets the loss t (cid:55)→ (cid:96) eﬀ ( C t ; ( x, y )) is non-decreasing in t .
As the coverage loss
L ( C t ) =
P ( Y / ∈ C t ( X ))
(and its empirical version) is instead non-increasing in t , the efﬁciency loss and the coverage loss always comes as a trade-off.
2. C ONFORMAL PREDICTION
Conformal prediction (Vovk et al., 2005; Lei & Wasserman, 2014) is a powerful technique for learning prediction sets with coverage guarantees.
The core of conformal prediction is its conformalization step, which turns any base prediction function (or training algorithm) into a prediction set.
We here brieﬂy review conformal prediction using the vanilla (split) conformal regression method of (Lei et al., 2018), and refer the readers to (Angelopoulos & Bates, 2021) for more examples.
Given any base predictor f : X → R (potentially learned on a training dataset D train ), conformal prediction outputs a prediction interval
[Equation]
where (cid:98) t ∈ R ≥ 0 is chosen as the (1 − α )
-quantile 1 of | y − f
( x ) | on a calibration dataset D cal with size n cal :=
| D cal | using the following conformalization step:
[Equation]
The main guarantee for the learned interval C (cid:98) t is that it achieves a (1 − α ) coverage guarantee of the form P D cal , ( X,Y ) ( Y ∈ C (cid:98) t ( X )) ≥ 1 − α (Lei et al., 2018, Theorem 2.2).
The proof relies on the exchangeability between the scores {| y i − f ( x i ) |} n cal i =1 and | Y − f
( X ) | , which allows this guarantee to hold in a distribution-free fashion (i.e. for any data distribution).
Conformal prediction as a constrained ERM with one parameter We start by a simple reinterpretation that the conformalization step (2) is equivalent to solving a constrained empirical risk minimization (ERM) problem with a single learnable parameter t (cf.
Appendix A for the proof).
Proposition 1 (Conformal regression as a constrained ERM with one learnable parameter) .
The parameter (cid:98) t ∈ R deﬁned in (2) is the solution to the following constrained ERM problem
[Equation]
C
{ C θ t } θ ∈ Θ t α
≥
ε
: ( (cid:98) θ, (cid:98) t ) ← arg min θ ∈ Θ ,t ∈T (cid:98) L eﬀ ( C θ,t ) :
= 1 n cal (cid:88) i ∈
D cal (cid:96) eﬀ
( C θ,t ( x i ) , y i ) subject to (cid:98) L coverage ( C θ,t ) :
= 1 n cal (cid:88) i ∈ D cal 1 { y i / ∈ C
θ,t ( x i ) } ≤ α + ε 0 .
Above, (cid:96) eﬀ ( C ; ( x, y ))
=
length( C ( x )) is the length of the interval C ( x ) .
Though simple, this re-interpretation suggests a limitation to the conformalization step (2) as well as its analogue in other existing conformal methods: It only learns a single parameter t , and thus cannot further optimize the efﬁciency due to the coverage-efﬁciency trade-off (cf.
Figure 1).
However, the form of the constrained ERM problem (3) suggests that it can be readily extended to more general function classes with more than one learnable parameters, which is the focus of this work.
C ONFORMAL PREDICTION WITH GENERAL FUNCTION CLASSES 3. A LGORITHM
Our algorithm, Conformal Prediction with General Function Classes ( CP-Gen ; full description in
Algorithm 1), is an extension of the constrained ERM problem (3) into the case of general function classes with multiple learnable parameters.
CP-Gen takes in a function class of prediction sets
[Equation]
where (as mentioned)
we assume that { C θ,t } t ∈T is a nested set for each θ ∈ Θ .
The parameter set Θ as well as the form of C θ,t in (4) can be arbitrary, depending on applications and the available base predictors at hand.
Given C , our algorithm then solves the constrained ERM problem (5) of ﬁnding the smallest interval among C subject to valid coverage on dataset D cal .
Compared with vanilla conformal prediction, Algorithm 1 allows more general tasks with an arbitrary function class and efﬁciency loss; for example, this encompasses several recent algorithms such as ﬁnite hyperparameter selection and linear aggregation (Yang & Kuchibhotla, 2021; Chen et al., 2021).
We remark that (5) includes an additional relaxation parameter ε 0 ≥ 0 for the coverage constraint.
This is for analysis (for Proposition 2(b) & 7(b)) only; our implementation uses ε 0 = 0 .
3. T HEORY
An important theoretical question about CP-Gen is whether it achieves coverage and efﬁciency guarantees on the population (test data).
This section showcases that, by standard generalization arguments, CP-Gen achieves approximate validity and near-optimal efﬁciency whenever function class is low-capacity in a certain sense.
We remark that our experiments use the modiﬁed algorithm CP-Gen-Recal (Section 3.3) which involves a reconformalization step.
Here we focus on
CP-Gen as we believe its theory could be more informative.
Let L { eﬀ , coverage } ( C θ,t ) :
= E [ (cid:96) { eﬀ , coverage } ( C θ,t ; ( X, Y ))] denote the population coverage and efﬁciency lossesfor any ( θ, t ) .
We deﬁne the following uniform concentration quantities:
[Equation]
[Table]
,t
[Table]
−
[Table]
C θ,t ) (cid:12)(cid:12)(cid:12) .
The following proposition connects the generalization of CP-Gen to the above uniform concentration quantities by standard arguments (see Appendix B for the proof.
We remark that the proof relies on D cal being i.i.d., which is slightly stronger than exchangeability assumption commonly assumed in the conformal prediction literature.)
Proposition 2 (Generalization of CP-Gen ) .
The prediction set C (cid:98) θ, (cid:98) t learned by Algorithm 1 satisﬁes
(a) (Approximately valid population coverage)
We have
[Equation]
i
e. the population coverage of C (cid:98) θ, (cid:98) t is at least 1 − α − ( ε 0 + ε coverage ) .
(b) (Near-optimal efﬁciency)
Suppose ε 0 ≥ ε coverage , then we further have
[Equation]
i.e. C (cid:98) θ, (cid:98) t achieves 2 ε eﬀ -near-optimal efﬁciency against any prediction set within C with at least (1 − α ) population coverage
Examples of good generalization Proposition 2 shows that CP-Gen acheives approximate coverage and near-optimal efﬁciency if the concentration terms ε eﬀ and eﬀ coverage are small.
In
Appendix C, we bound these on two example function classes: Finite Class (Proposition 4) and
VC/Rademacher Class (Proposition 5).
Both classes admit bounds of the form { ε eﬀ , eﬀ coverage } ≤ (cid:112)
Comp( C ) /n cal with high probability via standard concentration arguments, where Comp( C ) is a certain complexity measure of C .
Combined with Proposition 2, our CP-Gen algorithm with these classes achieve an 1 − α − (cid:112) Comp( C )
/n cal approximate coverage guarantee and (cid:112) Comp( C ) /n
cal near-optimal efﬁciency guarantee.
In particular, our Proposition 4 recovers the coverage guarantee for the ﬁnite-class selection algorithm of (Yang & Kuchibhotla, 2021, Theorem 1) though our efﬁciency guarantee is more general.
We remark that both examples above contain important applications.
The ﬁnite class contains e.g. optimizing over a K -dimensional hyperparameter to use via grid search, with e.g. (1 /δ )
K conﬁdence sets and thus Comp( C ) =
O (log((1 /δ ) K ))
=
O ( K log(1 /δ )) .
The VC/Rademacher class contains the important special case of linear classes with K base predictors (we defer the formal statement and proof to Appendix C.3).
Also, these examples are not necessarily exhaustive.
Our take-away is rather that we may expect CP-Gen to generalize well (and thus achieves good coverage and efﬁciency) more broadly in practice, for instance whenever it learns K (cid:28) n cal parameters.
3. A LGORITHM WITH VALID COVERAGE VIA RECONFORMALIZATION
Although CP-Gen enjoys theoretical bounds on the coverage and efﬁciency, a notable drawback is that it does not guarantee exactly valid (at least) 1 − α coverage like usual conformal prediction, and its approximate coverage bound depends on the uniform concentration quantity ε coverage that is not computable from the observed data without structural assumptions on the function class C .
To remedy this, we incorporate a simple reconformalization technique on another recalibration dataset D recal (e.g. a further data split), which guarantees valid ﬁnite-sample coverage by exchangeability.
We call this algorithm CP-Gen-Recal and provide the full description in Algorithm 2.
We remark that this reconformalization technique for obtaining guaranteed 1 − α coverage is widely used in the conformal prediction literature, e.g. (Angelopoulos et al., 2020).
However, to the best of our knowledge, there is no known analysis for our CP-Gen-Recal algorithm for general function classes, for which we provide a result below (formal statement and proof can be found in Proposition 7 & Appendix D).
The proof of the coverage bound is standard as in the conformal prediction literature, while the proof of the efﬁciency bound builds upon the result for CP-Gen (Proposition 2(b)) and handles additional concentration terms from the reconformalization step.
Proposition 3 (Coverage and efﬁciency guarantee for CP-Gen-Recal ; Informal version) .
The prediction set C (cid:98) θ, (cid:98) t recal learned by Algorithm 2 achieves (1 − α ) ﬁnite-sample coverage:
P D recal , ( X,Y ) ( Y ∈ C (cid:98) θ, (cid:98) t recal ) ≥ 1 − α .
Further, it achieves O ( ε eﬀ + ε coverage + 1 / √ n recal ) nearoptimal efﬁciency under additional regularity assumptions.
D IFFERENTIABLE OPTIMIZATION
Our (meta) algorithms CP-Gen and CP-Gen-Recal involve solving the constrained ERM problem (5).
One feasible case is when Θ is ﬁnite and small, in which we enumerate all possible θ ∈ Θ and ﬁnd the optimal (cid:98) t for each θ efﬁciently using quantile computation.
However, this optimization is signiﬁcantly more challenging when the underlying parameter set Θ is continuous and we wish to jointly optimize over ( θ, t ) :
The coverage loss (cid:98) L coverage ( C θ,t ) is non-differentiable and its “gradient” is zero almost everywhere as it uses the zero-one loss.
This makes the coverage constraint challenging to deal with and not amenable to any gradient-based algorithm.
To address this non-differentiability, we develop a gradient-based practical implementation by approximating the coverage constraint.
We ﬁrst rewrite each individual coverage loss into the form
[Equation]
where (cid:96) 01 ( z ) :
= 1 { z < 0 } is the zero-one loss.
(Such a rewriting is possible in most cases by taking s θ,t as a suitable “score-like” function; see Appendix E for instantiations in our experiments.)
Then, inspired by the theory of surrogate losses (Bartlett et al., 2006), we approximate (cid:96)
( z ) by the hinge loss (cid:96) hinge ( z ) =
[1 − z ] + which is (almost everywhere) differentiable with a non-trivial gradient.
We ﬁnd the hinge loss to perform better empirically than alternatives such as the logistic loss.
To deal with the (modiﬁed) constraint, we turn it into an exact penalty term with penalty parameter λ ≥ 0 , and use a standard primal-dual formulation (Bertsekas, 1997) to obtain an unconstrained min-max optimization problem on the Lagrangian:
[Equation]
(8) where (cid:98) L hinge ( C θ,t ) :
= 1 n cal (cid:80) n cal i =1 (cid:96) hinge ( s θ,t ( x i , y i )) is the empirical hinge loss on the calibration dataset D cal .
Our ﬁnal practical implementation of (5) solves the problem (8) by Stochastic Gradient
Descent-Ascent (with respect to ( θ, t ) and λ ) to yield an approximate solution ( (cid:98) θ, (cid:98) t ) .
We remark that in our experiments where we use the reconformalized version CP-Gen-Recal , we only keep the (cid:98) θ obtained from (8) and perform additional reconformalization to compute (cid:98) t recal to guarantee coverage.
We also emphasize that the approximation in (8) makes the problem differentiable at the cost of deviating from the true constrained ERM problem (5) and thus potentially may sacriﬁce in terms of the efﬁciency.
However, our experiments in Section 5 show that such an implementation can still improve the efﬁciency over existing approaches in practice, despite the approximation.
E XPERIMENTS
We empirically test our CP-Gen-Recal algorithm
(using the practical implementation (8)) on three representative real-data tasks.
The concrete construction of { C θ,t } will be described within each application.
Throughout this section we choose 1 − α = 90% as the nominal coverage level, and use the CP-Gen-Recal algorithm to guarantee coverage in expectation.
We provide ablations with α ∈ { 80% , 95% } in Appendix G.1 and G.2 and the CP-Gen algorithm in Appendix G.3.
Several additional ablations and analyses can be found in Appendix H.
[Table]
5. I MPROVED PREDICTION INTERVALS VIA CONFORMAL QUANTILE FINETUNING
Setup We consider regression tasks in which we use quantile regression (pinball loss) to train a base quantile predictor (cid:98) F ( x ) =
[ (cid:98) f lo ( x ) , (cid:98) f hi ( x )]
= (cid:98) θ (cid:62)
(cid:98) Φ( x ) on D train (with learning rate decay by monitoring validation loss on D cal ).
Here (cid:98) Φ : X → R d h is the learned representation function, (cid:98) θ 0 ∈ R
d h × 2 is the last linear layer ( d h denotes the last hidden dimension), and (cid:98) f
lo , (cid:98) f
hi are the learned { lower, upper } quantile functions (see Appendix E.1 for more details on the training procedure).
Given (cid:98) F , we learn a baseline prediction interval of the form [ (cid:98)
f lo ( x ) − t, (cid:98) f
hi ( x ) + t ] on D recal via
Conformalized Quantile Regression ( CQR ) (Romano et al., 2019).
We then attempt to improve the length over CQR by conformal quantile ﬁnetuning :
Fix the representation function (cid:98) Φ and ﬁnetune the linear layer θ using our CP-Gen-Recal algorithm, so that
C θ,t ( x ) =
[ θ (cid:62) lo (cid:98) Φ( x ) − t, θ (cid:62) hi (cid:98) Φ( x )+ t ] (where θ =
[ θ lo , θ hi ] ).
We learn a new (cid:98) θ on D cal via (8) (where (cid:96) eﬀ is chosen as the length), and then compute (cid:98) t recal on D recal as in Algorithm 2.
We perform the above on 9 real-world regression datasets with a 3-layer MLP with width d h = 64 , similar as (Romano et al., 2019; Feldman et al., 2021).
Additional details about the setup can be found in Appendix E.1.
We also test various tweaks of the CQR baseline (results provided in
Appendix H.2).
Results Table 1 compares the (test) coverage and length between CQR and the ﬁnetuned linear layer via our CP-Gen-Recal .
While both CQR and CP-Gen-Recal achieves valid 90% coverage, CP-Gen-Recal can systematically improve the length over CQR on all tasks.
Table 1 also reports the pinball loss for both the base (cid:98) θ 0 (cid:98) Φ( x ) as well as the ﬁne-tuned (cid:98) θ (cid:62) (cid:98) Φ( x ) on the test set
D test .
Intriguingly, our conformal ﬁnetuning made the pinball loss worse while managing to improve the length.
This suggests the unique advantage of our constrained ERM objective, as it rules out the simple explanation that the length improvement is just because of a lower test loss.
We remark that while CP-Gen-Recal improves the length over CQR , it comes at a cost in terms of worse conditional coverage (analysis presented in Appendix H.1).
5. M INIMUM - VOLUME PREDICTION SETS FOR MULTI - OUTPUT REGRESSION
Setup This task aims to learn a box-shaped prediction set for multi-output regression with a small volume .
Our learning task is regression with output dimension d out > 1 .
We ﬁrst learn a based predictor (cid:98) f : R d → R d out by minimizing the MSE loss on D train .
We then learn a box-shaped prediction set of the form C u ( x ) =
(cid:81)
d out i =
[ (cid:98) f i ( x ) − (cid:98) u i , (cid:98) f i ( x ) + (cid:98) u i ]
by
one of the following methods:
• ( Coord-wise ): Each (cid:98) u
i is obtained by vanilla conformalization (2) over the i -th output coordinate on D cal ∪ D recal .
To guarantee 1 − α coverage, each coordinate is conformalized at level 1 − α/d out , motivated by the union bound.
[Table]
• ( Coord-wise-Recal ): Perform the above on D cal to learn (cid:98) u i , and reconformalize an additional t ≥ 0 on D recal to reshape the prediction set proportionally :
[Equation]
• ( CP-Gen-Recal , ours): Optimize the volume directly over all u ∈ R d out using (8) on D cal , where (cid:96) eﬀ
( C u ; ( x, y ))
= (cid:81)
d out i =1 (2 u i ) is chosen as the volume loss.
We then reconformalize an additional (cid:98) t ≥ 0 on D recal to reshape the prediction set same as in (9).
Note that this reconformalization step is equivalent to Algorithm 2 with the re-parametrization (cid:98) u
(cid:55)→ ( (cid:98) θ, (cid:98) t ) where (cid:98) θ ∈ R d out − 1 >denotes the ratio between (cid:98) u i , and (cid:98) t ∈ R > 0 denotes a common scale
Our datasets are a collection of next-state prediction tasks with multi-dimensional continuous states in ofﬂine reinforcement learning (RL), constructed similarly as D4RL (Fu et al., 2020) with some differences.
Additional details about the dataset and experimental setup are in Appendix E.2.
We also test an additional Max-score-Conformal baseline (which uses vanilla conformal prediction with score function (cid:107) y − (cid:98) f ( x ) (cid:107) ∞ , equivalent to a hypercube -shaped predictor) in Appendix H.3, which we ﬁnd also performs worse than our CP-Gen-Recal .
Results Table 2 reports the (test) coverage and volume of the above three methods.
The Coord-wise method achieves valid coverage but is quite conservative (over-covers), which is as expected as the union bound is worst-case in nature and the coordinate-wise conformalization does not utilize the potential correlation between the output coordinates.
Coord-wise-Recal achieves approximately 90% coverage with a much smaller volume.
Our
CP-Gen-Recal also achieves valid 90% coverage but a further lower volume across all tasks.
This suggests that optimizing the volume over all possible u ∈ R d out data-dependently using our
CP-Gen-Recal is indeed more ﬂexible than pre-determined conformalization schemes such as
Coord-wise .
Additional experiment: label prediction sets for ImageNet We show that CP-Gen-Recal can learn label prediction sets for ImageNet with valid coverage and improved size over existing approaches, by ﬁnding an optimized set of ensemble weights over multiple base neural networks (Table 5).
Due to space limit, we defer the setup and results to Appendix F.
C ONCLUSION
This paper proposes Conformal Prediction with General Function Class, a conformal prediction algorithm that optimizes the efﬁciency metric subject to valid coverage over a general function class of prediction sets.
We provide theoretical guarantees for its coverage and efﬁciency in certain situations, and develop a gradient-based practical implementation which performs well empirically on several large-scale tasks.
We believe our work opens up many directions for future work, such as stronger theoretical guarantees via more structured function classes, further improving the gradientbased approximate implementation, or experiments on other uncertainty quantiﬁcation tasks.
A P ROOF OF P ROPOSITION
Recall that C t ( x ) =
[ f ( x ) − t, f ( x ) + t ] satisﬁes (cid:96) eﬀ
( C t ; ( x, y ))
=
length( C t ( x ))
= 2 t for any ( x, y ) .
Also, we have 1 { y / ∈ C t ( x ) } = 1 { y / ∈
[ f ( x ) − t, f ( x ) + t ] } = 1 {| y − f ( x ) | > t } , or equivalently 1 { y ∈ C t ( x ) }
= 1 {| y − f ( x ) | ≤ t } .
Therefore, problem (3) is equivalent to minimize t ≥t subject to 1 − (cid:98) L coverage ( C t ) :
= 1 n cal n cal (cid:88) i = 1 {| y
i − f ( x i ) | ≤ t } ≥ 1 − α.
By deﬁnition of quantiles, this problem is solved at (cid:98) t =
(cid:100) (1 − α ) n cal (cid:101) -th largest element of {| y i − f ( x i ) |} n cal i =1 , which is the desired result.
B P ROOF OF P ROPOSITION
(a) As C (cid:98) θ, (cid:98) t
solves problem
it
satisﬁes the constraint
(cid:98)
L
coverage ( C (cid:98) θ, (cid:98) t )
≤ α + ε
Therefore,
L coverage
[Equation]
C
[Equation] [Equation]
=
[Equation]
coverage
[Equation]
C
[Equation]
coverage
[Equation]
= α + ε 0 + ε coverage .
(b) Suppose ε coverage ≤ ε 0 .
Taking any ( θ, t ) ∈
Θ × T such that L coverage ( C θ,t
≤
α
, we have (cid:98) L coverage ( C θ,t )
≤ L coverage ( C θ,t ) + ε coverage ≤ α + ε coverage ≤ α + ε 0 .
This shows that ( θ, t ) lies within the constraint set of problem (5).
Thus as ( (cid:98) θ, (cid:98) t ) further minimizes the loss (cid:98) L eﬀ within the constraint set, we have (cid:98) L eﬀ ( C (cid:98) θ, (cid:98) t ) ≤ (cid:98)
L eﬀ ( C θ,t ) .
This shows that
[Equation]
As the above holds simultaneously for all ( θ, t ) ∈
Θ × T with at most α coverage loss, taking the sup over the left-hand side yields
[Equation]
C E XAMPLES OF GOOD GENERALIZATION FOR CP-Gen
We provide two concrete examples where the concentration terms ε eﬀ and ε coverage are small with high probability, in which case Proposition 2 guarantees that CP-Gen learns an prediction set with approximate validity and near-optimal efﬁciency.
Assumption A (Bounded and Lipschitz efﬁciency loss) .
The loss function (cid:96) eﬀ satisﬁes
A1 (Bounded loss).
| (cid:96) eﬀ
( C θ,t ; ( x, y ))
| ≤ M for all ( θ, t ) ∈
Θ × T and all
( x, y ) ∈
X × Y .
A2 ( t -Lipschitzness).
t (cid:55)→ (cid:96) eﬀ
( C θ,t ; ( x, y )) is L T -Lipschitz for all ( θ, x, y ) ∈
Θ × X × Y .
Assumption B (Bounded T ) .
The parameter space T ⊂ R is bounded: sup t ∈T | t | ≤ B T .
C. F INITE CLASS
Proposition 4 (Finite class) .
Suppose Θ is a ﬁnite set ( N Θ := | Θ | < ∞ ), and suppose Assumptions A1, A2, B hold.
Then, we have with probability at least 1 − δ that
[Equation]
where C > 0 is an absolute constant.
Proof.
We ﬁrst bound ε coverage .
Fix any θ ∈ Θ , deﬁne
[Equation]
(10) to be the smallest possible t ∈ T such that the C θ,t ( x ) contains y .
Observe that, as { C θ,t ( x ) } t ∈T are nested sets, the coverage event can be rewritten as 1 { y ∈ C θ,t ( x ) } = 1 { t ≥ t θ ( x, y ) } for any ( x, y ) .
Therefore, we have
[Equation]
where we have deﬁned F θ : R → [0 , 1] as the CDF of the random variable t θ ( X, Y ) and similarly (cid:98) F θ as the empirical CDF of the same random variable over the ﬁnite dataset D cal .
Applying the
Dvoretzky-Kiefer-Wolfowitz (DKW) inequality (Massart, 1990, Corollary 1) yields that
[Equation]
with probability at least 1 − δ .
Now, taking the union bound with respect to θ ∈ Θ (where for each θ we plug in tail probability δ/ 2 N Θ ) we get that with probability at least 1 − δ/ 2 ,
[Equation]
for some absolute constant C > 0 .
We next bound ε eﬀ .
Fix any θ ∈ Θ .
We have by standard symmetrization argument that
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
Above, (i) used the Lipschitzness Assumption A2 and the Rademacher contraction inequality (Vershynin, 2018, Exercise 6.7.7); (ii) used Assumption B, and (iii) used E
ε i (cid:104)(cid:12)(cid:12)(cid:12) 1 n cal (cid:80) n cal i =
ε i (cid:12)(cid:12)(cid:12)(cid:105) ≤ (cid:18) E ε i (cid:20)(cid:16) 1 n cal (cid:80) n cal i =
ε i (cid:17) 2 (cid:21)(cid:19) 1 / 2 = 1 / √ n cal .
(Above ε i iid ∼ Unif( {± 1 } ) are Rademacher variables.)
Next, as each loss | (cid:96) eﬀ ( C θ,t ; ( x, y ))
| ≤ M by Assumption A1, the random variable
[Equation]
L eﬀ ( C
[Equation]
,t ) − L eﬀ
[Equation]
C
[Equation] [Equation]
satisﬁes the M/n cal ﬁnite-difference property.
Therefore by McDiarmid’s Inequality, we have with probability at least 1 − δ that
[Equation]
Finally, by union bound over θ ∈ Θ (where we plug in δ/ 2 N Θ as tail probability into the above), we have with probability at least 1 − δ/ 2 that
[Equation]
(11) together with (12) is the desired result.
C. VC/R ADEMACHER CLASS
Next, for any class C , let VC( C ) :
=
VC( { ( x, y ) (cid:55)→ 1 { y / ∈
C
θ,t ( x ) } : θ ∈ Θ , t ∈ T } ) denote its VC dimension with respect to the coverage loss.
Proposition 5 (VC/Rademacher class) .
We have for some absolute constant C > 0 that
(a) Suppose VC( C )
=
K + 1 < ∞ , then with probability at least 1 − δ/ 2 , ε coverage ≤ C (cid:112) ( K + 1 + log(1 /δ ))
/n cal .
(b) Suppose Assumption A1 holds.
Then we have with probability at least 1 − δ/ 2 that
[Equation]
where R eﬀ n cal ( C ) :
=
[Equation]
is the Rademacher complexity of the class C with respect to (cid:96) eﬀ
(above ε i iid ∼ Unif( {± 1 } ) ).
Proof.
(a) By assumption, the class of Boolean functions { ( x, y ) (cid:55)→ 1 { y / ∈ C θ,t ( x ) }}
( θ,t ) ∈
Θ ×T has VC dimension K + 1 < ∞ .
Therefore by the standard Rademacher complexity bound for VC classes (Vershynin, 2018, Theorem 8.3.23) and McDiarmid’s Inequality, we have with probability at least 1 − δ/ 2 that
[Equation]
(b) We have by standard symmetrization argument that (below ε i iid ∼ Unif( {± 1 } )
denote
Rademacher variables)
[Equation]
Further by Assumption A
[Equation]
the quantity ε eﬀ satisﬁes M/n cal bounded-difference, so applying McDiarmid’s Inequality gives that with probability at least 1 − δ/ 2 ,
[Equation]
C. C ASE STUDY : L INEAR CLASS
In this section
[Equation]
we study prediction intervals with a speciﬁc linear structure and show that it satisﬁes the conditions of the VC/Rademacher class of Proposition 5.
Concretely, suppose we have a regression task ( Y = R ), and the prediction interval C θ,t ( x ) takes a linear form
[Equation]
) where θ ∈ Θ ⊂ R K , Φ hi , Φ lo : X → R K are feature maps such that Φ lo ( x ) i ≤ Φ hi ( x ) i for all i ∈ [ K ] , σ : X → R > 0 .
For intuitions, we can think of Φ { hi , lo } as pretrained representation functions and σ as an (optional) pretrained function for modeling the variability of y | x .
Note that this encompasses linear ensembling of several existing methods, such as vanilla conformal regression (Lei et al., 2018) by taking Φ hi =
Φ lo =
Φ where each Φ i : X → R is a base predictor, as well as Conformalized Quantile Regression (Romano et al., 2019) where each (Φ lo ,i , Φ hi ,i ) is a pair of learned lower and upper quantile functions.
Our goal is to ﬁnd an optimal linear function of this representation that yields the shortest prediction interval (with ﬁxed width) subject to valid coverage.
We assume that both the features and the parameters are bounded:
Assumption C (Bounded features and parameters) .
We have sup θ ∈ Θ (cid:107) θ (cid:107) ≤ B Θ , sup x ∈X (cid:107) Φ( x ) (cid:107) ≤
B Φ , sup x ∈X σ ( x ) ≤ B σ , and sup t ∈T | t | ≤ B T .
The following result shows that Proposition 5 is applicable on the linear class.
Corollary 6 (Coverage and length guarantees for linear class) .
For the ( K + 1) -dimensional linear class (13) , suppose Assumption C holds, and we take the efﬁciency loss to be the length of the interval: (cid:96) eﬀ ( C ; ( x, y )) :
= length( C ( x )) .
Then, we have with probability at least 1 − δ (over the calibration dataset D cal ) that
[Equation]
where C > 0 is an absolute constant.
Proof.
We verify the conditions of Proposition 5.
First, we have
[Equation]
The set within the indicator above is the union of two sets (cid:8) ( x, y ) :
y − θ (cid:62) Φ hi ( x ) − tσ ( x ) > 0 (cid:9) and (cid:8) ( x, y ) : θ (cid:62) Φ lo ( x ) − y − tσ ( x ) > 0 (cid:9) .
Note that each family of sets (over ( θ, t ) ∈
R K × R are linear halfspaces with feature dimension K + 2 ), and thus has VC-dimension ≤ K + 2 .
Applying the VC dimension bound for unions of sets (Van Der Vaart & Wellner, 2009, Theorem 1.1), we get
VC( C ) ≤ C (cid:48)
( K + 2 + K + 2) ≤ C ( K + 1) for some absolute constant C > 0 .
Therefore the condition of Proposition 5(a) holds from which we obtain the desired bound for ε coverage .
To bound ε eﬀ , we ﬁrst note that for any ( x, y ) ∈ X × R ,
[Equation]
and thus the boundedness assumption (Assumption A1) holds with M deﬁned above.
Next, we have the following bound on the Rademacher complexity
[Equation]
≤
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12) (cid:12) (cid:42)
[Equation]
n cal n cal (cid:88) i
=
ε i (Φ hi
[Table]
n
[Table] [Table]
B
[Table]
n
[Table]
σ
[Table]
x
[Table]
≤ C · B Θ B Φ + B T B σ √ n cal .
Applying Proposition 5(b), we get ε eﬀ ≤ C · [ B Θ B Φ + B T B σ ] · (cid:112) log(1 /δ )
/n cal with probability at least 1 − δ .
This is the desired bound for ε eﬀ .
D T HEORETICAL GUARANTEE FOR CP-Gen-Recal
In this section we state and prove the formal theoretical guarantee for the CP-Gen-Recal algorithm (Algorithm 2).
Deﬁne the score t θ ( X, Y ) :
= inf { t ∈ T : Y ∈ C θ,t ( X ) } and let F θ ( t ) :
= P ( Y ∈ C θ,t ( X ))
=
P ( t θ ( X, Y ) ≤ t ) denote its CDF.
Assumption D (Lower bounded density for score function) .
For any θ ∈ Θ , t θ ( X, Y ) has a positive density f θ ( t )
=
F (cid:48)
θ ( t ) > 0 on t ∈ T .
Further, let t θ, 1 − α :=
inf { t ∈ T : F θ ( t ) ≥ 1 − α } denote its (1 − α ) quantile, then there exists some constants c 0 ,
[Equation]
0 > 0 such that
[Equation]
Proposition 7 (Valid coverage and near-optimal efﬁciency for reconformalized algorithm) .
The following holds for Algorithm 2:
(a) (Valid coverage)
For any possible (cid:98) θ ∈ Θ learned in Line 1 and the resulting (cid:98) t recal , we have E
[Equation]
recal
[Equation]
coverage
[Equation]
C (cid:98)
[Equation]
(cid:98) t recal ) (cid:105) ≤ α, and thus P D recal , ( X,Y ) (cid:16) Y ∈
[Equation]
X ) (cid:17) ≥ 1 − α.
(b) (Efﬁciency)
Suppose
Assumptions
A2 and D
hold
max (cid:110)
ε coverage + 1 /n cal , 2 (cid:112) log(1 /δ )
/n recal (cid:111) ≤ c 0 δ 0 (recall the deﬁnition
of
ε coverage
in (
and
ε coverage
≤
ε
Then for δ
≤
5 , we have with probability at least 1 − δ that
[Equation]
Proof.
(a) As the learned parameter (cid:98) θ (and thus the family of nested sets C (cid:98) θ,t ) is independent of the recalibration dataset D recal , we have that the scores t (cid:98) θ ( x, y ) on dataset D recal and a new test point ( X, Y ) are exchangeable given any (cid:98) θ .
Therefore by (Gupta et al., 2019, Proposition 1), we have for any (cid:98) θ ∈ Θ that
[Equation]
or equivalently E
[Equation]
(b) For any θ ∈ Θ , deﬁne the score function t θ ( x, y ) the same as in (10), and similarly deﬁne the
CDF F θ ( t ) :
= P ( t θ ( X, Y ) ≤ t ) and its empirical counterpart (cid:98) F cal θ ( t ) and (cid:98) F recal θ ( t ) as the ﬁnitesample version on dataset D cal and D recal respectively.
We ﬁrst analyze (cid:98) t .
By the same derivation as in (11), we have
[Equation]
As ( (cid:98) θ, (cid:98) t ) solves the constrained ERM (5) and by the assumption that (cid:96) eﬀ ( C θ,t ; ( x, y )) is monotone in t , we have that (cid:98) t is the minimal value of t ∈ T such that (cid:98) F cal (cid:98) θ ( t ) ≥ 1 − α .
Therefore, (as | D cal | = n cal and { t θ ( x i , y i ) }
i ∈
D cal are almost surely distinct by Assumption D,) we have
[Equation]
This shows that
[Equation]
where we recall that t (cid:98) θ, 1 − α is the (1 − α ) (population) quantile of t (cid:98) θ
( X, Y ) .
Note that F (cid:48) θ ( t )
=
f θ ( t ) ≥ c 0 on t ∈ [ t (cid:98) θ, 1 − α − δ 0 , t (cid:98) θ, 1 − α + δ 0 ] by Assumption D. Further, ε coverage + 1 /n cal ≤ c 0 δ 0 .
Therefore, by monotonicity of F θ , we must have (cid:98) t ∈
[ t (cid:98) θ, 1 − α − δ 0 , t (cid:98) θ, 1 − α + δ 0 ] , and thus
[Equation]
We next analyze (cid:98) t recal .
As the dataset D recal is independent of (cid:98) θ , we can apply the DKW Inequality (Massart, 1990, Corollary 1) to obtain that
[Equation]
with probability at least 1 − δ .
Using a similar argument as above
[Equation]
we get (for δ ≤ 0 . 5 )
[Equation]
As
[Equation]
log
[Equation]
1 /δ )
/n recal ≤ c
δ
[Equation]
we can apply the similar argument as above to deduce that
[Equation]
Combining (14) and (15) and using the Lipschitzness of the efﬁciency loss (Assumption A2), we get
[Equation]
≤ L T · (cid:16)(cid:12) (cid:12)(cid:12)(cid:98) t recal − t (cid:98) θ, 1 − α (cid:12) (cid:12)(cid:12) + (cid:12) (cid:12)(cid:12) t (cid:98) θ, 1 − α − (cid:98) t (cid:12) (cid:12)(cid:12) (cid:17) ≤ CL T ·   ε coverage + n − 1 cal + (cid:115) log(1 /δ )
n recal 
 /c 0 .
Finally, as we assumed ε coverage ≤ ε 0 , the condition of Proposition 2(b) holds, so we have
L
[Equation]
inf
[Equation] [Equation]
Summing the preceding two bounds, we get
L eﬀ ( C
[Equation]
θ,
[Equation]
) ≤ inf
[Equation]
C
[Equation] [Equation]
CL
[Equation]

[Equation]
n
[Equation]
cal + (cid:115) log(1 /δ )
n recal 
 /c 0 .
which is the desired result.
E A DDITIONAL EXPERIMENTAL DETAILS E. C ONFORMAL QUANTILE FINETUNING
Datasets Our choice of the datasets follows (Feldman et al., 2021).
We provide information about these datasets in Table 3.
[Table]
All datasets are standardized so that inputs and labels have mean 0 and standard deviation 1 , and split into (train, cal, recal, test) with size 70%, 10%, 10%, 10% (varying with the random seed).
Base predictor and optimization
Our network architecture is a 3-layer MLP with width 64 and output dimension 2 (for the lower and upper quantile).
We use momentum SGD with initial learning rate 10 − 3 and momentum 0 .
9 , batch-size 1024, and run the optimization for a max of 10000 epochs.
A 10x learning rate decay is performed if the validation loss on D cal has not decreased in 10 epochs, and we stop the learning whenever the learning rate decay happens for 3 times.
The loss function used in training (cid:98) F = [ (cid:98) f
lo , (cid:98) f
hi ] is the summed pinball loss of level α/ 2 for (cid:98)
f lo and 1 − α/ 2 for (cid:98) f
hi , following (Romano et al., 2019):
[Equation]
where for any β ∈ (0 , 1) , (cid:96) β pinball is the pinball loss at level β :
[Equation]
Optimization details for CP-Gen-Recal For the conformal quantile ﬁnetuning procedure with our CP-Gen-Recal , we rewrite the miscoverage loss for the quantile-based prediction interval as
[Equation]
(In practice our θ also includes a trainable bias same as the original top linear layer; here we abuse notation slightly to allow easier presentation.)
We approximate the right-hand side above with the hinge loss to obtain the formulation (8).
To solve that optimization problem, we use SGD on ( θ, t ) with learning rate 0 .
01 and (ascent on) λ with learning rate 0 .
The batch-size here is 256 and the number of episodes is 1000.
To ensure t > 0 we use a log parametrization for t .
Finally, t recal is computed by the reconformalization step in Algorithm 2 on D recal .
E. M ULTI - OUTPUT REGRESSION
Datasets We generate ofﬂine datasets consisting of (state, action, next state) pairs within RL tasks within the OpenAI Gym (Brockman et al., 2016).
For each task, the data is generated by executing a medium-performing behavior policy that is extracted from standard RL training runs.
All tasks are continuous state and continuous action.
Table 4 summarizes the state and action dimension, along with the reward of the policies used for generating the data.
All datasets contain 200K examples.
All datasets are standardized so that inputs and labels have mean 0 and standard deviation 1 , and split into (train, cal, recal, test) with size 70%, 10%, 10%, 10% (varying with the random seed).
[Table]
Base predictor and optimization
Our network architecture is a 3-layer MLP with width 64, input dimension d in = d S + d A , and output dimension d out = d S .
We use momentum SGD with initial learning rate 10 − 3 , momentum 0 .
9 , and batch-size 512.
We run the optimization for 1000 epochs with a 10x learning rate decay at epoch 500.
The loss function for training the network is the standard MSE loss.
Optimization details for CP-Gen-Recal For the conformal quantile ﬁnetuning procedure with our CP-Gen-Recal , we rewrite the miscoverage loss for the box-shaped prediction set as
[Equation]
where we recall u ∈ R d out is the learnable parameter within the initial optimization stage of
CP-Gen-Recal as discussed in Section 5.2.
We approximate the right-hand side above with the hinge loss to obtain the formulation (8).
To solve that optimization problem, we use SGD on ( θ, t ) with learning rate 0 .
01 and (ascent on) λ with learning rate 0 .
The batch-size here is 1024 and the number of episodes is 1000.
To ensure u > 0 we use a log parametrization for u .
For the reconformalization step, we keep the (relative) ratios of the (cid:98) u obtained above (as the (cid:98) θ ), and then reconformalize an additional t recal > 0 on D recal via the proportional reshaping of (9).
E. D ETAILS FOR F IGURE
Figure 1 is obtained on one run of our conformal quantile ﬁnetuning experiment on the MEPSdataset, and illustrates the coverage-efﬁciency tradeoff.
Both ﬁgures there compute the coverage and length on the (unseen) test set D test , for better illustration.
Figure 1 Left plots the family [ (cid:98)
f lo
[Equation]
x ) − t, (cid:98) f
hi ( x ) + t ] used by Conformalized Quantile Regression.
Figure 1 Right plots the family
[Equation]
The speciﬁc function class of θ shown in the thinner lines is a ﬁnite set of linear interpolations of the original (cid:98) θ 0 obtained in QR and the new (cid:98) θ obtained by conformal quantile ﬁnetuning, with combination weights within {− 0 .
3 , − 0 .
The shaded region is then obtained by ﬁlling in the area.
F R ESULTS FOR LABEL PREDICTION SETS ON I MAGE N ET
Here we present the ImageNet label prediction set experiment abbreviated in Section 5.
Dataset and model We take K = 9 large-scale pretrained neural networks on the ImageNet training set (Deng et al., 2009).
Our models are { ResNeXt101, ResNet152, ResNet101, DenseNet161,
ResNet18, ResNet50, VGG16, Inception, ShufﬂeNet } , similar as in (Angelopoulos et al., 2020).
We then consider task of constructing label prediction sets with valid coverage and small cardinality.
We train and test out conformal procedures on the following two datasets, neither seen by the pretrained models:
(1) ImageNet-Val: The original validation set of ImageNet with 50000 images.
We randomly split (varying with seed) this into | D cal | = 10000 , | D recal | = 10000 , and | D test | = 30000 .
(2) ImageNet-V2 (Recht et al., 2019):
A new validation set following the roughly the same collection routines of the original images in ImageNet, however believed to have a mild distribution shift and thus slightly harder for classiﬁers pretrained on ImageNet.
This dataset contains 10000 images, which we randomly split (varying with seed) into | D cal | = 4000 , | D recal | = 1000 , and | D test | = 5000 .
Methods for learning prediction sets Our constructions of the prediction sets are based on the
Least Ambiguous Set-Valued Classiﬁer (LAC) method of (Sadinle et al., 2019), which turns any base predictor p where p ( ·| x ) denotes the predicted distribution of the L = 1000 labels into a prediction
set C t ( x ) via
C
[Equation]
: p
[Equation]
y | x ) > t } , where t is found by conformal prediction.
We consider learning a valid prediction set with smaller set size by ﬁnding an optimized ensemble weight of the K base predictors using our CP-Gen-Recal algorithm.
This means we learn prediction sets of the form
[Equation]
(cid:88) k
[Equation]
θ
[Equation] [Equation]
t (cid:41) ,
[Table]
where { p k } k ∈ [ K ] are the base predictors.
Our CP-Gen-Recal algorithm (and its practical implementation (8)) would solve a primal-dual optimization problem with the efﬁciency loss and hinge approximate coverage constraint to optimize ( θ, t ) .
However, here the efﬁciency loss we care about (the cardinality) is non-differentiable.
We make a further approximation by considering the L qq norm with q = 0 .
5 as the surrogate efﬁciency loss:
[Equation]
with the intuition that the q → 0 limit is exactly the cardinality of C θ,t ( x i ) .
Our ﬁnal optimization problem is then
[Equation]
(cid:96) hinge ( p θ ( y i | x i ) − t ) .
We solve this by SGD on ( θ, t ) and (ascent on) λ , with the softmax parameterization for θ ( θ ∈ ∆ K as an ensemble weight is a probability distribution) and log parametrization for t > 0 .
The learning rate is 10 − 2 for ( θ, t ) and 10 − 4 for λ .
We perform this optimization for 500 epochs over D cal with batch-size 256 for ImageNet-Val and 64 for ImageNet-V2.
After we obtain the iterates (cid:110)(cid:98) θ j (cid:111) (where j denotes the epoch count), we perform a further iterate selection of ﬁrst re-computing the (cid:98) t ( (cid:98) θ j ) by conformalizing on D cal , and then choosing the iterate j with the best average set size also on D cal , before feeding it into the reconformalization step with
D recal .
As the D recal is only used in the reconformalization step, such as method still guarantees valid coverage like the original Algorithm 2.
We compare our above algorithm against two baselines: conformalizing each individual model and reporting the best one, or conformalizing the uniform ensemble (which uses weights θ unif = 1 K 1 K ).
For these two baselines, for fairness of comparison, we allow them to use the whole D cal ∪ D recal as the calibration set, as their construction (apart from pre-training) is not data-dependent.
Results Table 5 shows that our algorithm is able to learn label prediction sets with valid coverage and improved set sizes over the baselines.
This demonstrates the advantage of our method even in applications where the efﬁciency loss (here set size) is non-differentiable and needs to be further approximated to allow gradient-based algorithms.
G A BLATION STUDIES G. C ONFORMAL QUANTILE FINETUNING
We report ablation results for the conformal quantile ﬁnetuning problem with nominal coverage level 1 − α ∈ { 80% , 95% } , and otherwise exactly the same setup as Section 5.1.
The conclusions are qualitatively the same as the 90% version presented in Table 1.
G. M ULTI - OUTPUT REGRESSION
We report ablation results for the multi-output regression problem with nominal coverage level 1 − α ∈ { 80% , 95% } , and otherwise exactly the same setup as Section 5.2.
The conclusions are qualitatively the same as the 90% version presented in Table 2, except for one dataset at level 95% .
[Table]
[Table]
G. C OMPARISON OF CP-Gen AND CP-Gen-Recal
We compare the performance of CP-Gen and CP-Gen-Recal on the multi-output regression tasks using the same setup as Section 5.2.
Recall that the vanilla CP-Gen optimizes both ( (cid:98) θ, (cid:98) t ) on D cal (we additionally reconformalize (cid:98) t on the same D cal to address the potential bias in (cid:98) t brought by the approximate optimization (8))
[Table]
whereas our main CP-Gen-Recal algorithm optimizes (cid:98) θ on D cal and reconformalizes (cid:98) t recal on D recal .
Table 10 reports the results.
Observe that, except for the volume on one dataset (Humanoid), there is no signiﬁcant difference in both the coverage and the volume for the two methods.
For practice we recommend CP-Gen-Recal whenever the exact coverage guarantee is important, yet this result shows that—perhaps originating from the fact that here n cal = 20000 is large—the coverage (generalization error) of CP-Gen is also nearly valid, which may be better than what our Propositionsuggests.
[Table] [Table]
[Table]
H A DDITIONAL EXPERIMENTS AND ANALYSES H. C ONDITIONAL COVERAGE OF CP-Gen-Recal
We analyze the improved length prediction intervals learned by CP-Gen-Recal (Section 5.1) by evaluating its conditional coverage metrics and comparing with the baseline CQR method.
As conditional coverage is hard to reliably estimate from ﬁnite data, we consider two proxy metrics proposed in (Feldman et al., 2021) that measure the independence between length and indicator of coverage :
• The correlation coefﬁcient (Corr) between the following two random variables: the interval size L = length( C ( X )) and the indicator of coverage V = 1 (cid:110) Y ∈ (cid:98) C ( X ) (cid:111) .
A (population) correlation of 0 is a necessary (but not sufﬁcient) condition of perfect conditional coverage (Feldman et al., 2021).
Here we measure the absolute correlation, which is smaller the better.
HSIC:
A more sophisticated correlation metric between L and V that takes into account nonlinear correlation structures.
A (population)
HSIC of 0 is a necessary and sufﬁcient condition of the independence between L and V .
We estimate HSIC on the ﬁnite test data using the method in (Feldman et al., 2021).
Table 11 reports the results.
Observe that while our CP-Gen-Recal improves the length, it achieves worse (higher) Correlation/HSIC than the baseline CQR , which is expected as
length and
conditional coverage often come as a trade-
off.
[Table]
H.
A LTERNATIVE TWEAKS FOR CQR BASELINE
Here we test two additional tweaked versions of the CQR baseline in the prediction interval experiment of Section 5.1:
• CQR - D train ∪ D cal : Use dataset D train ∪ D cal for training the base quantile regressor, then conformalize on D recal .
• CQR-PinballFt : Train the base quantile regressor on D train , and ﬁnetune the last linear layer on D cal using the pinball loss (same as training), and conformalize on D recal .
Optimization details about these two methods are described in Section H.2.1.
Result Table 12 reports the results for these two tweaked baselines, in comparison with our original baseline CQR as well as our proposed QR + CP-Gen-Recal .
Observe that using more training data ( CQR - D train ∪ D cal ) improves the length slightly on some datasets but not all.
In contrast, CQR-PinballFt is unable to improve either the pinball loss or the length over the base
CQR (observe that CQR-PinballFt uses the same set of training data as CQR - D train ∪ D cal but uses a less expressive model in the ﬁnetuning stage).
Overall, on almost all datasets (except for kin8nm), our CP-Gen-Recal still achieves the best length.
[Table]
H.2. O PTIMIZATION DETAILS
CQR - D train ∪ D cal : Our original CQR baseline used D cal for monitoring validation loss and automatically determining the early stopping (cf.
Section E.1), and D recal for conformalization.
To optimize on D train ∪ D cal , we do not use automatic learning rate decay and early stopping, but instead manually picked the number of epochs and corresponding learning rate decay schedule that is close to average runs of the CQR method on each dataset.
This choice ensures that our new baseline still gets to use see the exact same amount of data for conformalizing ( D recal ) and testing ( D test ), and has a optimization setup as close as possible the original CQR baseline.
More concretely, we optimize for 800 epochs for { MEPS 19, MEPS 20, MEPS 21 } , 1500 epochs for { Facebook 1, Facebook 2, blog data } , 6000 epochs for kin8nm, 350 epochs for naval, andepochs for bio.
For all datasets, the learning rate decays by 10x twice, at 90% and 95% of the total epochs.
CQR-PinballFt : We ﬁnetuned on D cal with 1000 epochs and batch size 256.
The learning rate was chosen within { 10 − 2 , 10 − 3 } and the results are not too different for these two choices (length difference is within 0 .
010 for these two choices, and there is no overall winner).
We presented the results with learning rate 10 − 3 .
H. A DDITIONAL Max-score-Conformal BASELINE FOR MULTI - OUTPUT REGRESSION
We test one additional baseline method for the multi-output regression experiment in Section 5.2:
• Max-score-Conformal : Here we consider the hypercube -shaped predictor
[Equation]
and use conformal prediction on D cal ∪ D recal to compute a conformalized (cid:98)
t
and the ﬁnal
prediction set C (cid:98) t ( x ) .
In other words, we perform standard conformal prediction with score function (cid:107) y − (cid:98) f
( x ) (cid:107) ∞ .
We remark that both the Coord-wise-Recal and the Max-score-Conformal baseline methods are special instances of CP-Gen-Recal with some ﬁxed θ 0 :
In our parametrization, u = ( θ, t ) , θ determines the shape (i.e. relative ratios between the u (cid:48)
i s ) whereas t determines the size .
Therefore, Max-score-Conformal can be thought of as choosing θ 0 to be the all-ones ratio (i.e. hypercube-shaped), whereas Coord-wise-Recal can be thought of as choosing θ 0 from a coordinate-wise one dimensional conformal prediction.
Result Table 13 reports the result for Max-score-Conformal .
Compared with the existing baseline Coord-wise-Recal , Max-score-Conformal achieves better volume on the
Cartpole dataset but worse volume on almost all other datasets (except for Swimmer where their volumes are similar).
Further, note that Max-score-Conformal achieves signiﬁcantly higher volumes for certain datsets (Ant, Humanoid).
Our inspection shows that this due to the fact that there are a certain number of hard-to-predict state dimensions (and many other easy-to-predict state dimensions) for these two datasets.
Therefore, Coord-wise-Recal which builds on
Coord-wise adapts to this structure and uses only a high length on these dimensions only, whereas the Max-score-Conformal method pays this max conformal score on all dimensions to yield an unnecessarily high volume.
We remark that our CP-Gen-Recal still performs signiﬁcantly better than both baselines.
[Table]
H. 100 RANDOM SEEDS AND STANDARD DEVIATION
Here we repeat the experiments in Section 5.1 & 5.2 with 100 random seeds and the exact same setups.
We report the mean and standard deviations in Table 14 for the prediction intervals experiment (Section 5.1), Table 15 for the mean and Table 16 for the standard deviation for the multi-output regression experiment (Section 5.2)
[Table] [Table]
[Table] [Table]