Conformal inference is (almost) free for neural networks trained with
  early stopping
Early stopping based on hold-out data is a popular regularization technique
designed to mitigate overfitting and increase the predictive accuracy of neural
networks. Models trained with early stopping often provide relatively accurate
predictions, but they generally still lack precise statistical guarantees
unless they are further calibrated using independent hold-out data. This paper
addresses the above limitation with conformalized early stopping: a novel
method that combines early stopping with conformal calibration while
efficiently recycling the same hold-out data. This leads to models that are
both accurate and able to provide exact predictive inferences without multiple
data splits nor overly conservative adjustments. Practical implementations are
developed for different learning tasks -- outlier detection, multi-class
classification, regression -- and their competitive performance is demonstrated
on real data.
Ziyi Liang, Yanfei Zhou, Matteo Sesia
# Conformal inference is (almost) free for neural networks

###### Abstract

Early stopping based on hold-out data is a popular regularization technique designed to mitigate overfitting and increase the predictive accuracy of neural networks. Models trained with early stopping often provide relatively accurate predictions, but they generally still lack precise statistical guarantees unless they are further calibrated using independent hold-out data. This paper addresses the above limitation with _conformalized early stopping_: a novel method that combines early stopping with conformal calibration while efficiently recycling the same hold-out data. This leads to models that are both accurate and able to provide exact predictive inferences without multiple data splits nor overly conservative adjustments. Practical implementations are developed for different learning tasks--outlier detection, multi-class classification, regression--and their competitive performance is demonstrated on real data.

## 1 Introduction

Deep neural networks can detect complex data patterns and leverage them to make accurate predictions in many applications, including computer vision, natural language processing, and speech recognition, to name a few examples. These models can sometimes even outperform skilled humans [1], but they still make mistakes. Unfortunately, the severity of these mistakes is compounded by the fact that neural networks are often prone to overfitting and may become overconfident [2]. Several training strategies have been developed to mitigate overconfidence, including dropout [3], batch normalization [4], weight normalization [5], data augmentation [6], and early stopping [7]; the latter is the focus of this paper.

Early stopping consists of continuously evaluating after each batch of stochastic gradient updates (or _epoch_) the predictive performance of the current model on _hold-out_ independent data. After a large number of gradient updates, only the intermediate model achieving the best performance on the hold-out data is utilized to make predictions. This strategy is often effective at mitigating overfitting and can produce relatively accurate predictions compared to fully trained models, but it does not fully resolve overconfidence because it does not lead to models with finite-sample guarantees.
A general framework for quantifying the predictive uncertainty of any _black-box_ machine learning model is that of conformal inference [8]. The key idea of conformal inference is to apply a pre-trained model to a _calibration_ set of hold-out observations drawn at random from the target population. If the calibration data are exchangeable with the test point of interest, the model performance on the calibration set can be translated into statistically rigorous predictive inferences. This framework is flexible and can accommodate different learning tasks, including out-of-distribution testing [9], classification [10], and regression [8]. For example, in the context of classification, conformal inference can give prediction sets that contain the correct label for a new data point with high probability. In theory, the quality of the trained model has no consequence on the _average_ validity of conformal inferences, but it does affect their reliability and usefulness on a case-by-case level. In particular, conformal uncertainty estimates obtained after calibrating an overconfident model may be too conservative for some test cases and too optimistic for others [11]. The goal of this paper is to combine conformal calibration with standard early stopping training techniques as efficiently as possible, in order to produce more reliable predictive inferences with a finite amount of available data.

Achieving high accuracy with deep learning often requires large training sets [12], and conformal inference makes the overall pipeline even more data-intensive. As high-quality observations can be expensive to collect, in some situations practitioners may naturally wonder whether the advantage of having principled uncertainty estimates is worth a possible reduction in predictive accuracy due to fewer available training samples. This concern is relevant because the size of the calibration set cannot be too small if one wants stable and reliable conformal inferences [13, 14]. In fact, very large calibration sets may be necessary to obtain stronger conformal inferences that are valid not only on average but also conditionally on some important individual features; see Vovk et al. [10], Romano et al. [15], and Barber et al. [16].

This paper resolves the above dilemma by showing that conformal inferences for deep learning models trained with early stopping can be obtained almost "for free"--without spending precious data. More precisely, we present an innovative method that blends model training with early stopping and conformal calibration using the same hold-out samples, essentially obtaining rigorous predictive inferences at no additional data cost compared to standard early stopping. It is worth emphasizing this result is not trivial. In fact, naively applying existing conformal calibration methods using the same hold-out samples utilized for early stopping would not lead to theoretically valid inferences, at least not without resorting to very conservative corrections.

The paper is organized as follows. Section 2 develops our _conformalized early stopping_ (CES) method, starting from outlier detection and classification, then addressing regression. Section 3 demonstrates the advantages of CES through numerical experiments. Section 4 concludes with a discussion and some ideas for further research. Additional details and results, including a theoretical analysis of the naive benchmark mentioned above, can be found in the Appendices, along with all mathematical proofs.

### Related work

Conformal inference [8, 17, 18] has become a very rich and active area of research [19, 20, 21, 22]. Many prior works studied the computation of efficient conformal inferences starting from pre-trained _black-box_ models, including for example in the context of outlier detection [9, 23, 24, 25], classification [10, 11, 26, 27, 28], and regression [8, 20, 29]. Other works have studied the general robustness of conformal
inferences to distribution shifts [30, 31] and, more broadly, to failures of the data exchangeability assumption [32, 33]. Our research is orthogonal, as we look inside the black-box model and develop a novel early-stopping training technique that is naturally integrated with conformal calibration. Nonetheless, the proposed method could be combined with those described in the aforementioned papers. Other recent research has explored different ways of bringing conformal inference into the learning algorithms [34, 35, 36, 37], and some of those works apply standard early stopping techniques, but they do not address our problem.

