Improving Conditional Coverage via Orthogonal Quantile Regression

1 Introduction . Learning algorithms are increasingly prevalent within consequential real-world systems , where reliability is an essential consideration : confidently deploying learning algorithms requires more than high prediction accuracy in controlled testbeds [ 1 , 2 ] . Consider , for example , estimating the effects of a drug for a specific person given their demographic information and medical measurements . In such a high-stakes setting , giving a point prediction for the drug ’ s effect is insufficient ; the decision-maker must know what the plausible range of effects for this specific individual . Instance-wise uncertainty quantification in such settings is critical [ 3–5 ] . One approach to this problem comes from the quantile regression and prediction interval literature [ 6–8 ] ; instead of a point prediction , we can return a range of outcomes that represent the plausible response for a given input . We would like these prediction intervals to achieve a pre-specified coverage level ( e.g. , 90 % ) for all inputs—that is , across all regions of feature space . Training models to satisfy this validity guarantee is challenging , however , particularly with complex models like neural networks [ 9 , 10 ] . In this work , we show how to generate prediction intervals that achieve coverage closer to the desired level evenly across all sub-populations . Technically , we achieve this by augmenting the quantile regression loss function with an additional term that promotes appropriately balanced coverage across the feature space . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Formally , consider a regression problem where we are given n training samples { ( Xi , Yi ) } ni=1 , where X ∈ Rp is a feature vector , and Y ∈ R is a response variable . At test time , we observe a feature vector Xn+1 and our goal is to predict the unknown value of Yn+1 and—importantly—to report on its uncertainty . In this work , we represent this uncertainty by constructing a prediction interval Ĉ ( Xn+1 ) ⊆ R that is likely to contain the response Yn+1 . In particular , we seek to produce intervals that contain the response with a user-specified probability 1− α that are valid across all regions of feature space : P [ Yn+1 ∈ Ĉ ( Xn+1 ) | Xn+1 = x ] ≥ 1− α , ( 1 ) a property known as conditional coverage . Notice that such prediction intervals are both correct in that they satisfy a coverage guarantee and also are adaptive , in that the size of the prediction intervals can change with the difficulty of the inputs : easy inputs give small intervals and hard inputs give large intervals . Returning to our medical example , consider predicting the outcome of a drug from age , gender , blood pressure , and so on . The conditional coverage requirement in ( 1 ) asks that intervals are correct for any age , gender , and health status combination . That is , no matter what an individual ’ s value of the features x , the uncertainty quantification must be valid . The conditional quantiles of Y | X = x are the natural way to produce intervals satisfying ( 1 ) . Let αlo = α/2 and αhi = 1− α/2 . Given the true conditional quantiles , qαlo ( x ) , qαhi ( x ) , we can build an oracle prediction interval satisfying ( 1 ) in the following way : C ( x ) = [ qαlo ( x ) , qαhi ( x ) ] . ( 2 ) In practice , the conditional quantiles are unknown but can be estimated with quantile regression , yielding the interval Ĉ ( x ) = [ q̂αlo ( x ) , q̂αhi ( x ) ] . This approach is attractive because quantile regression yields intervals that are adaptive to heteroscedasticity without requiring parametric assumptions [ 11– 13 ] , but these intervals might not satisfy the conditional coverage statement ( 1 ) , since q̂αlo ( x ) , q̂αhi ( x ) are merely estimations of the true quantiles [ 14 ] . Indeed , we observe in experiments 5 that traditional quantile regression often gives intervals with poor conditional coverage , prompting the present investigation . In this work , we propose a novel regularization scheme to push quantile regression algorithms towards solutions that better satisfies the conditional coverage requirement ( 1 ) . The core idea is to force the coverage and interval length to be approximately independent , since this independence must hold for the optimal oracle intervals in ( 2 ) . A method that constructs intervals whose coverage and length are dependent is either sometimes too conservative ( generating too wide intervals ) , sometimes too liberal ( yeilding too short intervals ) , or both . In addition to improved training schemes , we propose two new tools to check the validity of the resulting predictions in a meaningful way . Specifically , in Section 4 , we present two new interpretable metrics to asses the violation of conditional coverage , taking advantage of the orthogonality property identified above . We use these ( and other ) metrics in Section 5 to study our proposal on simulated data and nine real benchmark data sets . We find that our training scheme yields improvements when used together with both a classic [ 6 ] and a more recent [ 15 ] quantile regression method . A synthetic two-group example . We begin with a small synthetic experiment that demonstrates the challenges of constructing prediction intervals with accurate conditional coverage . We generate a dataset with two unbalanced subpopulations : 80 % of the samples belong to a majority group and the remaining 20 % to a minority group , where the conditional distribution Y | X of the minority group is more dispersed than the majority group . In our experiments , group membership is included as one of the features . See Section S3.1 of the Supplementary Material for a full description of the distribution . As a baseline method , we first fit a quantile neural network model ( vanilla QR ) by optimizing the pinball loss ( see Section 2 ) , attempting to estimate the low αlo = 0.05 and high αhi = 0.95 conditional quantiles . The left panel of Figure 1 shows the coverage obtained by the vanilla QR model across training epochs . The coverage on the test data is far lower than suggested by the training data , failing to reach the desired 90 % level and increasing as the training progresses . In particular , this gap remains large at the epoch in which the model achieves the minimal loss evaluated on an independent validation set . Here , the empirical test coverage measured over the majority and minority groups is equal to 80 % and 68 % , and the empirical average lengths evaluated over them are 1.55 , 6.45 , respectively . Next , we fit our proposed orthogonal QR model ( which we will formally introduce in Section 3 ) on the same data . The coverage rate across epochs is illustrated in the right panel of Figure 1 . In contrast to vanilla QR , the training and testing curves of the majority group overlap for the entire training epochs and both approximately reach the desired 90 % level . The minority group coverage has similar behavior , with a significantly smaller gap between the two curves compared to vanilla QR . Here , the model that corresponds to the best epoch achieves 86 % coverage rate with 1.62 average length on the majority group , and 83 % coverage rate with 9.33 average length on the minority group . Importantly , here we are able to train for more epochs before overfitting , which leads to a better final model . To conclude , in this example orthogonal QR prevents overfitting and leads to better conditional coverage . 2 Preliminaries and related work . 2.1 Quantile regression . The task of estimating the low and high conditional quantiles , q̂αlo ( x ) , q̂αhi ( x ) , can be expressed as a minimization problem taking the following form : ( q̂αlo , q̂αhi ) = argmin fαlo , fαhi∈F 1 n n∑ i=1 ` α ( Yi , fαlo ( Xi ) , fαhi ( Xi ) ) . ( 3 ) Above , ` α is a loss function designed to fit quantiles ; we discuss two examples next . The most common loss function for estimating a conditional quantile q̂α is called the pinball loss or check function [ 6 , 9 , 16 ] , expressed as ρα ( y , ŷ ) = { α ( y − ŷ ) y − ŷ > 0 , ( 1− α ) ( ŷ − y ) otherwise . ( 4 ) Observe that for the choice α = 1/2 the above becomes the L1 norm that is known to estimate the conditional median . By setting α 6= 1/2 we get a tilted version of the latter , with a degree controlled by the value of α . In the context of ( 3 ) , we can set ` pbα ( y , fαlo ( x ) , fαhi ( x ) ) = ραlo ( y , fαlo ( x ) ) + ραhi ( y , fαhi ( x ) ) to simultaneously estimate the low and high quantiles . Throughout this paper , we will refer to this procedure as vanilla QR . A recently-developed alternative to the pinball loss is the interval score loss [ 17 ] , defined as : ` intα ( y , fαlo ( x ) , fαhi ( x ) ) = ( fαhi ( x ) − fαlo ( x ) ) + 2 α ( fαlo ( x ) − y ) 1 [ y < fαlo ( x ) ] + 2 α ( y − fαhi ( x ) ) 1 [ y > fαhi ( x ) ] . ( 5 ) Note that the left-most term encourages short intervals and the two remaining components promote intervals with the right coverage . To improve statistical efficiency , it is recommended by [ 15 ] to simultaneously estimate all conditional quantiles by minimizing the following empirical risk function : Eα∼U [ 0,1 ] [ ` intα ( · ) ] .1 This is the approach we take in our experiments . Note that [ 15 ] proposed additional learning schemes for improving efficiency , such as group batching and ensemble learning ; see also [ 18 ] . These ideas are complementary to our proposal and may further improve the performance of our proposed orthogonal QR . Quantile regression is a large , active area of research , and we finish by pointing out a few representative strands of work in the literature . Estimates of conditional quantile functions using the pinball loss with specific models are proven to be asymptotically consistent under regularity conditions [ 16 ] . The work reported in [ 11–13 ] offer a non-parametric version of quantile regression . This line of research was further developed by [ 19 ] that presented a generalization to additive models that are non-parametric . The pinball loss and interval score loss can be also used to estimate conditional probability distribution [ 17 , 20 ] . Nevertheless , quantiles can be estimated in other ways rather than minimizing these two loss functions . These include quantile random forest [ 8 ] and the method proposed in [ 21 ] that iteratively estimates the conditional quantile through the Majorize-Minimize algorithm . Besides the regularization term we propose , there are other suggested penalties that are useful in different situations , such as using sparse modeling in high dimensional response [ 22–24 ] .

Summary: The authors propose a novel method to construct prediction intervals with pre-specified conditional coverage probability. Starting point of the paper is the observation that the length $|\hat{C}(X)|$ and conditional coverage indicator $1\{Y \in \hat{C}(X)\}$ are independent whenever the prediction interval $\hat{C}(X)$ has exact coverage probability.  The authors therefore suggest to compute lower and upper bounds of the prediction interval by solving a quantile regression problem augmented by a penalty term that penalizes correlation between $|\hat{C}(X)|$ and $1\{Y \in \hat{C}(X)\}$. The authors compare the performance of their procedure and competing methods on simulated data and nine benchmark data sets.