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: This paper considers the problem of calibrating quantiles outputted by a regression model. The data is assumed to draw from a non-iid source (e.g. time series with distribution shift). The paper proposes a method that adjusts the output quantiles to be better calibrated for a less-restricted definition of b-calibrated. Some further methods are also proposed to ensure the quantiles are monotonic and stable. Experiments are conducted to compare the proposed method with uncalibrated models and conformal predictions.