DICTION WITH GENERAL FUNCTION CLASSES
Year: 2022

Abstract:
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 efficiency (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 efficiency 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 finding the most efficient 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 efficiency 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 efficiency significantly 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 classification.

heading: 1 INTRODUCTION
text: 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.

heading: 1.1 RELATED WORK
text: 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).

heading: 2 PRELIMINARIES
text: 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.

heading: 2.1 CONFORMAL PREDICTION
text: 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 briefly 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 Dtrain), conformal prediction outputs a prediction interval
Ct̂(x) := [ f(x)− t̂, f(x) + t̂ ] , (1)
where t̂ ∈ R≥0 is chosen as the (1 − α)-quantile1 of |y − f(x)| on a calibration dataset Dcal with size ncal := |Dcal| using the following conformalization step:
t̂ = d(1− α)ncale -th largest of {|yi − f(xi)|}ncali=1 . (2) The main guarantee for the learned interval Ct̂ is that it achieves a (1 − α) coverage guarantee of the form PDcal,(X,Y )(Y ∈ Ct̂(X)) ≥ 1− α (Lei et al., 2018, Theorem 2.2). The proof relies on the exchangeability between the scores {|yi − f(xi)|}ncali=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 t̂ ∈ R defined in (2) is the solution to the following constrained ERM problem
minimize t≥0
L̂eff(Ct) := 1
ncal ∑ i∈Dcal `eff(Ct; (xi, yi)) = 2t
subject to L̂coverage(Ct) := 1
ncal ∑ i∈Dcal 1 {yi /∈ Ct(xi)} ≤ α. (3)
1Technically (2) requires the d(1− α)(nrecal + 1)e-th largest element to guarantee valid coverage (Vovk et al., 2005); here we choose the close d(1− α)nrecale-th largest to allow the following insight.
Algorithm 1 Conformal Prediction with General Function Class (CP-Gen) Input: Class of prediction sets C = {Cθ,t}θ∈Θ,t∈T ; target miscoverage level α ∈ (0, 1); ε0 ≥ 0.
Efficiency loss `eff ; Calibration dataset Dcal with size ncal. 1: Solve the following constrained ERM problem on dataset Dcal (with relaxation parameter ε0):
(θ̂, t̂)← arg min θ∈Θ,t∈T
L̂eff(Cθ,t) := 1
ncal ∑ i∈Dcal `eff(Cθ,t(xi), yi)
subject to L̂coverage(Cθ,t) := 1
ncal ∑ i∈Dcal 1 {yi /∈ Cθ,t(xi)} ≤ α+ ε0. (5)
Output: Prediction set Cθ̂,t̂.
Above, `eff(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 efficiency due to the coverage-efficiency 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.

heading: 3 CONFORMAL PREDICTION WITH GENERAL FUNCTION CLASSES

heading: 3.1 ALGORITHM
text: 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
C := {Cθ,t(x) : θ ∈ Θ, t ∈ T ⊂ R}, (4)
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 finding the smallest interval among C subject to valid coverage on dataset Dcal. Compared with vanilla conformal prediction, Algorithm 1 allows more general tasks with an arbitrary function class and efficiency loss; for example, this encompasses several recent algorithms such as finite 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.

heading: 3.2 THEORY
text: An important theoretical question about CP-Gen is whether it achieves coverage and efficiency guarantees on the population (test data). This section showcases that, by standard generalization arguments, CP-Gen achieves approximate validity and near-optimal efficiency whenever function class is low-capacity in a certain sense. We remark that our experiments use the modified 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{eff,coverage}(Cθ,t) := E[`{eff,coverage}(Cθ,t; (X,Y ))] denote the population coverage and efficiency lossesfor any (θ, t). We define the following uniform concentration quantities:
εeff := supθ∈Θ,t∈T ∣∣∣L̂eff(Cθ,t)− Leff(Cθ,t)∣∣∣, (6) εcoverage := supθ∈Θ,t∈T ∣∣∣L̂coverage(Cθ,t)− Lcoverage(Cθ,t)∣∣∣. (7)
Algorithm 2 Conformal Prediction with General Fn. Class and Recalibration (CP-Gen-Recal) Input: Class of prediction sets C = {Cθ,t}θ∈Θ,t∈T ; target miscoverage level α ∈ (0, 1); ε0 ≥ 0.
Efficiency loss `eff ; Calibration datasets Dcal, Drecal with size ncal, nrecal. 1: Run Algorithm 1 on dataset Dcal (with relaxation parameter ε0) to obtain (θ̂, t̂). 2: Keep θ̂, and reconformalize t ∈ T on the recalibration dataset Drecal:
t̂recal ← inf { t ∈ T : yi ∈ Cθ̂,t(xi) for at least d(1− α)(nrecal + 1)e examples (xi, yi) ∈ Drecal } .
Output: Prediction set Cθ̂,t̂recal .
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 Dcal 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 setCθ̂,t̂ learned by Algorithm 1 satisfies
(a) (Approximately valid population coverage) We have
Lcoverage(Cθ̂,t̂) ≤ α+ ε0 + εcoverage,
i.e. the population coverage of Cθ̂,t̂ is at least 1− α− (ε0 + εcoverage).
(b) (Near-optimal efficiency) Suppose ε0 ≥ εcoverage, then we further have Leff(Cθ̂,t̂) ≤ inf(θ,t)∈Θ×T
Lcoverage(Cθ,t)≤α
Leff(Cθ,t) + 2εeff ,
i.e. Cθ̂,t̂ achieves 2εeff -near-optimal efficiency 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 efficiency if the concentration terms εeff and effcoverage 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 {εeff , effcoverage} ≤√
Comp(C)/ncal 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−α− √ Comp(C)/ncal approximate coverage guarantee and √ Comp(C)/ncal near-optimal efficiency guarantee. In particular, our Proposition 4 recovers the coverage guarantee for the finite-class selection algorithm of (Yang & Kuchibhotla, 2021, Theorem 1) though our efficiency guarantee is more general.
We remark that both examples above contain important applications. The finite class contains e.g. optimizing over a K-dimensional hyperparameter to use via grid search, with e.g. (1/δ)K confidence 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 efficiency) more broadly in practice, for instance whenever it learns K ncal parameters.

heading: 3.3 ALGORITHM WITH VALID COVERAGE VIA RECONFORMALIZATION
text: Although CP-Gen enjoys theoretical bounds on the coverage and efficiency, 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 datasetDrecal (e.g. a further data split), which guarantees valid finite-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 efficiency bound builds upon the result for CP-Gen (Proposition 2(b)) and handles additional concentration terms from the reconformalization step. Proposition 3 (Coverage and efficiency guarantee for CP-Gen-Recal; Informal version). The prediction set Cθ̂,t̂recal learned by Algorithm 2 achieves (1 − α) finite-sample coverage: PDrecal,(X,Y )(Y ∈ Cθ̂,t̂recal) ≥ 1 − α. Further, it achieves O(εeff + εcoverage + 1/ √ nrecal) nearoptimal efficiency under additional regularity assumptions.

heading: 4 DIFFERENTIABLE OPTIMIZATION
text: Our (meta) algorithms CP-Gen and CP-Gen-Recal involve solving the constrained ERM problem (5). One feasible case is when Θ is finite and small, in which we enumerate all possible θ ∈ Θ and find the optimal t̂ for each θ efficiently using quantile computation. However, this optimization is significantly more challenging when the underlying parameter set Θ is continuous and we wish to jointly optimize over (θ, t): The coverage loss 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 first rewrite each individual coverage loss into the form
1 {y /∈ Cθ,t(x)} = 1 {sθ,t(x, y) < 0} = `01(sθ,t(x, y)). where `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 `01(z) by the hinge loss `hinge(z) = [1− z]+ which is (almost everywhere) differentiable with a non-trivial gradient. We find the hinge loss to perform better empirically than alternatives such as the logistic loss.
To deal with the (modified) 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:
min θ,t max λ≥0
L̂eff(Cθ,t) + λ [ L̂hinge(Cθ,t)− α ] + , (8)
where L̂hinge(Cθ,t) := 1ncal ∑ncal i=1 `hinge(sθ,t(xi, yi)) is the empirical hinge loss on the calibration datasetDcal. Our final practical implementation of (5) solves the problem (8) by Stochastic Gradient Descent-Ascent (with respect to (θ, t) and λ) to yield an approximate solution (θ̂, t̂). We remark that in our experiments where we use the reconformalized version CP-Gen-Recal, we only keep the θ̂ obtained from (8) and perform additional reconformalization to compute 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 sacrifice in terms of the efficiency. However, our experiments in Section 5 show that such an implementation can still improve the efficiency over existing approaches in practice, despite the approximation.

heading: 5 EXPERIMENTS
text: 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.

heading: 5.1 IMPROVED PREDICTION INTERVALS VIA CONFORMAL QUANTILE FINETUNING
text: Setup We consider regression tasks in which we use quantile regression (pinball loss) to train a base quantile predictor F̂ (x) = [f̂lo(x), f̂hi(x)] = θ̂>0 Φ̂(x) on Dtrain (with learning rate decay by monitoring validation loss on Dcal). Here Φ̂ : X → Rdh is the learned representation function, θ̂0 ∈ Rdh×2 is the last linear layer (dh denotes the last hidden dimension), and f̂lo, f̂hi are the learned {lower, upper} quantile functions (see Appendix E.1 for more details on the training procedure). Given F̂ , we learn a baseline prediction interval of the form [f̂lo(x) − t, f̂hi(x) + t] on Drecal via Conformalized Quantile Regression (CQR) (Romano et al., 2019).
We then attempt to improve the length over CQR by conformal quantile finetuning: Fix the representation function Φ̂ and finetune the linear layer θ using our CP-Gen-Recal algorithm, so that Cθ,t(x) = [θ > loΦ̂(x)− t, θ>hiΦ̂(x)+ t] (where θ = [θlo, θhi]). We learn a new θ̂ onDcal via (8) (where `eff is chosen as the length), and then compute t̂recal on Drecal as in Algorithm 2.
We perform the above on 9 real-world regression datasets with a 3-layer MLP with width dh = 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 finetuned 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 θ̂0Φ̂(x) as well as the fine-tuned θ̂>Φ̂(x) on the test set Dtest. Intriguingly, our conformal finetuning 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).