This paper is related to Yang and Kuchibhotla [38], which proposed a general theoretical adjustment for conformal inferences computed after model selection. That method could be utilized to account for early stopping without further data splits, as detailed in Appendix A1. However, we will demonstrate that even an improved version of such analysis remains overly conservative in the context of model selection via early stopping, and the alternative method developed in this paper performs much better in practice. Our solution is inspired by Liang, Sesia, and Sun [25], which deals with the problem of selecting the best model from an arbitrary machine learning toolbox to obtain the most powerful conformal p-values for outlier testing. The idea of Liang, Sesia, and Sun [25] extends naturally to the early stopping problem in the special cases of outlier detection and classification, but the regression setting requires substantial technical innovations. The work of Liang, Sesia, and Sun [25] is also related to Marandon et al. [39], although the latter is more distant from this paper because it focuses on theoretically controlling the false discovery rate [40] in multiple testing problems. Finally, this paper draws inspiration from Kim, Xu, and Barber [41], which shows that machine learning models trained with bootstrap (or bagging) techniques can also lead to valid conformal inferences essentially for free.

## 2 Methods

### Standard conformal inference and early stopping

Consider \(n\) data points, \(Z_{i}\) for \(i\in\mathcal{D}=[n]=\{1,\ldots,n\}\), sampled exchangeably (e.g., i.i.d.) from an unknown distribution \(P_{Z}\) with support on some space \(\mathcal{Z}\). Consider also an additional test sample, \(Z_{n+1}\). In the context of outlier detection, one wishes to test whether \(Z_{n+1}\) was sampled exchangeably from \(P_{Z}\). In classification or regression, one can write \(Z_{i}=(X_{i},Y_{i})\), where \(X_{i}\) is a feature vector while \(Y_{i}\) is a discrete category or a continuous response, and the goal is to predict the unobserved value of \(Y_{n+1}\) given \(X_{n+1}\) and the data in \(\mathcal{D}\).

The standard pipeline begins by randomly splitting the data in \(\mathcal{D}\) into three disjoint subsets: \(\mathcal{D}_{\mathrm{train}},\mathcal{D}_{\mathrm{es}},\mathcal{D}_{\mathrm{ cal}}\subset[n]\). The samples in \(\mathcal{D}_{\mathrm{train}}\) are utilized to train a model \(M\) via stochastic gradient descent, in such a way as to (approximately) minimize the desired loss \(\mathcal{L}\), while the observations in \(\mathcal{D}_{\mathrm{es}}\) and \(\mathcal{D}_{\mathrm{cal}}\) are held out. We denote by \(M_{t}\) the model learnt after \(t\) epochs of stochastic gradient descent, for any \(t\in[t_{\mathrm{max}}]\), where \(t_{\mathrm{max}}\) is a pre-determined maximum number of epochs. For simplicity, \(\mathcal{L}\) is assumed to be an additive loss, in the sense that its value calculated on the training data after \(t\) epochs is \(\mathcal{L}_{\mathrm{train}}(M_{t})=\sum_{i\in\mathcal{D}_{\mathrm{train}}} \ell(M_{t};Z_{i})\), for some appropriate function \(\ell\). For example, a typical choice for regression would be the squared-error loss: \(\ell(M_{t};Z_{i})=\left[Y_{i}-\hat{\mu}(X_{i};M_{t})\right]^{2}\), where \(\hat{\mu}(X_{i};M_{t})\) indicates the value of the regression function at \(X_{i}\), as estimated by \(M_{t}\). Similarly, the loss evaluated on \(\mathcal{D}_{\mathrm{es}}\) is denoted as \(\mathcal{L}_{\mathrm{es}}(M_{t})=\sum_{i\in\mathcal{D}_{\mathrm{es}}}\ell(M_ {t};Z_{i})\). After training for \(t_{\mathrm{max}}\) epochs, early stopping selects the model \(\hat{M}_{\mathrm{es}}\) that minimizes the loss on \(\mathcal{D}_{\mathrm{es}}\): \(\hat{M}_{\mathrm{es}}=\arg\min_{M_{t}\,:\,0\leq t\leq t_{\mathrm{max}}} \mathcal{L}_{\mathrm{es}}(M_{t})\). Conformal calibration of \(\hat{M}_{\mathrm{es}}\) is then conducted using the independent hold-out data set \(\mathcal{D}_{\mathrm{cal}}\), as
sketched in Figure 1 (a). This pipeline requires a three-way data split because: (i) \(\mathcal{D}_{\text{train}}\) and \(\mathcal{D}_{\text{es}}\) must be disjoint to ensure the early stopping criterion is effective at mitigating overfitting; and (ii) \(\mathcal{D}_{\text{cal}}\) must be disjoint from \(\mathcal{D}_{\text{train}}\cup\mathcal{D}_{\text{es}}\) to ensure the performance of the selected model \(\hat{M}_{\text{es}}\) on the calibration data gives us an unbiased preview of its future performance at test time.

### Preview of our contribution

