Distribution-free inference with hierarchical data

In an i.i.d. setting, we would work with data points that are sampled as $Z_{i}=(X_{i},Y_{i})\stackrel{{\scriptstyle\textnormal{iid}}}{{\sim}}P$ for some distribution $P$. For example, each data point $i$ might correspond to a randomly sampled individual; we would like to ask how features $X_{i}$ can be used to predict a response $Y_{i}$—for instance, if we are studying academic performance for a sample of children, $X_{i}$ might reflect features such as age, school quality rating, parents’ income level, etc, for the $i$th child in the sample, and $Y_{i}$ is the child’s test score.

In a hierarchical setting, we instead assume that the data are drawn via a hierarchical sampling procedure. For example, if our procedure for recruiting subjects in our study involves sampling classrooms within a city, and then sampling multiple children within each classroom, then we would expect some amount of dependence between subjects that are in the same classroom. Can we still use this training data set to build a model for predicting $Y$ from $X$, for a new child that comes from a new (previously unsampled) classroom? In this setting, statistical procedures that rely on an i.i.d. sample may no longer be valid. Parametric approaches might address this issue by, e.g., adding a random effects component to the parametric model. In this work, we instead seek to adapt to hierarchically sampled data within the framework of distribution-free inference.

Distribution-free inference has received much attention recently in the statistics and machine learning literature, with many methodological advances but also many results examining the limits of the distribution-free framework.

In the standard setting with exchangeable (e.g., i.i.d.) data, conformal prediction (see Vovk et al. (2005); Lei et al. (2018) for background) provides a universal framework for distribution-free prediction. Split conformal prediction (Vovk et al., 2005; Papadopoulos et al., 2002) offers a variant of conformal prediction that relies on data splitting to reduce the computational cost, but with some loss of statistical accuracy. Barber et al. (2021b) introduces the jackknife+, which provides a valid distribution-free prediction set without data splitting while having a more moderate computational cost than full conformal prediction.

Applications of the conformal prediction framework to settings beyond the framework of exchangeability have also been studied in the recent literature. For example, Tibshirani et al. (2019) and Podkopaev and Ramdas (2021) introduce extensions of conformal prediction that deal with covariate shift and label shift, respectively, while Barber et al. (2022) allow for unknown deviations from exchangeability. Xu and Xie (2021) considers applications to time series data, and Lei et al. (2015) considers functional data.

For the goal of having a useful bound for conditional coverage in the distribution-free setting, Vovk (2012) and Barber et al. (2021a) show impossibility results for coverage conditional on $X_{\textnormal{test}}$ in the case that the feature distribution is nonatomic. Izbicki et al. (2019) and Chernozhukov et al. (2021) propose methods that provide conditional coverage guarantees under strong additional assumptions. Conformalized quantile regression, introduced by Romano et al. (2019), can provide improved control of conditional miscoverage rates empirically, by accounting for heterogeneous spread of $Y$ given $X$.

Problems other than prediction have been studied as well in the distribution-free framework. Barber (2020) and Medarametla and Candès (2021) discuss limits on having a useful inference on conditional mean and median, respectively. Lee and Barber (2021) prove that for discrete feature distributions, we can make use of repeated feature observations to attain meaningful inference for the conditional mean.

Distribution-free inference for the hierarchical data setting was previously studied by (Dunn et al., 2022), proposing a range of different methods for the hierarchical setting. We describe their methods and results in detail in Section 2.3. (Dunn et al. (2022) also study the problem of prediction for a new data point observed within one of the groups in the training sample—e.g., a new child, sampled from a classroom that is already present in the study—but this is a fundamentally different inference problem and we do not consider it in this work.) The setting of repeated measurements is explored by Cheng et al. (2022), who contrast empirical risk minimization when the repeated measurements are aggregated prior to estimation and when they are kept separate. Our contributions complement their work with a framework for conformal prediction. Like Cheng et al. (2022), we find that non-aggregated repeated measurements can enable a more informative statistical analysis.

The remainder of this paper is organized as follows. In Section 2, we first give background on existing methods for the hierarchical setting (Dunn et al., 2022), and then present our proposed method, Hierarchical Conformal Prediction (HCP), and compare HCP with the existing methods on simulated data. In Section 3, we turn to the special case of data with repeated measurements; we examine the problem of conditionally valid coverage in this setting, and propose a variant of our method, HCP², that offers a stronger coverage guarantee, which we verify with simulations. Section 5 offers a short discussion of our findings and of open questions raised by this work. The proofs of our theorems, together with some additional methods and theoretical results, are deferred to the Appendix.