heading: 5.2 MINIMUM-VOLUME PREDICTION SETS FOR MULTI-OUTPUT REGRESSION
text: 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 dout > 1. We first learn a based predictor f̂ : Rd → Rdout by minimizing the MSE loss on Dtrain. We then learn a box-shaped prediction set of the form Cu(x) = ∏dout i=1 [f̂i(x)− ûi, f̂i(x) + ûi] by one of the following methods:
• (Coord-wise): Each ûi is obtained by vanilla conformalization (2) over the i-th output coordinate on Dcal ∪ Drecal. To guarantee 1 − α coverage, each coordinate is conformalized at level 1− α/dout, motivated by the union bound.
Our datasets are a collection of next-state prediction tasks with multi-dimensional continuous states in offline 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 ‖y−f̂(x)‖∞, equivalent to a hypercube-shaped predictor) in Appendix H.3, which we find 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 ∈ Rdout data-dependently using our CP-Gen-Recal is indeed more flexible 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 finding 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.

heading: 6 CONCLUSION
text: This paper proposes Conformal Prediction with General Function Class, a conformal prediction algorithm that optimizes the efficiency metric subject to valid coverage over a general function class of prediction sets. We provide theoretical guarantees for its coverage and efficiency 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 quantification tasks.

heading: A PROOF OF PROPOSITION 1
text: Recall that Ct(x) = [f(x) − t, f(x) + t] satisfies `eff(Ct; (x, y)) = length(Ct(x)) = 2t for any (x, y). Also, we have
1 {y /∈ Ct(x)} = 1 {y /∈ [f(x)− t, f(x) + t]} = 1 {|y − f(x)| > t} , or equivalently 1 {y ∈ Ct(x)} = 1 {|y − f(x)| ≤ t}. Therefore, problem (3) is equivalent to
minimize t≥0 t
subject to 1− L̂coverage(Ct) := 1
ncal ncal∑ i=1 1 {|yi − f(xi)| ≤ t} ≥ 1− α.
By definition of quantiles, this problem is solved at
t̂ = d(1− α)ncale -th largest element of {|yi − f(xi)|}ncali=1 , which is the desired result.

heading: B PROOF OF PROPOSITION 2
text: (a) As Cθ̂,t̂ solves problem (5), it satisfies the constraint L̂coverage(Cθ̂,t̂) ≤ α+ ε0. Therefore,
Lcoverage(Cθ̂,t̂) = L̂coverage(Cθ̂,t̂)︸ ︷︷ ︸ ≤α+ε0 +Lcoverage(Cθ̂,t̂)− L̂coverage(Cθ̂,t̂)
≤ α+ ε0 + sup (θ,t)∈Θ×T ∣∣∣Lcoverage(Cθ,t)− L̂coverage(Cθ,t)∣∣∣ = α+ ε0 + εcoverage.
(b) Suppose εcoverage ≤ ε0. Taking any (θ, t) ∈ Θ× T such that Lcoverage(Cθ,t) ≤ α, we have L̂coverage(Cθ,t) ≤ Lcoverage(Cθ,t) + εcoverage ≤ α+ εcoverage ≤ α+ ε0.
This shows that (θ, t) lies within the constraint set of problem (5). Thus as (θ̂, t̂) further minimizes the loss L̂eff within the constraint set, we have L̂eff(Cθ̂,t̂) ≤ L̂eff(Cθ,t). This shows that
Leff(Cθ̂,t̂)− Leff(Cθ,t)
≤ L̂eff(Cθ̂,t̂)− L̂eff(Cθ,t)︸ ︷︷ ︸ ≤0 +2 sup (θ,t)∈Θ×T ∣∣∣L̂eff(Cθ,t)− Leff(Cθ,t)∣∣∣ ≤ 2εeff .
As the above holds simultaneously for all (θ, t) ∈ Θ × T with at most α coverage loss, taking the sup over the left-hand side yields
Leff(Cθ̂,t̂)− inf(θ,t)∈Θ×T Lcoverage(Cθ,t)≤α Leff(Cθ,t) ≤ 2εeff .
C EXAMPLES OF GOOD GENERALIZATION FOR CP-Gen
We provide two concrete examples where the concentration terms εeff 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 efficiency. Assumption A (Bounded and Lipschitz efficiency loss). The loss function `eff satisfies A1 (Bounded loss). |`eff(Cθ,t; (x, y))| ≤M for all (θ, t) ∈ Θ× T and all (x, y) ∈ X × Y .
A2 (t-Lipschitzness). t 7→ `eff(Cθ,t; (x, y)) is LT -Lipschitz for all (θ, x, y) ∈ Θ×X × Y . Assumption B (Bounded T ). The parameter space T ⊂ R is bounded: supt∈T |t| ≤ BT .
C.1 FINITE CLASS
Proposition 4 (Finite class). Suppose Θ is a finite set (NΘ := |Θ| < ∞), and suppose Assumptions A1, A2, B hold. Then, we have with probability at least 1− δ that
εcoverage ≤ C √ log(NΘ/δ)/ncal and εeff ≤ C · [ M √ log(NΘ/δ) + LT BT ] / √ ncal,
where C > 0 is an absolute constant.
Proof. We first bound εcoverage. Fix any θ ∈ Θ, define tθ(x, y) := inf {t ∈ T : Cθ,t(x) 3 y} (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
sup t∈T ∣∣∣L̂coverage(Cθ,t)− Lcoverage(Cθ,t)∣∣∣ = sup
t∈T ∣∣∣∣∣ 1ncal ncal∑ i=1 1 {yi /∈ Cθ,t(xi)} − P(Y /∈ Cθ,t(X)) ∣∣∣∣∣ = sup
t∈T ∣∣∣∣∣ 1ncal ncal∑ i=1 1 {yi ∈ Cθ,t(xi)} − P(Y ∈ Cθ,t(X)) ∣∣∣∣∣ = sup
t∈T ∣∣∣∣∣ 1ncal ncal∑ i=1 1 {tθ(x, y) ≤ t} − P(x,y)(tθ(X,Y ) ≤ t) ∣∣∣∣∣ = sup
t∈T ∣∣∣F̂θ(t)− Fθ(t)∣∣∣, where we have defined Fθ : R → [0, 1] as the CDF of the random variable tθ(X,Y ) and similarly F̂θ as the empirical CDF of the same random variable over the finite dataset Dcal. Applying the Dvoretzky-Kiefer-Wolfowitz (DKW) inequality (Massart, 1990, Corollary 1) yields that
sup t∈T
∣∣∣F̂θ(t)− Fθ(t)∣∣∣ ≤ √ log(2/δ)
2ncal
with probability at least 1− δ. Now, taking the union bound with respect to θ ∈ Θ (where for each θ we plug in tail probability δ/2NΘ) we get that with probability at least 1− δ/2,
εcoverage = sup θ∈Θ sup t∈T ∣∣∣L̂coverage(Cθ,t)− Lcoverage(Cθ,t)∣∣∣ = sup θ∈Θ sup t∈T ∣∣∣F̂θ(t)− Fθ(t)∣∣∣ ≤ √ log(4NΘ/δ) 2ncal ≤ C √ log(NΘ/δ) ncal .
(11)
for some absolute constant C > 0.
We next bound εeff . Fix any θ ∈ Θ. We have by standard symmetrization argument that E [ sup t∈T ∣∣∣L̂eff(Cθ,t)− Leff(Cθ,t)∣∣∣] = E [ sup t∈T ∣∣∣∣∣ 1ncal ncal∑ i=1 `eff(Cθ,t; (xi, yi))− E[`eff(Cθ,t; (X,Y ))] ∣∣∣∣∣ ]
≤ 2E(xi,yi),εi [ sup t∈T ∣∣∣∣∣ 1ncal ncal∑ i=1 εi`eff(Cθ,t; (xi, yi)) ∣∣∣∣∣ ]
(i) ≤ 2LT · Eεi [ sup t∈T ∣∣∣∣∣ 1ncal ncal∑ i=1 εi · t ∣∣∣∣∣ ] = 2LT · Eεi [∣∣∣∣∣ 1ncal ncal∑ i=1 εi ∣∣∣∣∣ · supt∈T |t| ]
(ii)
≤ 2LT ·BT · Eεi [∣∣∣∣∣ 1ncal ncal∑ i=1 εi ∣∣∣∣∣ ] (iii) ≤ 2LT ·BT / √ ncal.
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 [∣∣∣ 1ncal ∑ncali=1 εi∣∣∣] ≤( Eεi [( 1 ncal ∑ncal i=1 εi )2])1/2 = 1/ √ ncal. (Above εi iid∼ Unif({±1}) are Rademacher variables.)
Next, as each loss |`eff(Cθ,t; (x, y))| ≤M by Assumption A1, the random variable
sup t∈T ∣∣∣L̂eff(Cθ,t)− Leff(Cθ,t)∣∣∣ satisfies the M/ncal finite-difference property. Therefore by McDiarmid’s Inequality, we have with probability at least 1− δ that
sup t∈T ∣∣∣L̂eff(Cθ,t)− Leff(Cθ,t)∣∣∣ ≤ E
[ sup t∈T ∣∣∣L̂eff(Cθ,t)− Leff(Cθ,t)∣∣∣]+ √ M2 log(1/δ) 2ncal
≤ C · LT BT +M
√ log(1/δ)
√ ncal
.
Finally, by union bound over θ ∈ Θ (where we plug in δ/2NΘ as tail probability into the above), we have with probability at least 1− δ/2 that
εeff = sup θ∈Θ sup t∈T ∣∣∣L̂eff(Cθ,t)− Leff(Cθ,t)∣∣∣ ≤ C · LT BT +M √
log(NΘ/δ)√ ncal .
(12)
(11) together with (12) is the desired result.
C.2 VC/RADEMACHER CLASS
Next, for any class C, let VC(C) := VC({(x, y) 7→ 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 √ (K + 1 + log(1/δ))/ncal.
(b) Suppose Assumption A1 holds. Then we have with probability at least 1− δ/2 that εeff ≤ C [ Reffncal(C) + √ M2 log(1/δ)/ncal ] ,
where Reffncal(C) := E(xi,yi),εi [ sup(θ,t)∈Θ×T ∣∣∣ 1ncal ∑ncali=1 εi`eff(Cθ,t; (xi, yi))∣∣∣] is the Rademacher complexity of the class C with respect to `eff (above εi iid∼ Unif({±1})).
Proof. (a) By assumption, the class of Boolean functions {(x, y) 7→ 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
εcoverage = sup (θ,t)∈Θ×T ∣∣∣∣∣ 1ncal ncal∑ i=1 1 {yi /∈ Cθ,t(xi)yi} − P(Y /∈ Cθ,t(X)) ∣∣∣∣∣ ≤ C √ K + 1
ncal +
√ log(2/δ)
2ncal ≤ C
√ K + 1 + log(1/δ)
ncal .
(b) We have by standard symmetrization argument that (below εi iid∼ Unif({±1}) denote Rademacher variables)
E[εeff ] = E [ sup
θ∈Θ,t∈T
∣∣∣L̂eff(Cθ,t)− Leff(Cθ,t)∣∣∣]
= E [ sup
θ∈Θ,t∈T ∣∣∣∣∣ 1ncal ncal∑ i=1 `eff(Cθ,t; (xi, yi))− E[`eff(Cθ,t; (X,Y ))] ∣∣∣∣∣ ]
≤ 2E(xi,yi),εi
[ sup
θ∈Θ,t∈T ∣∣∣∣∣ 1ncal ncal∑ i=1 εi`eff(Cθ,t; (xi, yi)) ∣∣∣∣∣ ] = 2Rn(C).
Further by Assumption A1, the quantity εeff satisfies M/ncal bounded-difference, so applying McDiarmid’s Inequality gives that with probability at least 1− δ/2,
εeff ≤ E[εeff ] +
√ 2M2 log(2/δ)
ncal ≤ C
[ Reffncal(C) + √ M2 log(1/δ)
n
] .
C.3 CASE STUDY: LINEAR CLASS
In this section, we study prediction intervals with a specific linear structure and show that it satisfies 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
Cθ,t(x) = [θ >Φlo(x)− tσ(x), θ>Φhi(x) + tσ(x)], (13)
where θ ∈ Θ ⊂ RK , Φhi,Φlo : X → RK 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 find an optimal linear function of this representation that yields the shortest prediction interval (with fixed 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θ∈Θ ‖θ‖ ≤ BΘ, supx∈X ‖Φ(x)‖ ≤ BΦ, supx∈X σ(x) ≤ Bσ , and supt∈T |t| ≤ BT .
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 efficiency loss to be the length of the interval: `eff(C; (x, y)) := length(C(x)). Then, we have with probability at least 1 − δ (over the calibration dataset Dcal) that
εcoverage ≤ C
√ K + 1 + log(1/δ)
ncal , and εeff ≤ C[BΘBΦ +BT Bσ] ·
√ log(1/δ)
ncal ,
where C > 0 is an absolute constant.
Proof. We verify the conditions of Proposition 5. First, we have 1 {y /∈ Cθ,t(x)} = 1 { max { y − θ>Φhi(x), θ>Φlo(x)− y } > tσ(x) } .
The set within the indicator above is the union of two sets { (x, y) : y − θ>Φhi(x)− tσ(x) > 0 }
and { (x, y) : θ>Φlo(x)− y − tσ(x) > 0 }
. Note that each family of sets (over (θ, t) ∈ RK ×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 ′(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 εeff , we first note that for any (x, y) ∈ X × R, |`eff(Cθ,t; (x, y))| = |length(Cθ,t(x))|
= θ>(Φhi(x)− Φlo(x)) + 2tσ(x) ≤ ‖θ‖ ‖Φhi(x)− Φlo(x)‖+ 2tσ(x) ≤ 2BΘBΦ + 2BT Bσ =: M,
and thus the boundedness assumption (Assumption A1) holds withM defined above. Next, we have the following bound on the Rademacher complexity
Reffncal(C) = E
[ sup
(θ,t)∈Θ×T ∣∣∣∣∣ 1ncal ncal∑ i=1 εi ( θ>(Φhi(xi)− Φlo(xi)) + 2tσ(xi) )∣∣∣∣∣ ]
≤ E [ sup θ∈Θ ∣∣∣∣∣ 〈 θ, 1 ncal ncal∑ i=1 εi(Φhi(xi)− Φlo(xi)) 〉∣∣∣∣∣ ] + E [ sup t∈T ∣∣∣∣∣2t · 1ncal ncal∑ i=1 εiσ(xi) ∣∣∣∣∣ ]
≤ sup θ∈Θ ‖θ‖ · E ∥∥∥∥∥ 1ncal ncal∑ i=1 εi(Φhi(xi)− Φlo(xi)) ∥∥∥∥∥ 2 1/2 + 2 sup t∈T |t| · E ( 1 ncal ncal∑ i=1 εiσ(xi) )21/2
≤ BΘ · E [ 1
ncal ‖Φhi(x1)− Φlo(x1)‖2
]1/2 + 2BT · E [ 1
ncal σ2(x1) ]1/2 ≤ C · BΘBΦ +BT Bσ√
ncal .
Applying Proposition 5(b), we get εeff ≤ C · [BΘBΦ +BT Bσ] · √
log(1/δ)/ncal with probability at least 1− δ. This is the desired bound for εeff .
D THEORETICAL GUARANTEE FOR CP-Gen-Recal
In this section we state and prove the formal theoretical guarantee for the CP-Gen-Recal algorithm (Algorithm 2).
Define 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 ′θ(t) > 0 on t ∈ T . Further, let tθ,1−α := inf {t ∈ T : Fθ(t) ≥ 1− α} denote its (1− α) quantile, then there exists some constants c0, δ0 > 0 such that
inf t∈[tθ,1−α−δ0,tθ,1−α+δ0] fθ(t) ≥ c0.
Proposition 7 (Valid coverage and near-optimal efficiency for reconformalized algorithm). The following holds for Algorithm 2:
(a) (Valid coverage) For any possible θ̂ ∈ Θ learned in Line 1 and the resulting t̂recal, we have EDrecal [ Lcoverage(Cθ̂,t̂recal) ] ≤ α, and thus PDrecal,(X,Y ) ( Y ∈ Cθ̂,t̂recal(X) ) ≥ 1− α. (b) (Efficiency) Suppose Assumptions A2 and D hold, max { εcoverage + 1/ncal, 2 √ log(1/δ)/nrecal } ≤
c0δ0 (recall the definition of εcoverage in (7)), and εcoverage ≤ ε0. Then for δ ≤ 0.5, we have with probability at least 1− δ that
Leff(Cθ̂,t̂recal) ≤ min(θ,t)∈Θ×T Lcoverage(Cθ,t)≤α Leff(Cθ,t) + 2εeff + CLT ·
εcoverage + 1 ncal + √ log(1/δ) nrecal /c0. Proof. (a) As the learned parameter θ̂ (and thus the family of nested sets Cθ̂,t) is independent of the recalibration dataset Drecal, we have that the scores tθ̂(x, y) on dataset Drecal and a new test point (X,Y ) are exchangeable given any θ̂. Therefore by (Gupta et al., 2019, Proposition 1), we have for any θ̂ ∈ Θ that
PDrecal,(X,Y ) ( Y ∈ Cθ̂,t̂recal(X) ) ≥ 1− α,
or equivalently EDrecal [ Lcoverage(Cθ̂,t̂recal) ] ≤ α.
(b) For any θ ∈ Θ, define the score function tθ(x, y) the same as in (10), and similarly define the CDF Fθ(t) := P(tθ(X,Y ) ≤ t) and its empirical counterpart F̂ calθ (t) and F̂ recalθ (t) as the finitesample version on dataset Dcal and Drecal respectively.
We first analyze t̂. By the same derivation as in (11), we have
sup t∈T
∣∣∣F̂ cal θ̂ (t)− Fθ̂(t) ∣∣∣ ≤ sup
(θ,t)∈Θ×T ∣∣∣F̂ calθ (t)− Fθ(t)∣∣∣ = εcoverage. As (θ̂, t̂) solves the constrained ERM (5) and by the assumption that `eff(Cθ,t; (x, y)) is monotone in t, we have that t̂ is the minimal value of t ∈ T such that F̂ cal
θ̂ (t) ≥ 1 − α. Therefore, (as
|Dcal| = ncal and {tθ(xi, yi)}i∈Dcal are almost surely distinct by Assumption D,) we have
1− α ≤ F̂ cal θ̂ (t̂) ≤ 1− α+ 1/ncal.
This shows that ∣∣∣Fθ̂(t̂)− Fθ̂(tθ̂,1−α)∣∣∣ = ∣∣Fθ̂(t̂)− (1− α)∣∣ ≤ εcoverage + 1/ncal, where we recall that tθ̂,1−α is the (1 − α) (population) quantile of tθ̂(X,Y ). Note that F ′ θ(t) = fθ(t) ≥ c0 on t ∈ [tθ̂,1−α − δ0, tθ̂,1−α + δ0] by Assumption D. Further, εcoverage + 1/ncal ≤ c0δ0. Therefore, by monotonicity of Fθ, we must have t̂ ∈ [tθ̂,1−α − δ0, tθ̂,1−α + δ0], and thus∣∣∣t̂− tθ̂,1−α∣∣∣ ≤ (εcoverage + 1/ncal)/c0. (14)
We next analyze t̂recal. As the dataset Drecal is independent of θ̂, we can apply the DKW Inequality (Massart, 1990, Corollary 1) to obtain that
sup t∈T
∣∣∣F̂ recal θ̂ (t)− Fθ̂(t) ∣∣∣ ≤
√ log(1/δ)
2nrecal
with probability at least 1− δ. Using a similar argument as above, we get (for δ ≤ 0.5)∣∣∣Fθ̂(t̂recal)− Fθ̂(tθ̂,1−α)∣∣∣ ≤ √ log(1/δ)
2nrecal +
1
nrecal ≤ 2
√ log(1/δ)
nrecal .
As 2 √
log(1/δ)/nrecal ≤ c0δ0, we can apply the similar argument as above to deduce that∣∣∣t̂recal − tθ̂,1−α∣∣∣ ≤ 2√log(1/δ)/nrecal/c0. (15) Combining (14) and (15) and using the Lipschitzness of the efficiency loss (Assumption A2), we get
Leff(Cθ̂,t̂recal)− Leff(Cθ̂,t̂) ≤ LT · ∣∣t̂recal − t̂∣∣ ≤ LT · (∣∣∣t̂recal − tθ̂,1−α∣∣∣+ ∣∣∣tθ̂,1−α − t̂∣∣∣)
≤ CLT · εcoverage + n−1cal + √ log(1/δ)
nrecal /c0. Finally, as we assumed εcoverage ≤ ε0, the condition of Proposition 2(b) holds, so we have
Leff(Cθ̂,t̂) ≤ inf(θ,t)∈Θ×T Lcoverage(Cθ,t)≤α Leff(Cθ,t) + 2εeff .
Summing the preceding two bounds, we get
Leff(Cθ̂,t̂recal) ≤ inf(θ,t)∈Θ×T Lcoverage(Cθ,t)≤α Leff(Cθ,t) + 2εeff + CLT ·
εcoverage + n−1cal + √ log(1/δ)
nrecal /c0. which is the desired result.

heading: E ADDITIONAL EXPERIMENTAL DETAILS
text: E.1 CONFORMAL QUANTILE FINETUNING
Datasets Our choice of the datasets follows (Feldman et al., 2021). We provide information about these datasets in Table 3.
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 onDcal 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 F̂ = [f̂lo, f̂hi] is the summed pinball loss of level α/2 for f̂lo and 1 − α/2 for f̂hi, following (Romano et al., 2019):
`(F̂ ; (xi, yi)) = ` α/2 pinball(f̂lo(xi)− yi) + ` 1−α/2 pinball(f̂hi(xi)− yi),
where for any β ∈ (0, 1), `βpinball is the pinball loss at level β:
`βpinball(t) = {− βt if t < 0, (1− β)t if t ≥ 0.
Optimization details for CP-Gen-Recal For the conformal quantile finetuning procedure with our CP-Gen-Recal, we rewrite the miscoverage loss for the quantile-based prediction interval as
1 {y /∈ Cθ,t(x)} = 1 { t−max { θ>loΦ̂(x)− y, y − θ>hiΦ̂(x) } < 0 } .
(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.1. 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, trecal is computed by the reconformalization step in Algorithm 2 on Drecal.
E.2 MULTI-OUTPUT REGRESSION
Datasets We generate offline 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).
Base predictor and optimization Our network architecture is a 3-layer MLP with width 64, input dimension din = dS + dA, and output dimension dout = dS . 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 finetuning procedure with our CP-Gen-Recal, we rewrite the miscoverage loss for the box-shaped prediction set as
1 {y /∈ Cu(x)} = 1 { y /∈
dout∏ i=1 [f̂i(x)− ui, f̂i(x) + ui]
} = 1 { 1− max
1≤i≤dout |yi − f̂i(x)|/ui < 0
} .
where we recall u ∈ Rdout 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.01. 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 û obtained above (as the θ̂), and then reconformalize an additional trecal > 0 on Drecal via the proportional reshaping of (9).
E.3 DETAILS FOR FIGURE 1
Figure 1 is obtained on one run of our conformal quantile finetuning experiment on the MEPS 19 dataset, and illustrates the coverage-efficiency tradeoff. Both figures there compute the coverage and length on the (unseen) test set Dtest, for better illustration. Figure 1 Left plots the family [f̂lo(x)− t, f̂hi(x) + t] used by Conformalized Quantile Regression. Figure 1 Right plots the family
Cθ,t(x) = [θ > loΦ̂(x)− t, θ>hiΦ̂(x)].
The specific function class of θ shown in the thinner lines is a finite set of linear interpolations of the original θ̂0 obtained in QR and the new θ̂ obtained by conformal quantile finetuning, with combination weights within {−0.3,−0.2, . . . , 1.0}. The shaded region is then obtained by filling in the area.

heading: F RESULTS FOR LABEL PREDICTION SETS ON IMAGENET
text: 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, ShuffleNet}, 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 |Dcal| = 10000, |Drecal| = 10000, and |Dtest| = 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 classifiers pretrained on ImageNet. This dataset contains 10000 images, which we randomly split (varying with seed) into |Dcal| = 4000, |Drecal| = 1000, and |Dtest| = 5000.
Methods for learning prediction sets Our constructions of the prediction sets are based on the Least Ambiguous Set-Valued Classifier (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 Ct(x) via
Ct(x) = {y ∈ [L] : p(y|x) > t},
where t is found by conformal prediction.
We consider learning a valid prediction set with smaller set size by finding an optimized ensemble weight of the K base predictors using our CP-Gen-Recal algorithm. This means we learn prediction sets of the form
Cθ,t(x) = { y ∈ [L] : pθ(y|x) :=
K∑ k=1 θkpk(y|x) > t
} ,
where {pk}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 efficiency loss and hinge approximate coverage constraint to optimize (θ, t). However, here the efficiency loss we care about (the cardinality) is non-differentiable. We make a further approximation by considering the Lqq norm with q = 0.5 as the surrogate efficiency loss:
`eff(θ, t; (xi, yi)) := L∑ y′=1 [pθ(y ′|xi)− t] q +,
with the intuition that the q → 0 limit is exactly the cardinality of Cθ,t(xi). Our final optimization problem is then
min θ∈∆K ,t>0 max λ>0
1
ncal ncal∑ i=1 L∑ y′=1 [pθ(y ′|xi)− t] q + + λ 1 ncal ncal∑ i=1 `hinge(pθ(yi|xi)− 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 Dcal with batch-size 256 for ImageNet-Val and 64 for ImageNet-V2.
After we obtain the iterates { θ̂j } (where j denotes the epoch count), we perform a further iterate
selection of first re-computing the t̂(θ̂j) by conformalizing on Dcal, and then choosing the iterate j with the best average set size also on Dcal, before feeding it into the reconformalization step with Drecal. As the Drecal 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 = 1K1K). For these two baselines, for fairness of comparison, we allow them to use the whole Dcal ∪Drecal 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 efficiency loss (here set size) is non-differentiable and needs to be further approximated to allow gradient-based algorithms.

heading: G ABLATION STUDIES
text: G.1 CONFORMAL QUANTILE FINETUNING
We report ablation results for the conformal quantile finetuning 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.2 MULTI-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%.
G.3 COMPARISON 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 (θ̂, t̂) on Dcal (we additionally reconformalize t̂ on the same Dcal to address the potential bias in t̂ brought by the approximate optimization (8)), whereas our main CP-Gen-Recal algorithm optimizes θ̂ on Dcal and reconformalizes t̂recal on Drecal.
Table 10 reports the results. Observe that, except for the volume on one dataset (Humanoid), there is no significant 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 ncal = 20000 is large—the coverage (generalization error) of CP-Gen is also nearly valid, which may be better than what our Proposition 2 suggests.

heading: H ADDITIONAL EXPERIMENTS AND ANALYSES
text: H.1 CONDITIONAL 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 finite data, we consider two proxy metrics proposed in (Feldman et al., 2021) that measure the independence between length and indicator of coverage:
• The correlation coefficient (Corr) between the following two random variables: the interval size L = length(C(X)) and the indicator of coverage V = 1 { Y ∈ Ĉ(X) } . A (population) corre-
lation of 0 is a necessary (but not sufficient) 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 nonlin-
ear correlation structures. A (population) HSIC of 0 is a necessary and sufficient condition of the independence between L and V . We estimate HSIC on the finite 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.
H.2 ALTERNATIVE 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-Dtrain ∪ Dcal: Use dataset Dtrain ∪ Dcal for training the base quantile regressor, then conformalize on Drecal. • CQR-PinballFt: Train the base quantile regressor on Dtrain, and finetune the last linear
layer on Dcal using the pinball loss (same as training), and conformalize on Drecal.
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-Dtrain ∪ Dcal) 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-Dtrain ∪ Dcal but uses a less expressive model in the finetuning stage). Overall, on almost all datasets (except for kin8nm), our CP-Gen-Recal still achieves the best length.
H.2.1 OPTIMIZATION DETAILS
CQR-Dtrain ∪ Dcal: Our original CQR baseline used Dcal for monitoring validation loss and automatically determining the early stopping (cf. Section E.1), and Drecal for conformalization. To optimize on Dtrain ∪Dcal, 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 (Drecal) and testing (Dtest), 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, and 2500 epochs for bio. For all datasets, the learning rate decays by 10x twice, at 90% and 95% of the total epochs.
CQR-PinballFt: We finetuned on Dcal 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.3 ADDITIONAL 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
Ct(x) = dout∏ i=1 [f̂i(x)− t, fi(x) + t], (16)
and use conformal prediction on Dcal ∪ Drecal to compute a conformalized t̂ and the final prediction set Ct̂(x). In other words, we perform standard conformal prediction with score function ‖y − f̂(x)‖∞.
We remark that both the Coord-wise-Recal and the Max-score-Conformal baseline methods are special instances of CP-Gen-Recalwith some fixed θ0: In our parametrization, u = (θ, t), θ determines the shape (i.e. relative ratios between the u′is) 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-Recalcan 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 significantly 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 significantly better than both baselines.
H.4 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).

Review summaries:
This paper proposes conformal methods for learning a predictive confidence set that is guaranteed to be nearly the most statistically efficient among predictive confidence sets produced by a parametric class of learning algorithms. To explain, many wrapper methods for assessing prediction uncertainty assume that the user has already committed to one particular learning algorithm and aim to quantify uncertainty in predictions produced by using *this* learning algorithm. By contrast, the methods proposed here assume that the learning algorithm has been specified only up to some parameter $\theta$ so that it makes sense to optimize $\theta$ with respect to the precision of the resulting predictive confidence sets.  From the point of view of practical implementation, the obvious challenge is the optimization of $\theta$ (as well as the associated prediction sets) over an arbitrary and large space of possible values of $\theta$, as the optimization involves a hard constraint. Thus, one main contribution of the paper is to propose a differentiable proxy that allows an efficient search over the parameter space. Regarding theoretical guarantees, the paper proves that the original, non-differentiable formulation of one of the methods achieves approximately valid coverage and nearly optimal statistical efficiency when the class of learning algorithms has a low capacity.

This paper considers to improve efficiency of conformal prediction (measured in the "size" of prediction sets). To this end, this paper uses multiple-learnable parameters by generalizing single-parameter conformal prediction. By doing so, the paper demonstrate that the proposed approach can improve the efficiency of conformal predictors while satisfying valid coverage. The efficacy of the proposed approach is demonstrated over regression, multi-output regression, and classification.

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. 

This paper introduces an extension of conformal prediction with a different formulation. Instead of guaranteed coverage for a finite calibration set, the authors solve a constrained optimization problem where the length of prediction intervals is minimized subject to a coverage constraint. As a result, coverage is only asymptotically guaranteed, but instead the length of the prediction intervals can be shortened compared to traditional split conformal prediction.   The authors present theoretical results in the form of coverage bounds for specific function classes.   In the experiments classical regression datasets are analyzed, as well as multi-output regression problems and one multi-class classification problem.