This paper develops a novel method to jointly carry out both early stopping and conformal inference using a single hold-out data set, denoted in the following as \(\mathcal{D}_{\text{es-cal}}\). The advantage of this approach is that it allows more samples to be allocated to \(\mathcal{D}_{\text{train}}\). This is not a straightforward problem. For example, one cannot naively apply standard conformal inference methods using the same hold-out set \(\mathcal{D}_{\text{es-cal}}\) previously used for early stopping, as detailed in Appendix A1. In that case, the early stopping decision would invalidate the conformal inferences by breaking the exchangeability between the calibration data and the test point, as the latter is not used to select the predictive model. As explained in Appendix A1.2, it is possible to correct conformal inferences obtained with this naive approach by adjusting the nominal coverage level conservatively, leveraging suitable concentration inequalities [38]. However, such theoretical corrections tend to be overly pessimistic in practice and may often be too conservative to be useful; this is demonstrated by the numerical experiments described in Section 3 and previewed here in Figure 2.

By contrast, the CES method proposed in this paper is based on the following idea inspired by Liang, Sesia, and Sun [25]. Valid conformal inferences can be obtained by calibrating \(\hat{M}_{\text{es}}\) using the same data set \(\mathcal{D}_{\text{es-cal}}\) used for model selection, as long as the test sample \(Z_{n+1}\) is also involved in the early stopping rule exchangeably with all other samples in \(\mathcal{D}_{\text{es-cal}}\). This concept, illustrated schematically in Figure 1 (b), is not obvious to translate into a practical method, however, for two reasons. First, the ground truth for the test point (i.e., its outlier status or its outcome label) is unknown. Second, the method may need to be repeatedly applied for a large number of distinct test points in a computationally efficient way, and one cannot re-train the model separately for each test point. In the next section, we will explain how to overcome these challenges in the special case of early stopping for outlier detection; then, the solution will be extended to the classification and regression settings.

Figure 1: Conformal inference for models trained with early stopping. (a) Conventional pipeline requiring a three-way sample split. (b) Conformalized early stopping, requiring a two-way split.
### Conformalized early stopping for outlier detection

Consider testing whether \(Z_{n+1}\), is an _inlier_, in the sense that it was sampled from \(P_{Z}\) exchangeably with the data in \(\mathcal{D}\). Following the notation of Section 2.1, consider a partition of \(\mathcal{D}\) into two subsets, \(\mathcal{D}_{\text{train}}\) and \(\mathcal{D}_{\text{es-cal}}\), chosen at random independently of everything else, such that \(\mathcal{D}=\mathcal{D}_{\text{train}}\cup\mathcal{D}_{\text{es-cal}}\). The first step of CES consists of training a deep one-class classifier \(M\) using the data in \(\mathcal{D}_{\text{train}}\) via stochastic gradient descent for \(t^{\text{max}}\) epochs, storing all parameters characterizing the intermediate model after each \(\tau\) epochs. We refer to \(\tau\in[t^{\text{max}}]\) as the _storage period_, a parameter pre-defined by the user. Intuitively, a smaller \(\tau\) increases the memory cost of CES but may also lead to the selection of a more accurate model. While the memory cost of this approach is higher compared to that of standard early-stopping training techniques, which only require storing one model at a time, it is not prohibitively expensive. In fact, the candidate models do not need to be kept in precious RAM memory but can be stored on a relatively cheap hard drive. As reasonable choices of \(\tau\) may typically be in the order of \(T=\lfloor t^{\text{max}}/\tau\rfloor\approx 100\), the cost of CES is not excessive in many real-world situations. For example, it takes approximately 100 MB to store a pre-trained standard ResNet50 computer vision model, implying that CES would require approximately 10 GB of storage in such applications--today this costs less than $0.25/month in the cloud.

After pre-training and storing \(T\) candidate models, namely \(M_{t_{1}},\ldots,M_{t_{T}}\) for some sub-sequence \((t_{1},\ldots,t_{T})\) of \([t^{\text{max}}]\), the next step is to select the appropriate early-stopped model based on the hold-out data in \(\mathcal{D}_{\text{es-cal}}\) as well as the test point \(Z_{n+1}\). Following the notation of Section 2.1, define the value of the one-class classification loss \(\mathcal{L}\) for model \(M_{t}\), for any \(t\in[T]\), evaluated on \(\mathcal{D}_{\text{es-cal}}\) as: \(\mathcal{L}_{\text{es-cal}}(M_{t})=\sum_{i\in\mathcal{D}_{\text{es-cal}}} \ell(M_{t};Z_{i})\). Further, for any \(z\in\mathcal{Z}\), define also \(\mathcal{L}_{\text{es-cal}}^{+1}(M_{t},z)\) as:

\[\mathcal{L}_{\text{es-cal}}^{+1}(M_{t},z)=\mathcal{L}_{\text{es-cal}}(M_{t}) +\ell(M_{t};z). \tag{1}\]

Therefore, \(\mathcal{L}_{\text{es-cal}}^{+1}(M_{t},Z_{n+1})\) can be interpreted as the cumulative value of the loss function calculated on an augmented hold-out data set including also \(Z_{n+1}\). Then, we select the model \(\hat{M}_{\text{ces}}(Z_{n+1})\)

Figure 2: Average performance, as a function of the sample size, of conformal inferences based on neural networks trained and calibrated with different methods, on the _bio_ regression data [42]. Ideally, the coverage of the conformal prediction intervals should be close to 90% and their width should be small. All methods shown here guarantee 90% marginal coverage.
minimizing \(\mathcal{L}_{\text{ess-cal}}^{+1}(M_{t},Z_{n+1})\):

\[\hat{M}_{\text{ces}}(Z_{n+1})=\operatorname*{arg\,min}_{M_{t_{j}}\;:1\leq j\leq T }\mathcal{L}_{\text{ess-cal}}^{+1}(M_{t_{j}},Z_{n+1}). \tag{2}\]