First, for background, consider the setting of i.i.d. or exchangeable data, $Z_{1},\dots,Z_{n},Z_{n+1}$, where $Z_{i}=(X_{i},Y_{i})$ for each $i$. Here $Z_{1},\dots,Z_{n}$ is the training data, while $Z_{\textnormal{test}}=Z_{n+1}$ is the test point. The split conformal prediction method (Vovk et al., 2005) constructs a distribution-free prediction interval as follows. As a preliminary step, we split the training data into two data sets of size $n_{0}+n_{1}=n$, with the first $n_{0}$ data points used for training, and the remaining $n_{1}$ data points used for calibration. Once the data is split, the first step is to use the training portion of the data (i.e., $Z_{1},\dots,Z_{n_{0}}$), we fit a score function

where $s(z)=s(x,y)$ measures the extent to which a data point $(x,y)$ conforms to the trends observed in the training data, with large values indicating a more unusual value of the data point. For example, if the response variables $Y$ lie in $\mathcal{Y}=\mathbb{R}$, one mechanism for defining this function is to first run a regression on the labeled data set $\{(X_{i},Y_{i}):i=1,\dots,n_{0}\}$, to produce a fitted model $\widehat{\mu}:\mathcal{X}\rightarrow\mathbb{R}$, and then define $s(x,y)=|y-\widehat{\mu}(x)|$, the absolute value of the residual. Next, the second step is to use the calibration set (i.e., $Z_{n_{0}+1},\dots,Z_{n}$) to determine a cutoff threshold for the score, and construct the corresponding prediction interval:

where $\delta_{t}$ denotes the point mass at $t$, while $Q_{1-\alpha}(\cdot)$ denotes the $(1-\alpha)$-quantile of a distribution.²²2Formally, for a distribution $P$ on $\mathbb{R}$, the quantile is defined as $Q_{1-\alpha}(P)=\inf\{t\in\mathbb{R}:\mathbb{P}_{{T\sim P}}\left\{{T\leq t}\right\}\geq 1-\alpha\}$.

Vovk et al. (2005) establish marginal coverage for the split conformal method: if the data $Z_{1},\dots,Z_{n},Z_{\textnormal{test}}$ satisfies exchangeability (that is, if the property (3) holds when we take $Z_{\textnormal{test}}=Z_{n+1}$), then

We note that this guarantee is marginal—in particular, it does not hold if we were to condition on the test features $X_{\textnormal{test}}$ (and also, holds only on average over the draw of the training and calibration data $Z_{1},\dots,Z_{n}$).

In the setting of i.i.d. or exchangeable data, the full conformal prediction method (Vovk et al., 2005) provides an alternative method that avoids the loss of accuracy incurred by sample splitting, at the cost of a greater computational burden. As a compromise between split and full conformal, cross-validation type methods are proposed by Vovk (2015); Vovk et al. (2018); Barber et al. (2021b). We give background on these alternative methods in Appendix A.

Now we return to the hierarchical data setting that is the problem of interest in this paper. Can split conformal prediction be applied to this setting? That is, given the observed data points $Z_{1,1},\dots,Z_{1,N_{1}},\dots,Z_{K,1},\dots,Z_{K,N_{K}}$, can we simply apply split conformal to this training data set to obtain a valid prediction interval? In fact, the hierarchical exchangeability property (Definition 1) is not sufficient. In general, the data may satisfy hierarchical exchangeability without satisfying the (ordinary, non-hierarchical) exchangeability property (3) that is needed for split conformal to have theoretical validity—for example, if the data consists of students sampled from $K$ many classrooms, then there may be higher correlation between students in the same classroom. (An exception is the trivial special case where $N_{1}=\dots=N_{K}=1$, i.e., all groups contain exactly one observation, in which case hierarchical exchangeability simply reduces to ordinary exchangeability.)

We now introduce our extension of split conformal prediction for the hierarchical data setting. Given training data $\tilde{Z}_{1},\dots,\tilde{Z}_{K}$, we split it into $K_{0}+K_{1}$ many groups, with $K_{0}$ groups used for training and $K_{1}=K-K_{0}$ groups for calibration. First, using data $\tilde{Z}_{1},\dots,\tilde{Z}_{K_{0}}$, we fit a score function $s:\mathcal{Z}\rightarrow\mathbb{R}$ (as before, a canonical choice is to use the absolute value of the residual for some fitted model, $s(z)=s(x,y)=|y-\widehat{\mu}(x)|$ for some fitted model $\widehat{\mu}$). Next, we use the calibration set to set a threshold $T$ for the prediction interval, as follows:

