Provably Calibrated Regression Under Distribution Drift

1 INTRODUCTION . Accurate uncertainty quantification is crucial for machine learning predictions used in high-stakes decision making . Typically , uncertainty is represented by probability distributions over the possible outcomes , and these probabilities should be calibrated . In the regression setup , for example , the true label should be below the predicted 95 % quantile for 95 % of the samples ( Gneiting et al. , 2007 ) . Calibration can convey confidence to decision makers because extreme values ( e.g . true label above 95 % quantile ) are guaranteed to be rare . In addition to calibration , these probabilities should be sharp ( i.e . concentrated and have low variance ) . Sharp probabilities are useful to decision makers because they are informative . If data is i.i.d. , then recalibration ( Kuleshov et al. , 2018 ) and conformal prediction ( Vovk et al. , 2020 ) algorithms can achieve low calibration error and good sharpness in the regression setup . However , the i.i.d . assumption is unlikely to hold in most time-series prediction tasks . Under distribution drift , it is possible to adapt regret minimization algorithms ( Cesa-Bianchi & Lugosi , 2006 ; Kuleshov & Ermon , 2017 ) to achieve calibration . However , regret minimization calibration algorithms are designed for the asymptotic regime , and empirically we show that they are effective only with large sample size , and have very poor calibration and sharpness with short time series ( e.g . 50 samples ) , limiting their practical utility . Our goal is to design an algorithm to achieve good calibration and sharpness in the online regression setup even for short time series . We start with an existing prediction algorithm that has good calibration and sharpness for i.i.d . data ( such as conformal prediction ) , and track the empirical frequency that the labels are below e.g . the 75 % quantile . If the empirical frequency is significantly below 75 % , we move future predictions up so more labels are below the prediction ; vice versa . We design a very efficient adjustment method that can guarantee near perfect calibration with only tens of samples . We call this the “ basic ” prediction algorithm shown in Figure 1 ( upper right ) . Our main technical contribution is an improved algorithm that addresses two short-comings of the “ basic ” algorithm without breaking guaranteed calibration . First , the predictions should be feasible , e.g . the 75 % quantile should never be smaller than the 25 % quantile . The basic algorithm might violate this requirement because we adjust each quantile separately with no constraint on their relation . Second , the predictions should be stable if the distribution drift stops ; empirically instability harms sharpness and proper scores . Our improved algorithm is based on an analogy between the prediction task and a mechanical system . We encode the feasibility constraints as the feasible states of a mechanical system , and use proportional-integral-derivative ( PID ) control ( Minorsky , 1922 ) to stabilize the system . Our final algorithm makes predictions that satisfy all three desirable properties : provable calibration , feasibility , and stability ( Figure 1 bottom right ) . We test our algorithm on two real world time series datasets ( stock earnings and COVID cases ) and a large benchmark of 17 common regression datasets . We reduce calibration error by approximately 2x compared too all baselines , while achieving comparable or better sharpness . We also simulate COVID response decisions based on predictions on the COVID dataset , and observe improved decision loss compared to baselines during periods of distribution drift ( e.g . when case numbers surge ) . 2 PROBLEM SETUP AND BACKGROUND . We consider regression problems . Let X ∈ X denote the input feature and Y denote the label . We assume that the label is bounded by B , i.e . Y can not take any value outside [ −B , B ] for some pre-specified B > 0 . We consider quantile predictions with a set of K equally spaced quantiles 0 < α1 < · · · < αK < 1 . For example , when K = 9 we choose α1 = 10 % , α2 = 20 % , · · · , αK = 90 % . A quantile prediction is a vector of K numbers ( denoted by Z ) where Zk should ideally predict the αk-th quantile of the label Y |X conditioned on the input feature X . We say a quantile prediction is feasible , if Yk < Yk+1 , ∀k , i.e . a lower quantile is always smaller than a higher quantile . In the limit of infinitely many quantiles K → ∞ , a quantile prediction is equivalent to a probability distribution , however , for technical reasons that will become evident later , we only consider finitely many quantiles in this paper . Our setup is an online prediction setup , where the forecaster sequentially makes the predictions . Consider a COVID prediction example : at time t , the forecaster observes some features Xt ( such as current vaccination rate ) and makes a prediction Zt about next week ’ s case number . After the forecaster makes a prediction , the true case number Y t is revealed ; then we move on to time step t + 1 and repeat this process . Following standard terminology we will use “ nature ” to refer to the process that generates the featuresXt and the true label Y t. Our setup is formalized by the following interaction between forecaster and nature : for t = 1 , 2 , · · · , T 1 . Nature reveals feature Xt 2 . Forecaster makes a quantile prediction Zt ∈ RK where Ztk is the αk-th quantile . 3 . Nature reveals label Y t ∈ [ −B , B ] We call a finite sequence of interactions a transcript , i.e . a sequence X1 , Z1 , Y 1 , · · · , XT , ZT , Y T . A forecasting algorithm ψ is a function that maps a transcript and a new feature to a prediction : ψ : X1 , Z1 , Y 1 , · · · , Xt−1 , Zt−1 , Y t−1 , Xt 7→ Zt . ( 1 ) We will consider two types of assumptions on nature . We say that nature is i.i.d . if ∀t , Xt , Y t are drawn from some i.i.d . random variable ( but we make no assumption on the probability law on this random variable ) . We say that nature is causal if nature chooses Xt , Yt without depending on the future . Specifically , Xt , Y t can only depend on the variables that precede it in the transcript ( i.e . Xt can depend on X1 , · · · , Y t−1 and Y t can only depend on X1 , · · · , Y t−1 , Xt , Zt ) . Causal is an extremely weak assumption . For example , nature can even adversarially choose the label Y t to increase prediction error after observing the forecaster ’ s prediction Zt . 2.1 CALIBRATION AND ACHIEVING CALIBRATION WITH CONFORMAL METHODS . A natural property that we can request is ( probabilistic ) calibration ( Gneiting et al. , 2007 ; Kuleshov et al. , 2018 ) . Intuitively , ∀k , the label Y t should be below the predicted αk-th quantile for an αk proportion of the samples . Formally , given a transcript of any length , consider the empirical frequency that the label Y t is below the predicted quantile Ztk up to time T ( for any T that ’ s less than the length of transcript ) FTk = ∑ 1≤t≤T I ( Y t ≤ Ztk ) . ( 2 ) Ideally , the label Y t should be below the αk-th quantile exactly αkT many times , i.e. , ∀k , FTk /T = αk . Achieving this perfectly is difficult , so we define an approximation . For any function b : N → R+ we say that a transcript is b-calibrated if ∀k = 1 , · · · , K , ∀T , |FTk /T − αk| ≤ b ( T ) . Correspondingly , a forecasting strategy is calibrated if the resulting transcript is calibrated : , Definition 1 . For any function b : N → R+ , a forecasting algorithm ψ is b-calibrated under i.i.d . ( or causal ) assumptions if for any nature ’ s strategy that is i.i.d . ( or causal ) , the resulting transcript is b-calibrated almost surely . Conformal Calibration Conformal calibration ( Vovk et al. , 2020 ) is a forecasting algorithm based on conformal prediction ideas ( Vovk et al. , 2005 ; Shafer & Vovk , 2008 ) . It is a wrapper algorithm that transforms a initial predictor ( such as an off-the-shelf predictor ) into new predictions that are provably calibrated under i.i.d . assumptions . Conformal calibration is based on the following intuition : for example , if 75 % of past labels are below the initial prediction ’ s mean , then we can use the initial prediction ’ s mean as the 75 % -quantile prediction in the future . More generally , we identify a statistic of the initial prediction , such that αk-proportion of the past labels are below that statistic . We use this statistic as the αk-th quantile of the future prediction . If the data is i.i.d. , conformal calibration is a very useful algorithm . This is because it has extremely strong calibration properties ( Proven in Proposition 1 of ( Vovk et al. , 2020 ) ) , where the random variables I ( Y t ≤ Ztk ) , t = 1 , 2 , · · · is a sequence of i.i.d . Bernoulli random variables with mean αk . This is impressive because even if we have access to the oracle forecaster ( i.e . Y tk is equal to the true conditional αk-th quantile of Y t|Xt ) , we can not achieve better calibration — I ( Y t ≤ Ztk ) is also a sequence of i.i.d . Bernoulli random variables with mean αk . Empirically , conformal prediction also has very good sharpness if the initial prediction function is reasonable ( Burnaev & Vovk , 2014 ) . A minor property ( but needed for proofs ) is that the predictions should bounded in [ −B , B ] . Conformal calibration satisfies this property when t > K , i.e . there are more samples than predicted quantiles . To bypass this limitation , we can initialize the conformal calibration algorithm with at least K offline samples . Alternatively , if we know B in advance , we can clip the prediction by B . 3 CALIBRATION UNDER DISTRIBUTION DRIFT . This section introduces our algorithm that guarantees calibration even when nature is not i.i.d . We start from a basic algorithm in Section 3.1 , and discuss its shortcomings . Then we move on to more advanced algorithms in Section 3.2 and 3.3 to address the shortcomings . 3.1 A BASIC CALIBRATED PREDICTION ALGORITHM . Conformal calibration achieves good calibration and sharpness when nature is i.i.d . When we are unsure if nature is i.i.d. , our basic idea is to use conformal calibration until it has failed to achieve calibration . Specifically , when nature is i.i.d. , I ( Y t ≤ Ztk ) is a sequence of Bernoulli random variables , so FTk is the sum of these Bernoulli random variables ( i.e . a binomial ) . We use b ∗ δ ( T ) to denote the confidence interval of a binomial distribution , i.e . Pr [ FTk /T ∈ αk ± b∗δ ( T ) ] = 1 − δ . We fix some small δ ( such as 0.05 ) and if under the conformal calibration algorithm FTk /T 6∈ αk ± b∗δ ( T ) , we know ( with 0.95 confidence ) that nature is not i.i.d , so we will make adjustments to salvage calibration . From a high level , F tk > αkt + b ∗ δ ( t ) t implies that the true label is below the αk-th quantile too often , so our prediction Ztk has been too large and we should reduce it ; vice versa ; There are some design freedom in choosing how much to reduce or increase the prediction . We choose an adjustment that grows exponentially with larger calibration error . The reason for this choice is to make very large adjustments when the calibration error is large , so we can tightly bound the calibration error . Formally , let Z̃tk denote the initial prediction generated by the conformal calibration prediction algorithm , the prediction of our algorithm Y tk is given by Zt+1k = Z̃ t+1 k + E t+1 k , where E t+1 k = 1− e β ( F tk−αkt−b ∗ δ ( t ) t ) F tk > αkt+ b ∗ δ ( t ) t eβ ( αkt−b ∗ δ ( t ) t−F t k ) − 1 F tk < αkt− b∗δ ( t ) t 0 otherwise ( 3 ) Intuitively , if F tk is within the correct range we make no adjustments ( i.e . Z t+1 k = Z̃ t+1 k ) ; when F tk falls outside the correct range , we make aggressive adjustments that increase exponentially with how much it falls outside the correct range . β > 0 is a hyper-parameter that controls how large the adjustments are ( we will show in Section 3.4 that these hyper-parameters are easy to choose ) . Note that Et+1k is for the ( t + 1 ) -th time step ( instead of t ) . This is because any adjustments can only depend on past information , so we will only have access to F tk at the ( t+ 1 ) -th time step . The following theorem shows that such adjustments indeed guarantee calibration . Theorem 1 . Let Z̃tk be generated by any forecasting algorithm and bounded in [ −B , B ] , for any δ > 0 , the prediction algorithm defined by Eq . ( 3 ) is b-calibrated where b ( T ) = b∗δ ( T ) + log ( 2B+1 ) +β βT . Note that the guarantee of Theorem 1 does not require that conformal calibration generate the initial prediction Z̃tk . The use of conformal calibration is motivated by its strong empirical performance ( Vovk et al. , 2020 ) when the data is i.i.d . or close to i.i.d . ( which is further supported by our experiments ) . The calibration guarantee is strong because under conformal calibration with i.i.d . data ( or the oracle forecaster ) we can expect to see b∗δ ( T ) calibration error ( with 1 − δ probability ) ; without i.i.d . assumptions , our algorithm can ensure that the calibration error only increases by O ( logB/T ) . Notably , the calibration error scales logarithmically with the assumed upper bound logB , so choosing a loose bound B does not significantly degrade the guarantee . The basic idea has two shortcomings illustrated in Figure 1 that we will address in the next sections . First , in a feasible quantile prediction , the 75 % quantile should never be smaller than the 25 % quantile , but our adjustments might violate this requirement because we adjust each quantile separately with no constraint . Second , the basic idea might lead to oscillation and instability . Intuitively , when the empirical frequency is incorrect , we make adjustments to correct it , but once the empirical frequency improves , we reduce the adjustments , which will cause the empirical frequency to become incorrect again , leading to oscillation . Empirically instability can hurt performance metrics such as sharpness .

Summary: For real-world applications where ML models are used, quantifying predictive uncertainty is an essential task as, for example, in safety-critical applications misclassification might have disastrous consequences. On the theory side, the analysis is typically performed under the i.i.d. assumption which could be unsatisfactory as deployed models inevitably encounter changes in the data generating distribution and/or certain dependencies that invalidate the results established for the i.i.d. setting. Focusing on regression, the authors design an adaptive (to distribution shifts) procedure that satisfies a certain calibration guarantee. The authors study the empirical performance of the proposed procedure on a collection of time-series and regression datasets.