Note that the computational cost of evaluating (2) is negligible compared to that of training the models.

Next, the selected model \(\hat{M}_{\text{ces}}(Z_{n+1})\) is utilized to compute a conformal p-value [24] to test whether \(Z_{n+1}\) is an inlier. In particular, \(\hat{M}_{\text{ces}}(Z_{n+1})\) is utilized to compute _nonconformity scores_\(\hat{S}_{i}(Z_{n+1})\) for all samples \(i\in\mathcal{D}_{\text{ess-cal}}\cup\{n+1\}\). These scores rank the observations in \(\mathcal{D}_{\text{ess-cal}}\cup\{n+1\}\) based on how the one-class classifier \(\hat{M}_{\text{ces}}(Z_{n+1})\) perceives them to be similar to the training data; by convention, a smaller value of \(\hat{S}_{i}(Z_{n+1})\) suggests \(Z_{i}\) is more likely to be an outlier. Suitable scores are typically included in the output of standard one-class classification models, such as those provided by the Python library PyTorch. For simplicity, we assume all scores are almost-surely distinct; otherwise, ties can be broken at random by adding a small amount of independent noise. Then, the conformal p-value \(\hat{u}_{0}(Z_{n+1})\) is given by the usual formula:

\[\hat{u}_{0}(Z_{n+1})=\frac{1+|i\in\mathcal{D}_{\text{ess-cal}}:\hat{S}_{i}\leq \hat{S}_{n+1}|}{1+|\mathcal{D}_{\text{ess-cal}}|}, \tag{3}\]

making the dependence of \(\hat{S}_{i}\) on \(Z_{n+1}\) implicit in the interest of space. This method, outlined by Algorithm A4 in Appendix A2, gives p-values that are exactly valid in finite samples, in the sense that they are stochastically dominated by the uniform distribution under the null hypothesis. The only technical requirement is that the learning algorithm should be invariant to permutations of the training data, which is a mild and realistic assumption.

**Theorem 1**.: _Assume \(Z_{1},\ldots,Z_{n},Z_{n+1}\) are exchangeable random samples, and let \(\hat{u}_{0}(Z_{n+1})\) be the output of Algorithm A4, as given in (3). Suppose the underlying machine learning algorithm is invariant to permutations of the training data. Then, \(\mathbb{P}\left[\hat{u}_{0}(Z_{n+1})\leq\alpha\right]\leq\alpha\) for any \(\alpha\in(0,1)\)._

### Conformalized early stopping for classification

The above CES method will now be extended to deal with \(K\)-class classification problems, for any \(K\geq 2\). Consider \(n\) exchangeable pairs of observations \((X_{i},Y_{i})\), for \(i\in\mathcal{D}=[n]\), and a test point \((X_{n+1},Y_{n+1})\) whose label \(Y_{n+1}\in[K]\) has not yet been observed. The goal is to construct an informative prediction set for \(Y_{n+1}\) given the observed features \(X_{n+1}\) and the rest of the data, assuming \((X_{n+1},Y_{n+1})\) is exchangeable with the observations indexed by \(\mathcal{D}\). An ideal goal would be to construct the smallest possible prediction set with guaranteed _feature-conditional coverage_ at level \(1-\alpha\), for any fixed \(\alpha\in(0,1)\). Formally, a prediction set \(\hat{C}_{\alpha}(X_{n+1})\subseteq[K]\) has feature-conditional coverage at level \(1-\alpha\) if \(\mathbb{P}[Y_{n+1}\in\hat{C}_{\alpha}(X_{n+1})\mid X_{n+1}=x]\geq 1-\alpha\), for any \(x\in\mathcal{X}\), where \(\mathcal{X}\) is the feature space. Unfortunately, perfect feature-conditional coverage is extremely difficult to achieve unless the feature space \(\mathcal{X}\) is very small. Therefore, in practice, one must be satisfied with obtaining relatively weaker guarantees, such as _label-conditional coverage_ and _marginal coverage_. Formally, \(\hat{C}_{\alpha}(X_{n+1})\) has \(1-\alpha\) label-conditional coverage if \(\mathbb{P}[Y_{n+1}\in\hat{C}_{\alpha}(X_{n+1})\mid Y_{n+1}=y]\geq 1-\alpha\), for any \(y\in[K]\), while marginal coverage corresponds to \(\mathbb{P}[Y_{n+1}\in\hat{C}_{\alpha}(X_{n+1})]\geq 1-\alpha\). Label-conditional coverage is stronger than marginal coverage, but both criteria are useful because the latter is easier to achieve with smaller (and hence more informative) prediction sets.
We begin by focusing on label-conditional coverage, as this follows most easily from the results of Section 2.3. This solution will be extended in Appendix A3 to target marginal coverage. The first step of CES consists of randomly splitting \(\mathcal{D}\) into two subsets, \(\mathcal{D}_{\text{train}}\) and \(\mathcal{D}_{\text{es-cal}}\), as in Section 2.3. The samples in \(\mathcal{D}_{\text{es-cal}}\) are further divided into subsets \(\mathcal{D}^{y}_{\text{es-cal}}\) with homogeneous labels; that is, \(\mathcal{D}^{y}_{\text{es-cal}}=\{i\in\mathcal{D}_{\text{es-cal}}:Y_{i}=y\}\) for each \(y\in[K]\). The data in \(\mathcal{D}_{\text{train}}\) are utilized to train a neural network classifier via stochastic gradient descent, storing the intermediate candidate models \(M_{t}\) after each \(\tau\) epochs. This is essentially the same approach as in Section 2.3, with the only difference being that the neural network is now designed to perform \(K\)-class classification rather than one-class classification. Therefore, this neural network should have a soft-max layer with \(K\) nodes near its output, whose values corresponding to an input data point with features \(x\) are denoted as \(\hat{\pi}_{y}(x)\), for all \(y\in[K]\). Intuitively, we will interpret \(\hat{\pi}_{y}(x)\) as approximating (possibly inaccurately) the true conditional data-generating distribution; i.e., \(\hat{\pi}_{y}(x)\approx\mathbb{P}\left[Y=y\mid X=x\right]\).