Compared to the original split conformal interval constructed in (5), we see that the constructions follow the same format, but the score threshold for HCP in general is different, since this new construction places different amounts of weight on different training data points, depending on the sizes of the various groups. However, as a special case, if $N_{k}=1$ for all $k$ (i.e., one observation in each group), then HCP simply reduces to the split conformal method.

The HCP prediction interval offers the following distribution-free guarantee:

In Appendix A, we provide analogous hierarchical versions of the jackknife+ and full conformal methods, to avoid the loss of accuracy incurred by sample splitting.

In this section, we summarize the work of Dunn et al. (2022) for the hierarchical data setting and compare to our proposed HCP method. Dunn et al. (2022) introduce four approaches for distribution-free inference with hierarchical data, which they refer to as Pooling CDFs, Double Conformal, Subsampling Once, and Repeated Subsampling. Here we only consider the split conformal versions of their methods, in the context of the supervised learning problem (i.e., predicting $Y$ from $X$), to be directly comparable to the setting of our work.

In the remainder of this section, we rewrite their methods to be consistent with our notation and problem setting for ease of comparison. As we detail below, the comparison between HCP and Dunn et al. (2022)’s four methods reveals:

• •

  Dunn et al. (2022)’s theoretical guarantee for Pooling CDFs offers only an asymptotic result. In Proposition 1, we show that Pooling CDFs is actually exactly equivalent to running HCP but with a slightly higher error level $\alpha$, and consequently, establish a novel finite-sample coverage guarantee.

• •

  Double Conformal was shown by Dunn et al. (2022) to have finite-sample coverage at level $1-\alpha$, but in practice is substantially overly conservative.

• •

  Subsampling Once was also shown by Dunn et al. (2022) to have finite-sample coverage at level $1-\alpha$, but in practice can lead to more variable performance due to discarding most of the calibration data.

• •

  Dunn et al. (2022)’s theoretical guarantee for Repeated Subsampling establishes coverage at the weaker level $1-2\alpha$, even though empirical performance shows that undercoverage does not occur. In our work (Proposition 2), we establish that Repeated Subsampling is equivalent to running HCP on a bootstrapped version of the original calibration data set, and consequently, obtain a novel guarantee of finite-sample coverage at level $1-\alpha$.

Overall, our proposed framework of hierarchical exchangeability allows us to build a new understanding of the finite-sample performance of two of Dunn et al. (2022)’s methods, Pooling CDFs and Repeated Subsampling—that is, the two whose empirical performance was observed to be the best, without high variance or overcoverage.

As for HCP, for all of Dunn et al. (2022)’s methods the first step is to use the training portion of the data (groups $k=1,\dots,K_{0}$) to train a score function $s:\mathcal{Z}\rightarrow\mathbb{R}$, and the next step is to define a prediction interval $\widehat{C}(X_{\textnormal{test}})=\left\{y\in\mathcal{Y}:s(X_{\textnormal{test}},y)\leq T\right\}$, for some threshold $T$ determined by the calibration set. We next give details for how $T$ is selected for each of their proposed methods.

We next present simulation results to illustrate the comparison between HCP and the methods proposed by Dunn et al. (2022).⁶⁶6Code to reproduce this simulation is available at https://github.com/rebeccawillett/Distribution-free-inference-with-hierarchical-data.

In the setting of i.i.d. data (i.e., without repeated measurements), when the features $X$ have a continuous or nonatomic distribution, Vovk et al. (2005); Lei and Wasserman (2014); Barber et al. (2021a) establish that conditionally-valid predictive inference (that is, guarantees of the form $\mathbb{P}\left\{{Y_{\textnormal{test}}\in\widehat{C}(X_{\textnormal{test}})}\ \middle|\ {X_{\textnormal{test}}}\right\}\geq 1-\alpha$) are impossible to achieve with finite-width prediction intervals. However, access to repeated measurements may provide additional information that allows us to move beyond this impossibility result. In this section, we discuss possible targets of distribution-free inference beyond the marginal coverage guarantee, by making use of the repeated measurements.