For any model \(M_{t}\), any \(x\in\mathcal{X}\), and any \(y\in[K]\), define the augmented loss \(\mathcal{L}^{+1}_{\text{es-cal}}(M_{t},x,y)\) as:

\[\mathcal{L}^{+1}_{\text{es-cal}}(M_{t},x,y)=\mathcal{L}_{\text{es-cal}}(M_{t} )+\ell(M_{t};x,y). \tag{4}\]

Concretely, a typical choice for \(\ell\) is the cross-entropy loss: \(\ell(M_{t};x,y)=-\log\hat{\pi}_{y}^{t}(x)\), where \(\hat{\pi}^{t}\) denotes the soft-max probability distribution estimated by model \(M_{t}\). Intuitively, \(\mathcal{L}^{+1}_{\text{es-cal}}(M_{t},x,y)\) is the cumulative value of the loss function calculated on an augmented hold-out data set including also the imaginary test sample \((x,y)\). Then, for any \(y\in[K]\), CES selects the model \(\hat{M}_{\text{ces}}(X_{n+1},y)\) minimizing \(\mathcal{L}^{+1}_{\text{es-cal}}(M_{t},X_{n+1},y)\) among the \(T\) stored models:

\[\hat{M}_{\text{ces}}(X_{n+1},y)=\operatorname*{arg\,min}_{M_{t_{j}}:1\leq j \leq T}\mathcal{L}^{+1}_{\text{es-cal}}(M_{t_{j}},X_{n+1},y). \tag{5}\]

The selected model \(\hat{M}_{\text{ces}}(X_{n+1},y)\) is then utilized to compute a conformal p-value for testing whether \(Y_{n+1}=y\). In particular, we compute nonconformity scores \(\hat{S}^{y}_{i}(X_{n+1})\) for all \(i\in\mathcal{D}^{y}_{\text{es-cal}}\cup\{n+1\}\), imagining that \(Y_{n+1}=y\). Different types of nonconformity scores can be easily accommodated, but in this paper, we follow the _adaptive_ strategy of Romano, Sesia, and Candes [11]. The computation of these nonconformity scores based on the selected model \(\hat{M}_{\text{ces}}\) is reviewed in Appendix A4. Here, we simply note the p-value is given by:

\[\hat{u}_{y}(X_{n+1})=\frac{1+|i\in\mathcal{D}^{y}_{\text{es-cal}}:\hat{S}^{y} _{i}\leq\hat{S}^{y}_{n+1}|}{1+|\mathcal{D}^{y}_{\text{es-cal}}|}, \tag{6}\]

again making the dependence of \(\hat{S}^{y}_{i}\) on \(X_{n+1}\) implicit. Finally, the prediction set \(\hat{C}_{\alpha}(X_{n+1})\) is constructed by including all possible labels for which the corresponding null hypothesis cannot be rejected at level \(\alpha\):

\[\hat{C}_{\alpha}(X_{n+1})=\{y\in[K]:\hat{u}_{y}(X_{n+1})\geq\alpha\}\,. \tag{7}\]

This method, outlined by Algorithm A5 in Appendix A2, guarantees label-conditional coverage at level \(1-\alpha\).

**Theorem 2**.: _Assume \((X_{1},Y_{1}),\ldots,(X_{n+1},Y_{n+1})\) are exchangeable, and let \(\hat{C}_{\alpha}(X_{n+1})\) be the output of Algorithm A5, as given in (7), for any given \(\alpha\in(0,1)\). Suppose the underlying machine learning algorithm is invariant to permutations of the training data. Then, \(\mathbb{P}[Y_{n+1}\in\hat{C}_{\alpha}(X_{n+1})\mid Y_{n+1}=y]\geq 1-\alpha\) for any \(y\in[K]\)._
### Conformalized early stopping for regression

This section extends CES to regression problems with a continuous outcome. As in the previous sections, consider a data set containing \(n\) exchangeable observations \((X_{i},Y_{i})\), for \(i\in\mathcal{D}=[n]\), and a test point \((X_{n+1},Y_{n+1})\) with a latent label \(Y_{n+1}\in\mathbb{R}\). The goal is to construct a reasonably narrow _prediction interval_\(\hat{C}_{\alpha}(X_{n+1})\) for \(Y_{n+1}\) that is guaranteed to have marginal coverage above some level \(1-\alpha\), i.e., \(\mathbb{P}[Y_{n+1}\in\hat{C}_{\alpha}(X_{n+1})]\geq 1-\alpha\), and can also practically achieve reasonably high feature-conditional coverage. Developing a CES method for this problem is more difficult compared to the classification case studied in Section 2.4 due to the infinite number of possible values for \(Y_{n+1}\). In fact, a naive extension of Algorithm A5 would be computationally unfeasible in the regression setting, for the same reason why full-conformal prediction [8] is generally impractical. The novel solution described below is designed to leverage the particular structure of an early stopping criterion based on the squared-error loss evaluated on hold-out data. Focusing on the squared-error loss makes CES easier to implement and explain using classical _absolute residual_ nonconformity scores [8, 43]. However, similar ideas could also be repurposed to accommodate other scores, such as those based on quantile regression [29], conditional distributions [44, 45], or conditional histograms [46].

As usual, we randomly split \(\mathcal{D}\) into \(\mathcal{D}_{\text{train}}\) and \(\mathcal{D}_{\text{es-cal}}\). The data in \(\mathcal{D}_{\text{train}}\) are utilized to train a neural network via stochastic gradient descent, storing the intermediate models \(M_{t}\) after each \(\tau\) epoch. The approach is similar to those in Sections 2.3-2.4, although now the output of a model \(M_{t}\) applied to a sample with features \(x\) is denoted by \(\hat{\mu}_{t}(x)\) and is designed to approximate (possibly inaccurately) the conditional mean of the unknown data-generating distribution; i.e., \(\hat{\mu}_{t}(x)\approx\mathbb{E}\left[Y\mid X=x\right]\). (Note that we will omit the superscript \(t\) unless necessary to avoid ambiguity). For any model \(M_{t}\) and any \(x\in\mathcal{X}\), \(y\in[K]\), define

\[\mathcal{L}_{\text{es-cal}}^{+1}(M_{t},x,y)=\mathcal{L}_{\text{es-cal}}(M_{t })+[y-\hat{\mu}_{t}(X_{n+1})]^{2}. \tag{8}\]

Consider now the following optimization problem,

\[\hat{M}_{\text{ces}}(X_{n+1},y)=\operatorname*{arg\,min}_{M_{t_{j}}\,:\,1\leq j \leq T}\mathcal{L}_{\text{es-cal}}^{+1}(M_{t_{j}},X_{n+1},y), \tag{9}\]

which can be solved simultaneously for all \(y\in\mathbb{R}\) thanks to the amenable form of (8). In fact, each \(\mathcal{L}_{\text{es-cal}}^{+1}(M_{t},x,y)\) is a simple quadratic function of \(y\); see the sketch in Figure 3. This implies \(\hat{M}_{\text{ces}}(X_{n+1},y)\) is a step function, whose parameters can be computed at cost \(\mathcal{O}(T\log T)\) with an efficient divide-and-conquer algorithm designed to find the lower envelope of a family of parabolas [47, 48]; see Appendix A5.

Therefore, \(\hat{M}_{\text{ces}}(X_{n+1},y)\) has \(L\) distinct steps, for some \(L=\mathcal{O}(T\log T)\) that may depend on \(X_{n+1}\), and it can be written as a function of \(y\) as:

\[\hat{M}_{\text{ces}}(X_{n+1},y)=\sum_{l=1}^{L}m_{l}(X_{n+1})\mathbbm{1}\left[y \in(k_{l-1},k_{l}]\right], \tag{10}\]

where \(m_{l}(X_{n+1})\in[T]\) represents the best model selected within the interval \((k_{l-1},k_{l}]\) such that \(m_{l}(X_{n+1})\neq m_{l-1}(X_{n+1})\) for all \(l\in[L]\). Above, \(k_{1}\leq k_{2}\leq\cdots\leq k_{L}\) denote the _knots_ of \(\hat{M}_{\text{ces}}(X_{n+1},y)\), which also depend on \(X_{n+1}\) and are defined as the boundaries in the domain of \(y\) between each consecutive pair of steps, with the understanding that \(k_{0}=-\infty\) and \(k_{L+1}=+\infty\)
Then, for each step \(l\in[L]\), let \(\mathcal{B}_{l}\) indicate the interval \(\mathcal{B}_{l}=(k_{l-1},k_{l}]\) and, for all \(i\in\mathcal{D}_{\text{es-cal}}\), evaluate the nonconformity score \(\hat{S}_{i}(X_{n+1},\mathcal{B}_{l})\) for observation \((X_{i},Y_{i})\) based on the regression model indicated by \(m_{l}(X_{n+1})\); i.e.,

\[\hat{S}_{i}(X_{n+1},\mathcal{B}_{l})=|Y_{i}-\hat{\mu}_{m_{l}(X_{n+1})}(X_{i})|. \tag{11}\]

Let \(\hat{Q}_{1-\alpha}(X_{n+1},\mathcal{B}_{l})\) denote the \(\lceil(1-\alpha)(1+|\mathcal{D}_{\text{es-cal}}|)\rceil\)-th largest value among all nonconformity scores \(\hat{S}_{i}(X_{n+1},\mathcal{B}_{l})\), assuming for simplicity that there are no ties; otherwise, ties can be broken at random. Then, define the interval \(\hat{C}_{\alpha}(X_{n+1},\mathcal{B}_{l})\) as that obtained by applying the standard conformal prediction method with absolute residual scores based on the regression model \(\hat{\mu}_{m_{l}(X_{n+1})}(X_{n+1})\):

\[\hat{C}_{\alpha}(X_{n+1},\mathcal{B}_{l})=\hat{\mu}_{m_{l}(X_{n+1})}(X_{n+1}) \pm\hat{Q}_{1-\alpha}(X_{n+1},\mathcal{B}_{l}). \tag{12}\]

Finally, the prediction interval \(\hat{C}_{\alpha}(X_{n+1})\) is given by:

\[\hat{C}_{\alpha}(X_{n+1})=\text{Convex}\left(\cup_{l=1}^{L}\{\mathcal{B}_{l} \cap\hat{C}_{\alpha}(X_{n+1},\mathcal{B}_{l})\}\right), \tag{13}\]