Consider the miscoverage rate conditional on all the observations we have—both the observed training and calibration data, $\mathcal{D}=\{(X_{k},Y_{k,1},\dots,Y_{k,N_{k}})\}_{k=1,\dots,K}$, and the test feature vector, $X_{\textnormal{test}}$. Specifically, define

where the probability is taken with respect to the distribution of the test response $Y_{\textnormal{test}}$ (which is drawn from $P_{Y|X}$ at $X_{\textnormal{test}}=x$).⁷⁷7Note that the random variable $\alpha_{\mathcal{D}}(X_{\textnormal{test}})$ is therefore the coverage conditional on both the observed training and calibration data $\mathcal{D}$, and test feature $X_{\textnormal{test}}$, meaning that we are now examining an even stronger form of conditional coverage. However, it is well known that conditioning on $\mathcal{D}$ is straightforward for split conformal prediction (Vovk, 2012), that is, the primary challenge in bounding $\alpha_{\mathcal{D}}(X_{\textnormal{test}})$ comes from conditioning on $X_{\textnormal{test}}$. The marginal coverage guarantee

can then equivalently be expressed as

by the tower law.

For the goal of controlling conditional miscoverage, an ideal target would be

This condition, if achieved, would ensure that the resulting prediction set provides a good coverage for any value of new input $X_{\textnormal{test}}$, rather than only on average over $X_{\textnormal{test}}$. However, even in the setting of repeated measurements, if $X$ is continuously distributed (or more generally, nonatomic, i.e., $P_{X}(x)=0$ for all $x\in\mathcal{X}$), the guarantee (13) is not achievable by any nontrivial (finite-width) prediction interval.

Consequently, we would like to find a coverage target that is stronger than the basic marginal guarantee (12) but weaker than the unachievable conditional guarantee (13). As a compromise between the two, we consider the following guarantee, which we refer to as second-moment coverage:

This condition is a compromise between the previous two: second-moment coverage (14) is strictly weaker than the (unattainable) conditional coverage guarantee (13), and strictly stronger than the marginal coverage guarantee (12).

Another way of understanding the guarantee (14) as a mechanism for controlling conditional miscoverage is to look at the tail probability. By Markov’s inequality, for any constant $c>1$, the marginal coverage guarantee leads to

while the stronger guarantee (14) implies

Hence, we can expect more uniformly small conditional miscoverage rate from the second-moment coverage guarantee.

Now we discuss how the stronger guarantee (14) can be achieved in the distribution-free sense, in the setting of repeated measurements. We propose a modification of our HCP method, which we call HCP²—this name refers to the fact that the new goal is second-moment coverage. (As before, here we aim to extend the split conformal method to address this goal; extensions of other methods—namely, full conformal and jackknife+—are deferred to the Appendix.)

To begin, let $K_{1}^{\geq 2}=\sum_{k=K_{0}+1}^{K}{\mathbbm{1}}_{{N_{k}\geq 2}}$, the number of data points in the calibration set for which we have at least two repeated measurements.

where $s_{k,(1)}\leq\dots\leq s_{k,(N_{k})}$ are the order statistics of the scores $s(Z_{k,1}),\dots,s(Z_{k,N_{k}})$ in the $k$th calibration group.

To understand the intuition behind this construction, we can first verify that the threshold $T$ in this definition can equivalently be written as

With this equivalent definition in place, we can now explain the main idea behind the construction of the HCP² method. For intuition, suppose $Y\mid X$ has a continuous distribution. For each group $k$, since the $Y_{k,i}$’s are i.i.d. conditional on $X=X_{k}$, if $q_{1-a}$ is equal to the $(1-a)$-quantile of the distribution of $s(X,Y)$ conditional on $X=X_{k}$, then for constructing HCP (and proving marginal coverage), we observe that $\mathbb{P}\left\{{s(X_{k},Y_{k,i})>q_{1-a}}\right\}=a$ by construction. But, if we have at least two measurements $1\leq i<i^{\prime}\leq N_{k}$ at this particular $k$, then we can also see that $\mathbb{P}\left\{{\min\{s(X_{k},Y_{k,i}),s(X_{k},Y_{k,i^{\prime}})\}>q_{1-a}}\right\}=a^{2}$. In other words, by examining the pairwise minimums across the scores at each value of $X$, we can learn about $\alpha_{\mathcal{D}}(X)^{2}$, rather than learning only about $\alpha_{\mathcal{D}}(X)$.