where \(\text{Convex}(\cdot)\) denotes the convex hull of a set. This procedure is summarized in Algorithm A6 and it is guaranteed to produce prediction sets with valid marginal coverage.

**Theorem 3**.: _Assume \((X_{1},Y_{1}),\ldots,(X_{n+1},Y_{n+1})\) are exchangeable, and let \(\hat{C}_{\alpha}(X_{n+1})\) be the output of Algorithm A6, as given by (13), for any given \(\alpha\in(0,1)\). Suppose the underlying machine learning algorithm is invariant to permutations of the training data. Then, \(\mathbb{P}[Y_{n+1}\in\hat{C}_{\alpha}(X_{n+1})]\geq 1-\alpha\)._

The intuition behind the above method is as follows. Each intermediate interval \(\hat{C}_{\alpha}(X_{n+1},\mathcal{B}_{l})\)

Figure 3: Squared-error loss on test-augmented hold-out data for three alternative regression models \(M_{1},M_{2}\) and \(M_{3}\), as a function of the place-holder outcome \(y\) for the test point. The CES method utilizes the best model for each possible value of \(y\), which is identified by the lower envelope of these three parabolas. In this case, the lower envelope has two knots, \(k_{1}\) and \(k_{3}\).
for \(l\in[L]\), may be thought of as being computed by applying, under the null hypothesis that \(Y_{n+1}\in\mathcal{B}_{l}\), the classification method from Section 2.4 for a discretized version of our problem based on the partition \(\{\mathcal{B}_{l}\}_{l=1}^{L}\). Then, leveraging the classical duality between confidence intervals and p-values, it becomes clear that taking the intersection of \(\mathcal{B}_{l}\) and \(\hat{C}_{\alpha}(X_{n+1},\mathcal{B}_{l})\) essentially amounts to including the "label" \(\mathcal{B}_{l}\) in the output prediction if the null hypothesis \(Y_{n+1}\in\mathcal{B}_{l}\) cannot be rejected. The purpose of the final convex hull operation is to CES outputs a contiguous prediction interval, which is what we originally stated to seek.

Although it is unlikely, Algorithm A6 may sometimes produce an empty set, which is an uninformative and potentially confusing output. A simple solution consists of replacing any empty output with the naive conformal prediction interval computed by Algorithm A3 in Appendix A1, which leverages an early-stopped model selected by looking at the original calibration data set without the test point. This approach is outlined by Algorithm A9 in Appendix A6. As the intervals given by Algorithm A9 always contain those output by Algorithm A6, it follows that Algorithm A9 also enjoys guaranteed coverage; see Corollary A3.

## 3 Numerical experiments

### Outlier detection

The use of CES for outlier detection is demonstrated using the _CIFAR10_ data set [49], a collection of 60,000 32-by-32 RGB images from 10 classes including common objects and animals. A convolutional neural network with ReLU activation functions is trained on a subset of the data to minimize the cross-entropy loss. The maximum number of epochs is set to be equal to 50. The trained classification model is then utilized to compute conformity scores for outlier detection with the convention that cats are inliers and the other classes are outliers. In particular, a nonconformity score for each \(Z_{n+1}\) is defined as 1 minus the output of the soft-max layer corresponding to the label "cat". This can be interpreted as an estimated probability of \(Z_{n+1}\) being an outlier. After translating these scores into a conformal p-value \(\hat{u}_{0}(Z_{n+1})\), the null hypothesis that \(Z_{n+1}\) is a cat is rejected if \(\hat{u}_{0}(Z_{n+1})\leq\alpha=0.1\).

The total number of samples utilized for training, early stopping, and conformal calibration is varied between 500 and 2000. In each case, CES is applied using 75% of the samples for training and 25% for early stopping and calibration. Note that the calibration step only utilizes inliers, while the other data subsets also contain outliers. The empirical performance of CES is measured in terms of the probability of falsely rejecting a true null hypothesis--the false positive rate (FPR)--and the probability of correctly rejecting a false null hypothesis--the true positive rate (TPR). The CES method is compared to three benchmarks. The first benchmark is naive early stopping with the best (_hybrid_) theoretical correction for the nominal coverage level described in Appendix A1.2. The second benchmark is early stopping based on data splitting, which utilizes 50% of the available samples for training, 25% for early stopping, and 25% for calibration. The third benchmark is full training without early stopping, which simply selects the model obtained after the last epoch. The test set consists of 100 independent test images, half of which are outliers. All results are averaged over 100 trials based on independent data subsets.

Figure 4 summarizes the performance of the four methods as a function of the total sample size; see Table A1 in Appendix A7 for the corresponding standard errors. All methods control the FPR below 10%, as expected, but CES achieves the highest TPR. The increased power of CES
compared to data splitting is not surprising, as the latter relies on a less accurate model trained on less data. By contrast, the naive benchmark trains a model more similar to that of CES, but its TPR is not as high because the theoretical correction for the naive conformal p-values is overly pessimistic. Finally, full training is the least powerful competitor for large sample sizes because its underlying model becomes more and more overconfident as the training set grows.

### Multi-class classification

The same _CIFAR10_ data [49] are utilized here to demonstrate the performance of CES for a 10-class classification task. These experiments are conducted similarly to those in Section 3.1. The only difference is that now the soft-max output of the convolutional neural network is translated into conformal prediction sets, as explained in Appendix A4, instead of conformal p-values. The CES method is compared to the same three benchmarks adopted in Section 3.1. All prediction sets are calibrated to guarantee 90% marginal coverage, and their performances are evaluated based on cardinality.

Figure 5 summarizes the results averaged over 100 independent realizations of these experiments, while Table A2 in Appendix A7 reports on the corresponding standard errors. While all approaches

Figure 4: Average performance, as a function of the sample size, of conformal inferences for outlier detection based on neural networks trained and calibrated with different methods, on the _CIFAR10_ data [49]. Ideally, the TPR should be as large as possible while maintaining the FPR below 0.1.

Figure 5: Average performance, as a function of the sample size, of conformal prediction sets for multi-class classification based on neural networks trained and calibrated with different methods, on the _CIFAR10_ data [49]. Ideally, the coverage should be close to 90% and the cardinality should be small.
always achieve the nominal coverage level, the CES method is able to do so with the smallest, and hence most informative, prediction sets. As before, the more disappointing performance of the data splitting benchmark can be explained by the more limited amount of data available for training, that of the naive benchmark by the excessive conservativeness of its theoretical correction, and that of the full training benchmark by overfitting.

### Regression

We now apply the CES method to the following 3 public-domain regression data sets from the UCI Machine Learning repository [50]: physicochemical properties of protein tertiary structure (_bio_) [42], hourly and daily counts of rental bikes (_bike_) [51], and concrete compressive strength (_concrete_) [52]. These data sets were previously also considered by Romano, Patterson, and Candes [29], to which we refer for further details. As in the previous sections, we compare CES to the usual three benchmarks: naive early stopping with the _hybrid_ theoretical correction for the nominal coverage level, early stopping based on data splitting, and full model training without early stopping. All methods utilize the same neural network with two hidden layers of width 128 and ReLU activation functions, trained for up to 1000 epochs. The models are calibrated in such a way as to produce conformal prediction sets with guaranteed 90% marginal coverage for a test set of 100 independent data points. The total sample size available for training, early stopping and calibration is varied between 200 and 2000. These data are allocated for specific training, early-stopping, and calibration operations as in Sections 3.1-3.2. The performance of each method is measured in terms of marginal coverage, worst-slab conditional coverage [53]--estimated as described in Sesia and Candes [14]--and average width of the prediction intervals. All results are averaged over 100 independent experiments, each based on a different random sample from the original raw data sets.

Figure 2 summarizes the performance of the four alternative methods on the _bio_ data, as a function of the total sample size; see Table A8 in Appendix A7 for the corresponding standard errors. These results show that all methods reach 90% marginal coverage in practice, as anticipated by the mathematical guarantees, although the theoretical correction for the naive early stopping method appears to be overly conservative. The CES method clearly performs best, in the sense that it leads to the shortest prediction intervals while also achieving approximately valid conditional coverage. By contrast, the conformal prediction intervals obtained without early stopping have significantly lower conditional coverage, which is consistent with the prior intuition that fully trained neural networks can sometimes suffer from overfitting. More detailed results from these experiments can be found in Table A3 in Appendix A7. Analogous results corresponding to the _bike_ and _concrete_ data sets can be found in Figures A8-A9 and Tables A4-A5 in Appendix A7.

Finally, it must be noted that the widths of the prediction intervals output by the CES method in these experiments are very similar to those of the corresponding intervals produced by naively applying early stopping without data splitting and without the theoretical correction described in Appendix A1. This naive approach was not taken as a benchmark because it does not guarantee valid coverage, unlike the other methods. Nonetheless, it is interesting to note that the rigorous theoretical properties of the CES method do not come at the expense of a significant loss of power compared to this very aggressive heuristic, and in this sense, one may say that the conformal inferences computed by CES are "almost free".
Discussion

This paper has focused on early stopping and conformal calibration because these are two popular techniques designed to mitigate overconfidence that were previously combined inefficiently, but the relevance of our methodology is much more general. In fact, similar ideas have already been utilized in the context of outlier detection to tune hyper-parameters and select the most promising candidate from an arbitrary toolbox of machine learning models [25]. The techniques developed in this paper, however, also allow one to calibrate, without further data splits, the most promising model selected in a data-driven way from an arbitrary machine learning toolbox in the context of multi-class classification and regression.

As detailed in Appendix A1, the naive benchmark that uses the same hold-out data twice, both for standard early stopping and standard conformal calibration, is not theoretically valid without conservative corrections. Nonetheless, we have observed this naive approach often performs similarly to CES in practice. Of course, the naive benchmark may sometimes fail, and thus we would advise practitioners to apply the theoretically principled CES whenever its additional memory costs are not prohibitive. However, the empirical evidence suggests the naive benchmark may not be a completely unreasonable heuristic when CES is not applicable.

Finally, it is worth noting that CES for regression was implemented in this paper using classical nonconformity scores [8, 43] that are not designed to deal efficiently with heteroscedastic data [14]. However, the general idea can be extended to accommodate virtually any other type of nonconformity score, including existing options based on quantile regression [29], conditional distributions [44, 45], or conditional histograms [46]. The reason why this paper has focused on the classical absolute residual scores is that they are more intuitive to apply in conjunction with an early stopping criterion based on the squared-error loss. In the future, it would be interesting to extend the CES method as to accommodate early stopping criteria based on alternative loss functions, including for example the pinball loss for quantile regression.

Software implementing the algorithms and data experiments are available online at [https://github.com/ZiyiLiang/Conformalized_early_stopping](https://github.com/ZiyiLiang/Conformalized_early_stopping).

### Acknowledgements

The authors thank the Center for Advanced Research Computing at the University of Southern California for providing computing resources to carry out numerical experiments. M. S. and Y. Z. are supported by NSF grant DMS 2210637. M. S. is also supported by an Amazon Research Award.
