Original: Instance-Based Uncertainty Estimation for Gradient-Boosted Regression Trees
Revised: Instance-Based Uncertainty Estimation for Gradient-Boosted Regression Trees

Original: Anonymous Author(s)

Original: AffiliationAddress

Original: Abstract
Revised: Abstract

Original: Gradient-boosted regression trees (GBRTs) are hugely popular for solving tabularregression problems, but provide no estimate of uncertainty.
Revised: Gradient-boosted regression trees (GBRTs) are hugely popular for solving tabular regression problems, but provide no estimate of uncertainty.

Original: We propose InstanceBased Uncertainty estimation for Gradient-boosted regression trees (IBUG), asimple method for extending any GBRT point predictor to produce probabilisticpredictions.
Revised: We propose InstanceBased Uncertainty estimation for Gradient-boosted regression trees (IBUG), a simple method for extending any GBRT point predictor to produce probabilistic predictions.

Original: IBUG computes a non-parametric distribution around a predictionusing the k -nearest training instances, where distance is measured with a treeensemble kernel.
Revised: IBUG computes a non-parametric distribution around a prediction using the k -nearest training instances, where distance is measured with a treeensemble kernel.

Original: The runtime of IBUG depends on the number of training examplesat each leaf in the ensemble, and can be improved by sampling trees or traininginstances.
Revised: The runtime of IBUG depends on the number of training examples at each leaf in the ensemble, and can be improved by sampling trees or training instances.

Original: Empirically, we find that IBUG achieves similar or better performancethan the previous state-of-the-art across 22 benchmark regression datasets.
Revised: Empirically, we find that IBUG achieves similar or better performance than the previous state-of-the-art across 22 benchmark regression datasets.

Original: We alsofind that IBUG can achieve improved probabilistic performance by using differentbase GBRT models, and can more flexibly model the posterior distribution of aprediction than competing methods.
Revised: We also find that IBUG can achieve improved probabilistic performance by using different base GBRT models, and can more flexibly model the posterior distribution of a prediction than competing methods.

Original: We also find that previous methods sufferfrom poor probabilistic calibration on some datasets, which can be mitigated usinga scalar factor tuned on the validation data. 
Revised: We also find that previous methods suffer from poor probabilistic calibration on some datasets, which can be mitigated using a scalar factor tuned on the validation data. Source code is available at https: //github.com/jjbrophy47/ibug .

Original:  Introduction
Revised:  Introduction

Original: Despite the impressive success of deep-learning models on unstructured data (e.g., images, audio,text), gradient-boosted trees [24] remain the preferred choice for tabular or structured data [50].
Revised: Despite the impressive success of deep learning models on unstructured data (e g., images, audio, text) gradient-boosted trees [26] remain the preferred choice for tabular or structured data [52].

Original: Infact, Kaggle CEO Anthony Goldbloom recently described gradient-boosted trees as the most “glaringdifference” between what is used on Kaggle and what is “fashionable in academia” [32].
Revised: In fact Kaggle CEO Anthony Goldbloom recently described gradient-boosted trees as the most “glaring difference” between what is used on Kaggle and what is “fashionable in academia” [34].

Original: Our focus is on tabular data for regression tasks, which vary widely from financial [1] and retailproduct forecasting [41] to weather [27, 26] and clinic-mortality prediction [6].
Revised: Our focus is on tabular data for regression tasks, which vary widely from financial [1] and retailproduct forecasting [43] to weather [29, 28] and clinic-mortality prediction [6].

Original: Gradient-boostedregression trees (GBRTs) are known to make accurate point predictions [40] but provide no estimateof the prediction uncertainty, which is desirable for both forecasting practitioners [9, 66] and theexplainable AI (XAI) community [23, 67, 2] in general.
Revised: Gradient-boosted regression trees (GBRTs) are known to make accurate point predictions [42] but provide no estimate of the prediction uncertainty, which is desirable for both forecasting practitioners [9, 68] and the explainable AI (XAI) community [25, 69, 2] in general.

Original: Recently, Duan et al.
Revised: Recently, Duan et al.

Original: [20] and Sprangerset al.
Revised: [22], Sprangers et al.

Original: [62] introduced NGBoost and PGBM (respectively), new gradient boosting algorithms thatprovide state-of-the-art probabilistic predictions.
Revised: [64], and Malinin et al. [44] introduced NGBoost, PGBM, and CBU (CatBoost [52] with uncertainty), respectively, new gradient boosting algorithms that provide state-of-the-art probabilistic predictions.

Original: However, NGBoost tends to underperform as apoint-predictor, and PGBM is limited in the types of distributions it can use to model the output.
Revised: However, NGBoost tends to underperform as a point-predictor, and PGBM and CBU are limited in the types of distributions they can use to model the output.

Original: We introduce a simple yet effective method for enabling any GBRT point-prediction model toproduce probabilistic predictions.
Revised: We introduce a simple yet effective method for enabling any GBRT point-prediction model to produce probabilistic predictions.

Original: Our proposed approach, Instance-Based Uncertainty estimationfor Gradient-boosted regression trees ( IBUG ), has two key components: 1) We leverage the factthat GBRTs accurately model the conditional mean and use this point prediction as the mean in aprobabilistic forecast; and 2)
Revised: Our proposed approach, Instance-Based Uncertainty estimation for Gradient-boosted regression trees ( IBUG ), has two key components: 1) We leverage the fact that GBRTs accurately model the conditional mean and use this point prediction as the mean in a probabilistic forecast; and 2)

Original: We identify the k training examples with the greatest affinity to the test instance and use these examples to estimate the uncertainty of the test prediction.
Revised: We identify the k training examples with the greatest affinity to the target instance and use these examples to estimate the uncertainty of the target prediction.

Original: We definethe affinity between two instances as the number of times both instances appear in the same leafthroughout the ensemble.
Revised: We define the affinity between two instances as the number of times both instances appear in the same leaf throughout the ensemble.

Original: Thus, our method acts as a wrapper around any given GBRT model.
Revised: Thus, our method acts as a wrapper around any given GBRT model.

Original: In experiments on 21 regression benchmark datasets and one synthetic dataset, we demonstratethe effectiveness of IBUG to deliver on par or improved probabilistic performance as compared toexisting state-of-the-art methods while maintaining state-of-the-art point-prediction performance.
Revised: In experiments on 21 regression benchmark datasets and one synthetic dataset, we demonstrate the effectiveness of IBUG to deliver on par or improved probabilistic performance as compared to existing state-of-the-art methods while maintaining state-of-the-art point-prediction performance.

Original: We also show that probabilistic predictions can be improved by applying IBUG to different GBRTmodels, something that NGBoost and PGBM cannot do.
Revised: We also show that probabilistic predictions can be improved by applying IBUG to different GBRT models, something that NGBoost and PGBM cannot do.

Original: Additionally, IBUG can use the traininginstances closest to the test example to directly model the output distribution using any parametric ornon-parametric distribution; again, something NGBoost and PGBM cannot do.
Revised: Additionally, IBUG can use the training instances closest to the target example to directly model the output distribution using any parametric or non-parametric distribution, something NGBoost, PGBM, and CBU cannot do.

Original: Finally, we show thatsampling trees dramatically improves runtime efficiency for computing training-example affinitieswithout having a significant detrimental impact on the resulting probabilistic predictions, allowing IBUG to scale to larger datasets.
Revised: Finally, we show that sampling trees dramatically improves runtime efficiency for computing training-example affinities without having a significant detrimental impact on the resulting probabilistic predictions, allowing IBUG to scale to larger datasets.

Original:  Notation & Background
Revised:  Notation & Background

Original: We assume an instance space X ⊆ R p and target space Y ⊂ R .
Revised: We assume an instance space X ⊆ R p and target space Y ⊂ R .

Original: Let D := { ( x i , y i ) } n i =1 be a trainingdataset in which each instance x
Revised: Let D := { ( x i , y i ) } ni =1 be a training dataset in which each instance x

Original: i ∈
Revised: i ∈

Original: X is a p -dimensional vector ( x ji ) pj =1 and y i ∈ Y .
Revised: X is a p -dimensional vector ( x j i ) p j =1 and y i ∈ Y .

Original: 2.1 Gradient-Boosted Regression Trees
Revised: 2.1 Gradient-Boosted Regression Trees

Original: Gradient-boosting [24] is a powerful machine-learning algorithm that iteratively adds weak learnersto construct a model f : X → R that minimizes some empirical risk L : R × R → R .
Revised: Gradient-boosting [26] is a powerful machine-learning algorithm that iteratively adds weak learners to construct a model f : X → R that minimizes some empirical risk L : R × R → R .

Original: The modelis defined by a recursive relationship: f 0 ( x ) = γ , . .
Revised: The model is defined by a recursive relationship: f 0 ( x ) = γ , . .

Original: , f t ( x ) =
Revised: , f t ( x ) =

Original: f t − 1 ( x ) + η m t ( x ) in which f 0 isthe base learner, γ is an initial estimate, f t is the model at iteration t , m t is the weak learner addedduring iteration t to improve the model, and η is the learning rate.
Revised: f t − 1 ( x ) + η m t ( x ) in which f 0 is the base learner, γ is an initial estimate, f t is the model at iteration t , m t is the weak learner added during iteration t to improve the model, and η is the learning rate.

Original: Gradient-boosted regression trees (GBRTs) typically choose ℓ to be the mean squared error (MSE), γas 1 n (cid:80) ni
Revised: Gradient-boosted regression trees (GBRTs) typically choose ℓ to be the mean squared error (MSE), γ as 1 n (cid:80) ni

Original: =

Revised: =

Original: y i (mean output of the training instances), and regression trees as weak learners.
Revised: y i (mean output of the training instances), and regression trees as weak learners.

Original: Each weak learner is typically chosen to approximate the negative gradient [42]: m t =arg min m 1 n (cid:80) ni =1 ( − g it − m ( x i ))
Revised: Each weak learner is typically chosen to approximate the negative gradient [44]: m t = arg min m 1 n (cid:80) ni =1 ( − g it − m ( x i ))

Original: 2 in which g it =
Revised: 2 in which g it =

Original: ∂ℓ ( y i , ˆ y i ) ∂
Revised: ∂ℓ ( y i , ˆ y i ) ∂

Original: ˆ y i is the functional gradient of the i thtraining instance at iteration t with respect to ˆ y i =
Revised: ˆ y i is the functional gradient of the i th training instance at iteration t with respect to ˆ y i =

Original: f t − 1 ( x i ) .
Revised: f t − 1 ( x i ) .

Original: The weak learner at iteration t recursively partitions the instance space into M t disjoint regions{ r jt } M t j =1 .
Revised: The weak learner at iteration t recursively partitions the instance space into M t disjoint regions { r jt } M t j =1 .

Original: Each region is called a leaf, and the parameter value for leaf j at tree 1 t is typicallydetermined (given a fixed structure) using a one-step Newton-estimation method [34]: θ jt =− (cid:80) i ∈ I
Revised: Each region is called a leaf, and the parameter value for leaf j at tree 2 t is typically determined (given a fixed structure) using a one-step Newton-estimation method [36]: θ jt = − (cid:80) i ∈ I

Original: jtw i g it / ( (cid:80) i ∈ I
Revised: jt g it / ( (cid:80) i ∈ I

Original: jt w
Revised: jt h

Original: i h it + λ ) in which I jt = { ( x i , y i )
Revised: it + λ ) in which I jt = { ( x i , y i )

Original: | x i ∈
Revised: | x i ∈

Original: r jt } ni
Revised: r jt } ni

Original: =1 is the instance set of leaf jfor tree t , h
Revised: =1 is the instance set of leaf j for tree t , h

Original: it is the second derivative of the i th training instance w.r.t.
Revised: it is the second derivative of the i th training instance w.r.t.

Original: ˆ y i , and λ is a regularizationconstant.
Revised: ˆ y i , and λ is a regularization constant.

Original: Thus, the output of m t can be written as follows: m t ( x ) = (cid:80) M t j =1 θ j t 1 [ x ∈ r j t ] in which 1is the indicator function.
Revised: Thus, the output of m t can be written as follows: m t ( x ) = (cid:80) M t j =1 θ jt 1 [ x ∈ r jt ] in which 1 is the indicator function.

Original: The final GBRT model generates a prediction for a test example x te bysumming the values of the leaves x te traverses to across all T iterations:
Revised: The final GBRT model generates a prediction for a target example x te by summing the values of the leaves x te traverses to across all T iterations:

Original: ˆ y te =
Revised: ˆ y te =

Original: (cid:80) Tt =1 m t ( x te ) .
Revised: (cid:80) T t =1 m t ( x te ) .

Original: 2.2 Probabilistic Regression
Revised: 2.2 Probabilistic Regression

Original: Our focus is on probabilistic regression—estimating the conditional probability distribution P ( y | x )for some target variable y given some input vector x ∈ X .
Revised: Our focus is on probabilistic regression—estimating the conditional probability distribution P ( y | x ) for some target variable y given some input vector x ∈ X .

Original: Unfortunately, traditional GBRTmodels only output scalar values.
Revised: Unfortunately, traditional GBRT models only output scalar values.

Original: Under a squared-error loss function, these scalar values can beinterpreted as the conditional mean in a Gaussian distribution with some (unknown) constant variance
Revised: Under a squared-error loss function, these scalar values can be interpreted as the conditional mean in a Gaussian distribution with some (unknown) constant variance.

Original: However homoscedasticity  a strong assumption and unknown constant variance has little value in a probabilistic prediction; thus, in order to allow heteroscedasticity, the predicted distribution needs at least two parameters to convey both the magnitude and uncertainty of the prediction [20]
Revised: However, homoscedasticity is a strong assumption and unknown constant variance has little value in a probabilistic prediction; thus, in order to allow heteroscedasticity, the predicted distribution needs at least two parameters to convey both the magnitude and uncertainty of the prediction [22].

Original: NGBoost is a recent method by Duan et al.
Revised: Natural Gradient Boosting (NGBoost) is a recent method by Duan et al.

Original: [20] that tackles the aforementioned problems byestimating the parameters of a desired distribution using a multi-parameter boosting approach thattrains a separate ensemble for each parameter of the distribution.
Revised: [22] that tackles the aforementioned problems by estimating the parameters of a desired distribution using a multiparameter boosting approach that trains a separate ensemble for each parameter of the distribution.

Original: NGBoost employs the naturalgradient to be invariant to parameterization, but requires the inversion of many small matrices(each the size of the number ofdesired parameters) to do so.
Revised: NGBoost employs the natural gradient to be invariant to parameterization, but requires the inversion of many small matrices (each the size of the number of parameters) to do so.

Original: Empirically, NGBoost generatesstate-of-the-art probabilistic predictions, but tends to underperform as a point predictor.
Revised: Empirically, NGBoost generates state-of-the-art probabilistic predictions, but tends to underperform as a point predictor.

Original: Even more recently, Sprangers et al.
Revised: More recently, Sprangers et al.

Original: [62] introduced PGBM, a single model that optimizes for pointperformance, but is able to generate accurate probabilistic predictions.
Revised: [64] introduced Probabilistic Gradient Boosting Machines (PGBM), a single model that optimizes for point performance, but can also generate accurate probabilistic predictions.

Original: PGBM treats leaf values asstochastic random variables, using sample statistics to model the mean and variance of each leaf value.
Revised: PGBM treats leaf values as stochastic random variables, using sample statistics to model the mean and variance of each leaf value.

Original: PGBM predicts the output mean and variance of a test example using the estimated parameters ofeach leaf it is assigned to.
Revised: PGBM estimates the output mean and variance of a target example using the estimated parameters of each leaf it is assigned to.

Original: The predicted mean and variance is then used as parameters in a specifieddistribution to generate a probabilistic prediction.
Revised: The predicted mean and variance are then used as parameters in a specified distribution to generate a probabilistic prediction.

Original: PGBM has been shown to produce state-of-the-artprobabilistic predictions; however, computing the necessary leaf statistics during training can becomputationally expensive, especially as the number of leaves in the ensemble increases (see §5.6).
Revised: PGBM has been shown to produce state-of-the-art probabilistic predictions; however, computing the necessary leaf statistics during training can be computationally expensive, especially as the number of leaves in the ensemble increases (see §5.6).

Original: Also, since only the mean and variance are predicted for a given test example, PGBM is limited todistributions using only location and scale to model the output. 
Revised: Also, since only the mean and variance are predicted for a given target example, PGBM is limited to distributions using only location and scale to model the output. Finally, Malinin et al.

Revised: [44] introduce CatBoost with uncertainty (CBU), a method that estimates uncertainty using ensembles of GBRT models.

Revised: Similar to NGBoost, multiple ensembles are learned to output the mean and variance.

Revised: However, CBU also constructs a virtual ensemble —a set of overlapping partitions of the learned GBRT trees—to estimate the uncertainty of a prediction by taking the mean of the variances output from the virtual ensemble.

Revised: Their approach uses a recently proposed stochastic gradient Langevin boosting algorithm [72] to sample from the true posterior via the virtual ensmble (in the limit); however, their formulation of uncertainty is limited only to the first and second moments, similar to PGBM.

Original: To address the shortcomings of existing approaches, we introduce a simple method that performswell on both point and probabilistic performance, and can more flexibly model the output than eitherapproach.
Revised: To address the shortcomings of existing approaches, we introduce a simple method that performs well on both point and probabilistic performance, and can more flexibly model the output than previous approaches.

Original: Our method can also be applied to any GBRT model, adding additional flexibility.
Revised: Our method can also be applied to any GBRT model, adding additional flexibility.

Original:  Instance-Based Uncertainty
Revised:  Instance-Based Uncertainty

Original: Instance-based methods such as k -nearest neighbors have been around for decades and have beenuseful for many different machine learning tasks [48].
Revised: Instance-based methods such as k -nearest neighbors have been around for decades and have been useful for many different machine learning tasks [50].

Original: However, defining neighbors based on a fixedmetric like euclidean distance may lead to suboptimal performance, especially as the dimensionalityof the dataset increases.
Revised: However, defining neighbors based on a fixed metric like Euclidean distance may lead to suboptimal performance, especially as the dimensionality of the dataset increases.

Original: More recently, it has been shown that random forests can be used as anadaptive nearest neighbors method [38] which identifies the most similar examples to a given instanceusing the learned model structure.
Revised: More recently, it has been shown that random forests can be used as an adaptive nearest neighbors method [18, 40] which identifies the most similar examples to a given instance using the learned model structure.

Original: This supervised tree kernel can more effectively measure thesimilarity between examples, and has been used for clustering [44] and local linear modeling [8] aswell as instance-[11] and feature-based attribution explanations [49], for example.
Revised: This supervised tree kernel can more effectively measure the similarity between examples, and has been used for clustering [46] and local linear modeling [8] as well as instance-[11] and feature-based attribution explanations [51], for example.

Original: In this work, we apply the idea of a supervised tree kernel to help model the uncertainty of agiven GBRT prediction.
Revised: In this work, we apply the idea of a supervised tree kernel to help model the uncertainty of a given GBRT prediction.

Original: Our approach, I nstance- B ased U ncertainty estimation for G radient-boostedregression trees (IBUG), identifies the neighborhood of similar training examples to a test exampleusing the structure of the GBRT, and then uses those instances to generate a probabilistic prediction.
Revised: Our approach, I nstance- B ased U ncertainty estimation for G radient-boosted regression trees (IBUG), identifies the neighborhood of similar training examples to a target example using the structure of the GBRT, and then uses those instances to generate a probabilistic prediction.

Original: IBUG works for any GBRT, and can more flexibly model the output than competing methods.
Revised: IBUG works for any GBRT, and can more flexibly model the output than competing methods.

Original: 3.1 Estimating Prediction Variance Using High-Affinity Neighbors
Revised: 3.1 Identification of High-Affinity Neighbors

Original: Given a GBRT model f and a test example x te , IBUG first models the magnitude of the test predictionusing the scalar output of f : µ ˆ y te = f

Original: ( x te ) since GBRTs already produce accurate point predictions.

Original: To estimate the variance of ˆ y te , IBUG computes the output variance of the k training examples withthe largest affinity to the test example.
Revised: At its core, IBUG uses the k training examples with the largest affinity to the target example to model the conditional output distribution.

Original: We define  affinity between two examples simply as the numberof times each instance appears in the same leaf across all trees in f .
Revised: Given a GBRT model f , we define the affinity between two examples simply as the number of times each instance appears in the same leaf across all trees in f .

Original: Thus, the affinity of the i thtraining example x i to the test example  can be written as:
Revised: Thus, the affinity of the i th training example x i to a target example x te can be written as:

Original: [Equation]
Revised: [Equation]

Original: (1) in which R t ( x i ) is the leaf
Revised: in which R t ( x i ) is the leaf

Original: x i is assigned to for tree t .
Revised: x i is assigned to for tree t .

Original: Alg.
Revised: Alg.

Original: 1 summarizes the procedure forcomputing affinity scores for all training examples.
Revised: 1 summarizes the procedure for computing affinity scores for all training examples.

Original: This metric clusters similar examples togetherbased on the learned model representation (i.e., the tree structure).
Revised: This metric clusters similar examples together based on the learned model representation (i.e., the tree structures).

Original: Intuitively, if two examplesappear in the same leaf in every tree throughout the ensemble, then both examples are predictedin an identical manner.
Revised: Intuitively, if two examples appear in the same leaf in every tree throughout the ensemble, then both examples are predicted in an identical manner.

Original: One may also view eq.
Revised: One may also view Eq.

Original: (1) as an indication of which training examples most often affect the leaf values x te is assigned to and thus implicitly which examples are likelyto have a big effect on the prediction ˆ y te .
Revised: (1) as an indication of which training examples most often affect the leaf values x te is assigned to and thus implicitly which examples are likely to have a big effect on the prediction ˆ y te .

Revised: This similarity metric is similar to the random forest kernel [18], however, unlike random forests, GBRTs are typically constructed to a shallower depth, resulting in more training examples assigned to the same leaf (see §C.5 for additional details about leaf density in GBRTs).

Revised: 3.2 Modeling the Output Distribution

Original: Once IBUG identifies the k nearest training instancesto x te using eq.
Revised: IBUG has a multitude of choices when modeling the conditional output distribution.

Original: (1), we then compute the output variance of those instances to obtain σ 2ˆ y te . Wecan then use  µ ˆ y te and σ 2ˆ y te to generate a probabilistic prediction for x te , which we show in §3.3.
Revised: The simplest and most common approach is to model the output assuming a Gaussian distribution [22, 64]. We use the scalar output of f : µ ˆ y te = f ( x te ) to model the conditional mean since GBRTs already produce accurate point predictions.

Original: 7: return A 
Revised: Then, we use the k training instances with the largest affinity to x te —we denote this set A ( k ) —to compute the variance σ 2ˆ y te .

Original: 3.2 Calibrating Prediction Variance
Revised: Calibrating prediction variance.

Original: The k -nearest neighbors generally do a good job of determining the relative uncertainty of differentpredictions, but on some datasets, the resulting variance is systematically too large or too small.
Revised: The k -nearest neighbors generally do a good job of determining the relative uncertainty of different predictions, but on some datasets, the resulting variance is systematically too large or too small.

Original: Tocorrect for this, we apply an additional affine transformation before making the prediction:
Revised: To correct for this, we apply an additional affine transformation before making the prediction:

Original: [Equation]
Revised: [Equation]

Original: (2) where γ and δ are tuned on validation data after k has been selected.
Revised: where γ and δ are tuned on validation data after k has been selected.

Original: Instead of exhaustively searchingover all values of γ or δ , we use either the multiplicative factor (tuning γ with δ = 0 ) or the additivefactor (tuning δ with γ = 1 ), and choose between them using their performance on validation data.
Revised: Instead of exhaustively searching over all values of γ or δ , we use either the multiplicative factor (tuning γ with δ = 0 ) or the additive factor (tuning δ with γ = 1 ), and choose between them using their performance on validation data.

Original: We find this simple calibration step consistently improves probabilistic performance for not only
Revised: We find this simple calibration step consistently improves probabilistic performance for not only

Original: IBUG, but for PGBM as well, and at a relatively small cost compared to training the model.
Revised: IBUG, but competing methods as well, and at a relatively small cost compared to training the model.

Original: 3.3 Modeling the Posterior Distribution
Revised: Flexible posterior modeling.

Original: Given a continuous distribution D , we can generate a probabilistic prediction using µ ˆ y te and σ 2ˆ y te asfollows: ˆ D te = D ( µ ˆ y te , σ 2ˆ y te ) . Note D is limited to distributions using location and scale (PGBM canonly model these types of distributions), and is typically chosen to be the normal distribution [20, 62].
Revised: In general, we can generate a probabilistic prediction using µ ˆ y te and σ 2ˆ y te for any distribution that uses location and scale (note PGBM and CBU can only model these types of distributions).

Original: However, IBUG has additional flexibility when modeling the posterior. Once the k training instanceswith the largest affinity to x te is found—we denote this set A ( k ) —we can use those samples todirectly fitD ∗ , any  distribution  including those with high-order moments:
Revised: However, IBUG can additionally use A ( k ) to directly fit any continuous distribution D , including those with high-order moments:

Original: [Equation]
Revised: [Equation]

Original: Eq.
Revised: Eq.

Original: (3) is defined to cover the fact that D∗ can be fit directly with A ( k ) using MLE (maximumlikelihood estimation)
Revised: (3) is defined such that D can be fit directly with A ( k ) using MLE (maximum likelihood estimation)

Original: [45], or may also use µ ˆ y te or σ 2ˆ y te as fixed parameter values and then use A ( k )to fit any other parameters of the distribution.
Revised: [47], or may be fit using µ ˆ y te or σ 2ˆ y te as fixed parameter values with A ( k ) fitting any other parameters of the distribution.

Original: Overall, directly fitting all or some additional parametersin D∗ —for example, the shape parameter in a Weibull distribution—is a benefit over PGBM, whichmust optimize for a suitable shape value using a gridsearch-like approach with extra validation data.
Revised: Overall, directly fitting all or some additional parameters in D —for example, the shape parameter in a Weibull distribution—is a benefit over PGBM and CBU, which can only optimize for a global shape value using a gridsearch-like approach with extra validation data.

Original: Note that NGBoost can model any parameterized distribution, but must specify this choice before training; in contrast, IBUG can optimize this choice after training.
Revised: Note that NGBoost can model any parameterized distribution, but must specify this choice before training; in contrast, IBUG can optimize this choice after training.

Original: Additionally, IBUG maychoose D∗ to be a non-parametric density estimator such as KDE (kernel density estimation) 
Revised: Additionally, IBUG may choose D to be a non-parametric density estimator such as KDE (kernel density estimation) [62], which PGBM,

Original: [60],something neither PGBM or NGBoost can do.
Revised: CBU, and NGBoost cannot do.

Original: In summary, Alg. 2 provides pseudocode for generating a probabilistic prediction with IBUG.
Revised: In summary, Alg. 2 provides pseudocode for generating a probabilistic prediction with IBUG.

Original: Note
Revised: Note

Original: Algs. 1 and 2 work for any GBRT model, allowing practitioners to employ IBUG to adapt multipledifferent point predictors into probabilistic estimators and select the model with the best performance.
Revised: Algs. 1 and 2 work for any GBRT model, allowing practitioners to employ IBUG to adapt multiple different point predictors into probabilistic estimators and select the model with the best performance.

Original: Empirically, we show using different base models for IBUG can result in improved probabilisticperformance than using just one (§5.3).
Revised: Empirically, we show using different base models for IBUG can result in improved probabilistic performance than using just one (§5.3).

Revised: IBUG is a nearest neighbors approach and thus seems well-suited to estimating aleatoric uncertainty— remaining uncertainty due to irreducible error or the inherent stochasticity in the system [33]—since it can quantify the range of outcomes to be expected given the observed features.

Revised: However, we use predictions on held-out data to tune the number of nearest neighbors and the variance calibration hyperparameters; thus, we effectively optimize prediction uncertainty encompassing both aleatoric uncertainty and epistemic uncertainty—error due to the imperfections of the model and the training data [20, 44].

Revised: The evaluation measures in our experiments thus also focus on predictive uncertainty.

Original:  Computational Efficiency
Revised:  Computational Efficiency

Original: Training efficiency.
Revised: Training efficiency.

Original: Since IBUG works with standard GBRT models, it inherits the training efficiency of modern GBRT implementations such as XGBoost [12], LightGBM [34], and CatBoost [50].
Revised: Since IBUG works with standard GBRT models, it inherits the training efficiency of modern GBRT implementations such as XGBoost [12], LightGBM [36], and CatBoost [52].

Original: Prediction efficiency.
Revised: Prediction efficiency.

Original: If there are T trees in the ensemble and each leaf has at most n l traininginstances assigned to it, then IBUG’s prediction time is O ( Tn l ) , since it considers each instancein each leaf.
Revised: If there are T trees in the ensemble and each leaf has at most n l training instances assigned to it, then IBUG’s prediction time is O ( Tn l ) , since it considers each instance in each leaf.

Original: Note training instances that do not appear in a leaf with the test instancex te do notincrease prediction time; what matters most is thus the number of instances at each leaf.
Revised: Note training instances that do not appear in a leaf with the target instance do not increase prediction time; what matters most is thus the number of instances at each leaf.

Original: We find LightGBM often induces regression trees with large leaves—in some cases, over half the dataset isassigned to a single leaf.
Revised: We find LightGBM often induces regression trees with large leaves—in some cases, over half the dataset is assigned to a single leaf (see §C.5 for details).

Original: Thus, prediction time still grows with the size of the dataset, as is typical forinstance-based methods.
Revised: Thus, prediction time still grows with the size of the dataset, as is typical for instance-based methods.

Original: This higher prediction time is the price IBUG pays for greater flexibility.
Revised: This higher prediction time is the price IBUG pays for greater flexibility.

Original: Prediction efficiency can be increased at training time by using deeper GBRTs with fewer instances ineach leaf, after training by subsampling the instances considered for predictions, or at prediction timeby sampling the trees used to compute affinities.
Revised: Prediction efficiency can be increased at training time by using deeper GBRTs with fewer instances in each leaf, after training by subsampling the instances considered for predictions, or at prediction time by sampling the trees used to compute affinities.

Original: We explore this last option in the next subsection.
Revised: We explore this last option in the next subsection.

Original: 4.1 Sampling Trees
Revised: 4.1 Sampling Trees

Original: The most expensive operation when generating a probabilistic predictionfor x te with IBUG iscomputing the affinity vector (eq. 1).
Revised: The most expensive operation when generating a probabilistic prediction with IBUG is computing the affinity vector (Eq. 1).

Original: In order to increase prediction efficiency, we can insteadapproximate the affinity vector by randomly sampling τ < T trees without replacement:
Revised: In order to increase prediction efficiency, we can instead work with a subset of the trees τ < T in the ensemble.

Original: [Equation] (4) in which P ( T ) is the uniform distribution over all T trees.
Revised: We can build this subset by sampling trees uniformly at random, taking the first trees learned (representing the largest gradient steps), or the last trees learned (representing the fine-tuning steps).

Original: The runtime complexity of ˜ A isthus O ( τn l ) , and can provide significant speedups when τ ≪ T .
Revised: By sampling trees, the runtime complexity reduces to O ( τn l ) , which provides significant speedups when τ ≪ T .

Revised: In our empirical evaluation, we find that taking a subset of the first trees learned generally works best, significantly increasing prediction efficiency while maintaining accurate probabilistic predictions (§5.6).

Revised: 4.2 Accelerated k Tuning

Original: Selecting τ < T introducesa tradeoff between computational complexity and predictive accuracy.
Revised: Choosing an appropriate value of k is critical for generating accurate probabilistic predictions in

Original: However, in our empirical evaluation, we find that sampling a relatively small fraction of the total number of trees can significantlyincrease prediction efficiency while maintaining accurate probabilistic predictions (§5.6).
Revised: IBUG. Thus, we aim to tune k using a held-out validation dataset D val ⊂ D and an appropriate probabilistic scoring metric such as negative log likelihood (NLL). Unfortunately, typical tuning procedures would result in the same affinity vectors being computed—an expensive operation—for each candidate value of k .

Original: Experiments
Revised: To mitigate this issue, we perform a custom tuning procedure that reuses computed affinity vectors for all values of k . More specifically, IBUG computes an affinity vector A for a given validation example x val , and then sorts A in descending order (i.e., largest affinity first). Then, IBUG takes the top k training instances, and generates and scores the resulting probabilistic prediction. For each subsequent value of k , the same sorted affinity list can be used, avoiding duplicate computation.

Revised: We summarize this procedure in Alg. 3.

Revised: Once k is chosen, we may encounter a new unseen target instance in which the variance of the khighest affinity training examples for that target example is zero or extremely small.

Revised: In this case, we set the predicted target variance to ρ , which is set during tuning to the minimum (nonzero) variance computed over all predictions in the validation set for the chosen k .

Revised: In practice, we find instances of abnormally low variance to be rare with appropriately chosen values of k .

Revised: do 3:

Revised: A ← C OMPUTE A FFINITIES f

Revised: ▷ Algorithm4:

Revised: A ← Argsort A in descending order 5: for k ∈ K do ▷

Revised: A ( k ) ← Take first k training instances( A k ) 7:

Revised: ˆ D ky j ← P

Revised: P REDICT ( j

Revised: A

Revised: D )

Revised: ▷ Algorithm8:

Revised: S kj ← V

Revised: j

Revised: ▷ Save validation score 9: k ← Select best k from S 10: ρ ← Select minimum σ from

Original:  In this section, we demonstrate IBUG’s ability to produce competitive probabilistic and point predictions as compared to current state-of-the-art methods on a large set of regression datasets (§5.1, §5.2).
Revised: D k 11: return k , ρ  Experiments In this section, we demonstrate IBUG’s ability to produce competitive probabilistic and point predictions as compared to current state-of-the-art methods on a large set of regression datasets (§5.1, §5.2).

Original: Then, we show that IBUG can use different base models to improve probabilistic performance (§5.3),flexibly model the posterior distribution (§5.4), and use approximations to speed up  predictions  (§5.6).
Revised: Then, we show that IBUG can use different base models to improve probabilistic performance (§5.3), flexibly model the posterior distribution (§5.4), and use approximations to speed up probabilistic predictions while maintaining competitive performance (§5.6).

Original: Implementation and Reproducibility.
Revised: Implementation and Reproducibility.

Original: We implement IBUG in Python, using Cython—a Pythonpackage allowing the development of C extensions—to store a unified representation of the modelstructure.
Revised: We implement IBUG in Python, using Cython—a Python package allowing the development of C extensions—to store a unified representation of the model structure.

Original: IBUGcurrently supports all modern gradient boosting frameworks including XGBoost [12], LightGBM [34], and CatBoost [50].
Revised: IBUG supports all modern gradient boosting frameworks including XGBoost [12], LightGBM [36], and CatBoost [52].

Original: Experiments are run on publicly available datasets using an Intel(R)
Revised: Experiments are run on publicly available datasets using an Intel(R)

Original: Xeon(R) CPU E5-2690 v4 @ 2.6GHz with 60GB of RAM @ 2.4GHz.
Revised: Xeon(R) CPU E5-2690 v4 @ 2.6GHz with 60GB of RAM @ 2.4GHz.

Original: Links to all datasources as well as the code for IBUG and all experiments iscurrently available online anonymously.
Revised: Links to all data sources as well as the code for IBUG and all experiments is available at https://github.com/jjbrophy47/ibug .

Original: 5.1 Methodology
Revised: 5.1 Methodology

Original: We now compare IBUG’s probabilistic and point prediction to NGBoost [20] and PGBM [62] on 21benchmark regression datasets and one synthetic dataset.
Revised: We now compare IBUG’s probabilistic and point predictions to NGBoost [22], PGBM [64], and CBU [44] on 21 benchmark regression datasets and one synthetic dataset.

Original: Additional dataset details are in §B.1.
Revised: Additional dataset details are in §B.1.

Original: Metrics.
Revised: Metrics.

Original: We compute the average  negative log likelihood (NLL) and continuous ranked probabilityscore (CRPS) [27, 73] over the test set to evaluate probabilistic performance.
Revised: We compute the average continuous ranked probability score (CRPS ↓ ) and negative log likelihood (NLL ↓ ) [29, 76] over the test set to evaluate probabilistic performance.

Original: To evaluate pointperformance, we use root mean squared error (RMSE).
Revised: To evaluate point performance, we use root mean squared error (RMSE ↓ ).

Original:  Seethe Appendix §B for detailed descriptions.[Table] IBUG
Revised: For all metrics, lower is better. See §B for detailed descriptions.

Original: [Table]

Original: Protocol.
Revised: Protocol.

Original: We follow the same protocol as Sprangers et al.
Revised: We follow a similar protocol to Sprangers et al.

Original: [62] and Duan et al.
Revised: [64] and Duan et al.

Original: First, wecreate 20 random folds for each dataset.For each fold, we randomly sample 10% of the data touse as a test set.
Revised: We use 10-fold cross-validation to create 10 90/10 train/test folds for each dataset.

Original: The remaining 90%  is randomly split into an 80/20 train/validation set to tune anyhyperparameters.
Revised: For each fold, the 90% training set is randomly split into an 80/20 train/validation set to tune any hyperparameters.

Original: Once the hyperparameters are tuned, the model is retrained using the entire 90%training set.
Revised: Once the hyperparameters are tuned, the model is retrained using the entire 90% training set.

Original: For probabilistic predictions, a normal distribution is used to model the output.
Revised: For probabilistic predictions, a normal distribution is used to model the output.

Original: Significance testing.
Revised: Significance Testing.

Original: To determine which of two methods performs better on a given dataset undera given metric, we use a two-tailed paired t-test over the 20 random folds with a significance levelof p < 0 . Wealso report counts of the number of datasets in which a given method performedsignificantly better (“Win”), worse (“Loss”), or not  different (“Tie”) relative to a comparator.
Revised: We report counts of the number of datasets in which a given method performed better (“Win”), worse (“Loss”), or not statistically different (“Tie”) relative to a comparator using a two-sided paired t-test over the 10 random folds with a significance level of 0.05.

Original: Hyperparameters.
Revised: Hyperparameters.

Original: We tune NGBoost the same way as in Duan et al.
Revised: We tune NGBoost the same way as in Duan et al.

Original: Since PGBM and IBUGboth optimize a point prediction metric, we tune their hyperparameters similarly.
Revised: Since PGBM, CBU, and IBUG optimize a point prediction metric, we tune their hyperparameters similarly.

Original: We also tunevariance calibration parameters γ and δ for each method (§3.2).
Revised: We also tune variance calibration parameters γ and δ for each method (§3.2).

Original: Exact  values evaluated and selectedare in §B.2.
Revised: Exact hyperparameter values evaluated and selected are in §B.2.

Original: Unless specified otherwise, we use LightGBM [34] as the base model for IBUG. 
Revised: Unless specified otherwise, we use CatBoost [52] as the base model for IBUG. [Table]

Original: 5.2 Probabilistic and Point Prediction
Revised: 5.2 Probabilistic and Point Performance

Original: We first compare IBUG’s probabilistic and point predictions to each baseline on each dataset; see
Revised: We first compare IBUG’s probabilistic and point predictions to each baseline on each dataset. See

Original: Table 1 for detailed results.
Revised: Table 1 for detailed CRPS results; due to space constraints, results for additional probabilistic metrics (e [Table] g., NLL) as well as point performance results are in §B.3.

Original: Our main findings are as follows:
Revised: Our main findings are as follows:

Original: • On probabilistic performance, IBUG performs equally well or better than NGBoostandPGBM. For NLL, IBUG wins against NGBoost and PGBM on 17 and 13 (out of 22) datasets [Table] [Figure]
Revised: • On probabilistic performance, IBUG performs equally well or better than NGBoost and PGBM, winning on 17 and 11 (out of 22) datasets respectively, while losing on only 2 and 2 (respectively).

Original: [Figure]

Original: [Figure] [Figure] k [Figure]
Revised: Since CBU and IBUG performance is similar, we combine the two approaches, averaging their outputs; we denote this simple ensemble IBUG+CBU .

Original: [Figure] [Figure] respectively, while losing on only  2 and 3. For CRPS, IBUGwins on 18 and 9 datasets,respectively, while again losing on only 2 and 3.
Revised: Surprisingly, IBUG+CBU works very well, losing on only a maximum of 2 datasets when faced head-to-head against any other method; these results suggest IBUG and CBU are complimentary approaches.

Original: • On point performance, IBUG and PGBM performed significantly better than NGBoost; thisis consistent with previous work and is perhaps unsurprising since NGBoost is optimized 
Revised: • On point performance, PGBM, CBU, and IBUG performed significantly better than NGBoost; this is consistent with previous work and is perhaps unsurprising since NGBoost is optimized for probabilistic performance, not point performance.

Original: for probabilistic performance, not point performance.
Revised: However, IBUG generally performed better than PGBM, winning on 13 datasets and losing on only 1 dataset; and performed slightly better than CBU, winning on 6 datasets with no losses. We also compare IBUG with two additional baselines— k NN and BART [13]—shown in §C.2–C. due to space constraints. We find IBUG generally outperforms these methods in both probabilistic and point performance.

Original: However, IBUG generally performed better than PGBM, winning on 10 datasets and losing on only 1 dataset.
Revised: Overall, the results in this section suggest IBUG generates both competitive probabilistic and point predictions compared to existing methods. [Table] [Figure]

Original:  5.3 Different Base Models
Revised: [Figure] 5.3 Different Base Models

Original: Here we experiment using different base models for IBUG besides LightGBM [34]; specifically, weuse XGBoost [12] and CatBoost [50], two popular gradient boosting frameworks.
Revised: Here we experiment using different base models for IBUG besides CatBoost [52]; specifically, we use LightGBM [36] and XGBoost [12], two popular gradient boosting frameworks.

Original: Table 2 showsthat using a different base model can result in improved probabilistic performance.
Revised: Table 2 shows that using a different base model can result in improved probabilistic performance.

Original: This highlights
Revised: This highlights

Original: IBUG’s agnosticism to GBRT types, enabling practitioners tobe apply IBUG to future models withimproved point prediction performance.
Revised: IBUG’s agnosticism to GBRT type, enabling practitioners to apply IBUG to future models with improved point prediction performance.

Original: We show IBUG can even improve probabilistic performancefor PGBM by using PGBM as the base model for IBUG (see §C.1 for details).

Original: 5.4 Posterior Modeling
Revised: 5.4 Posterior Modeling

Original: One of the unique benefits of IBUG is the ability to directly model the output using empiricalsamples (Figure 1), giving practitioners a better sense of the output distribution for specific predictions.
Revised: One of the unique benefits of IBUG is the ability to directly model the output using empirical samples (Figure 1), giving practitioners a better sense of the output distribution for specific predictions.

Original: IBUG can optimize a distribution after training, and has more flexibility in the types of distributionsit can model—from distributions using just location and scale to those with high-order moments In most cases, calibration maintains or improves performance; it is especially helpful for PGBM. [Table] [Figure]
Revised: IBUG can optimize a distribution after training, and has more flexibility in the types of distributions it can model—from distributions using just location and scale to those with high-order moments as well as non-parametric density estimators.

Original: [Table] [Figure] as well as non-parametric density estimators. To test this flexibility, we model each probabilisticprediction using the following distributions: normal, skewnormal, lognormal, Laplace, student t,logistic, Gumbel, Weibull, and KDE; we then select the distribution with the best average NLL onthe validation set, and evaluate its probabilistic performance on the test set.
Revised: To test this flexibility, we model each probabilistic prediction using the following distributions: normal, skewnormal, lognormal, Laplace, student t, logistic, Gumbel, Weibull, and KDE; we then select the distribution with the best average NLL on the validation set, and evaluate its probabilistic performance on the test set.

Original: Figure 1 demonstrates  the selected distributions for the MEPS and Wine datasets achieve betterprobabilistic performance than assuming a normal distribution.
Revised: Figure 1 demonstrates that the selected distributions for the MEPS and Wine datasets achieve better probabilistic performance than assuming normality.

Original: Qualitatively, the empirical densitiesof A ( k ) for a randomly sampled set of test instances reaffirms the selected distributions. {+As an additional comparison, we report CBU achieves a test NLL of 3 . 699 ± 0 . 038 and 1 . 025 ± 0 . 
+}
Revised: Qualitatively, the empirical densities of A ( k ) for a randomly sampled set of test instances reaffirms the selected distributions. {+As an additional comparison, we report CBU achieves a test NLL of 3 . 699 ± 0 . 038 and 1 . 025 ± 0 . 
+}

Revised: for the MEPS and wine datasets (respectively) using a normal distribution, while IBUG achieves − 6 . 887 ± 0 .

Revised: 260 and 0 .

Original: Additionally,on the MEPS dataset,we note the selected Weibull distribution takes a shape parameter, which IBUGestimates directly on a per prediction basis using A ( k ) and MLE.
Revised: 785 ± 0 . 025 using Weibull and KDE estimation (respectively). For the MEPS dataset, the selected Weibull distribution takes a shape parameter, which IBUG estimates directly on a per prediction basis using A ( k ) and MLE.

Original: In contrast, PGBM  would have tooptimize for a single value over a validation set, which may be suboptimal for individual predictions. 
Revised: In contrast, PGBM or CBU would need to optimize a global shape value using a validation set, which is likely to be suboptimal for individual predictions. In all cases, calibration maintains or improves performance; it is especially helpful for CBU [Table] T

Original: 5.5 Variance Calibration
Revised: [Table] l T [Table]

Revised: n

Revised: T m

Original:  Table 3 shows probabilistic performance comparisons of each method against itself with and withoutvariance calibration.
Revised: [Table] [Table] 5.5 Variance Calibration Table 3 shows probabilistic performance comparisons of each method against itself with and without variance calibration.

Original: In the majority of cases, variance calibration (§3.2) either maintains or improvesperformance for all methods, especially PGBM.
Revised: In all cases, variance calibration (§3.2) either maintains or improves performance for all methods, especially CBU.

Original: Overall, these results suggest that variance calibrationshould be a standard procedure for probabilistic prediction, unless using a method that has particularlywell-calibrated predictions to begin with.
Revised: Overall, these results suggest that variance calibration should be a standard procedure for probabilistic prediction, unless using a method that has particularly wellcalibrated predictions to begin with.

Original: We therefore use variance calibration in all of our results.
Revised: We therefore use variance calibration in all of our results.

Original: Additionally, §C.2 shows performance results for all methods without variance calibration.
Revised: Additionally, §C.1 shows performance results for all methods without variance calibration.

Original: Overall,we observe similar  trends as when applying calibration (Table 1); the W-T-L ratio for IBUG vs.
Revised: Overall, we observe similar relative performance trends as when applying calibration (Table 1). [Figure] [Table] [Figure]

Revised: [Table]

Revised: [Figure]

Original: NGBoost and PGBM is 12-5-5 and 15-6-1 (respectively) for NLL and 18-2-2 and 14-4-4 for CRPS.
Revised: [Table] [Figure] [Table] [Figure]

Original: 5.6 Sampling Trees
Revised: 5.6 Sampling Trees

Original: Table 2 shows the runtime for each method broken down into total training time (including tuning)and prediction time per test example.
Revised: Figure 2 shows the runtime for each method broken down into total training time (including tuning) and prediction time per test example.

Original: On average, IBUG has similar training times to NGBoost, but isorders of magnitude faster than PGBM.
Revised: On average, IBUG has similar training times to PGBM and CBU, but on some datasets, IBUG is roughly an order of magnitude faster than PGBM.

Original: For predictions, IBUG is significantly slower than NGBoost. [Figure] i
Revised: For predictions, IBUG is similar to PGBM but relatively slow compared to NGBoost and CBU.

Original: [Figure] i [Figure] However, by sampling τ < T trees when computing the affinity vector, IBUG can significantlyreduce prediction time (Figure 3).
Revised: However, by sampling τ < T trees when computing the affinity vector, IBUG can significantly reduce prediction time.

Revised: Figure 3 shows results when sampling trees first-to-last, which typically works best over all tree-sampling strategies (alternate sampling strategies are evaluated in §C.4).

Original: We observe as τ decreases,  average prediction time decreasessignificantly while probabilistic performance remains relatively stable until τ/T reaches roughly5–20%, at which point probabilistic performance  starts to decrease more rapidly.
Revised: As τ decreases, we observe average prediction time decreases roughly 1-2 orders of magnitude while probabilistic performance remains relatively stable until τ/T reaches roughly 1–5%, at which point probabilistic performance sometimes starts to decrease more rapidly.

Original: Note for the Bike, California, Obesity, and Protein datasets, IBUG can reach the same average prediction time as PGBMwhile maintaining  better probabilistic performance than both NGBoost and PGBM.
Revised: Note for the Ames and Life datasets, IBUG can reach the same average prediction time as NGBoost while maintaining the same or better probabilistic performance than NGBoost, PGBM, and CBU.

Original: These resultsdemonstrate that if speed is a concern, IBUG can approximate the affinity computation to speed upprediction times while maintaining competitive probabilistic performance.
Revised: These results demonstrate that if speed is a concern, IBUG can approximate the affinity computation to speed up prediction times while maintaining competitive probabilistic performance.

Original:   Related Work
Revised:  Additonal Related Work

Original: More traditional approaches to probabilistic regression include generalized additive models forlocation, scale, and shape (GAMLSS), which allow for a flexible choice of distribution for the targetvariable but are restricted toa pre-specified model form [54].
Revised: Traditional approaches to probabilistic regression include generalized additive models for location, scale, and shape (GAMLSS), which allow for a flexible choice of distribution for the target variable but are restricted to pre-specified model forms [56].

Original: Prophet [66] also produces probabilisticestimates for generalized additive models, but has been shown to underperform as compared to morerecent approaches [58, 3].
Revised: Prophet [68] also produces probabilistic estimates for generalized additive models, but has been shown to underperform as compared to more recent approaches [60, 3].

Original: Bayesian methods [46, 28] naturally generate uncertainty estimates byintegrating over the posterior; but, exact solutions are limited to simple models, and more complexmodels such as Bayesian Additive Regression Trees (BART) [13, 39] require computationallyexpensive sampling techniques (e.g., MCMC [4]) to provide approximate solutions.
Revised: Bayesian methods [48, 30] naturally generate uncertainty estimates by integrating over the posterior; however, exact solutions are limited to simple models, and more complex models such as Bayesian Additive Regression Trees (BART) [13, 41] require computationally expensive sampling techniques (e.g., MCMC [4]) to provide approximate solutions.

Original: Other approaches to probabilistic regression tasks include conformal predictions [59, 65, 5] whichproduce confidence intervals via empirical errors obtained in the past, and quantile regression [30,35, 43, 55].
Revised: Other approaches to probabilistic regression tasks include conformal predictions [61, 67, 5] which produce confidence intervals via empirical errors obtained in the past, and quantile regression [32, 37, 45, 57].

Original: Similar to PGBM, distributional forests (DFs) [57] estimate distributional parametersin each leaf, and average these estimates over all trees in the forest.
Revised: Similar to PGBM, distributional forests (DFs) [59] estimate distributional parameters in each leaf, and average these estimates over all trees in the forest.

Original: Deep learning approachesfor probabilistic regression [52, 71, 3] have increased recently, with notable approaches such as
Revised: Deep learning approaches for probabilistic regression [54, 74, 3] have increased recently, with notable approaches such as

Original: DeepAR [56] and methods based on transformer architectures [36, 37].
Revised: DeepAR [58] and methods based on transformer architectures [38, 39].

Original:  Conclusion
Revised:  Conclusion

Original: IBUG uses ideas from instance-based learning to enable probabilistic predictions for any GBRT pointpredictor.
Revised: IBUG uses ideas from instance-based learning to enable probabilistic predictions for any GBRT point predictor.

Original: IBUG generates probabilistic predictions by using the k -nearest training instances to thetest instance found using the structure of the trees in the ensemble.
Revised: IBUG generates probabilistic predictions by using the k -nearest training instances to the test instance found using the structure of the trees in the ensemble.

Original: Our results on 22 regressiondatasets demonstrate this simple wrapper produces competitive probabilistic and point predictions tocurrent state-of-the-art methods, most notably NGBoost [20] and PGBM [62].
Revised: Our results on 22 regression datasets demonstrate this simple wrapper produces competitive probabilistic and point predictions to current state-of-the-art methods, most notably NGBoost [22], PGBM [64], and CBU [44].

Original: We also show that IBUG can more flexibly model the posterior distribution of a prediction using any parametric ornon-parametric density estimator.
Revised: We also show that IBUG can more flexibly model the posterior distribution of a prediction using any parametric or non-parametric density estimator.

Original: IBUG is also fast to train compared to competing methods.

Original: IBUG’sone limitation is relatively slow prediction time.
Revised: IBUG’s one limitation is relatively slow prediction time.

Original: However, we show that approximations in the searchfor the k -nearest training instances can significantly speed up prediction time.
Revised: However, we show that approximations in the search for the k -nearest training instances can significantly speed up prediction time; predictions are also easily parallelizable in IBUG.

Original: Checklist
Revised: For future work, we plan to investigate other approximations to the affinity vector such as subsampling or even reweighting training instances, which may lead to significant speed ups of IBUG. Acknowledgments and Disclosure of Funding We would like to thank Zayd Hammoudeh for useful discussions and feedback and the reviewers for their constructive comments that improved this paper. This work was supported by a grant from the

Original: For all authors...
Revised: Air Force Research Laboratory and the Defense Advanced Research Projects Agency (DARPA)— agreement number FA8750-16-C-0166, subcontract K001892-00-S05, as well as a second grant from DARPA, agreement number HR00112090135. This work benefited from access to the University of Oregon high-performance computer, Talapas.

Original:  (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?

Original: [Yes] (b) Did you describe the limitations of your work?
Revised: [Yes] (b) Did you describe the limitations of your work?

Original: [Yes] See §7.
Revised: [Yes] See §7.

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: (c) Did you discuss any potential negative societal impacts of your work?

Original: [Yes] See §A.1.
Revised: [Yes] See §A.1.

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Original: [Yes] 
Revised: [Yes] 2.

Original: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
Revised: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?

Original: [N/A] (b) Did you include complete proofs of all theoretical results?
Revised: [N/A] (b) Did you include complete proofs of all theoretical results?

Original: [N/A] 
Revised: [N/A] 3.

Original: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)?
Revised: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Original: [Yes] See §5.
Revised: [Yes] See §5.

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

Original: [Yes] See §5.1.
Revised: [Yes] See §5.1.

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
Revised: [Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?

Original: [Yes] See §5. 
Revised: [Yes] See §5. 4.

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?

Original: [Yes] See §B.1.
Revised: [Yes] See §B.1.

Original: (b) Did you mention the license of the assets?
Revised: (b) Did you mention the license of the assets?

Original: [No] The license of each asset can be found by following its corresponding citation.
Revised: [No] The license of each asset can be found by following its corresponding citation.

Original: (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [Yes] See §5.
Revised: [Yes] See §5.

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?

Original: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable[Table] information or offensive content?
Revised: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?

Original: [N/A] [Table]
Revised: [N/A] 5.

Original: If you used crowdsourcing or conducted research with human subjects...[Table] (a) Did you include the full text of instructions given to participants and screenshots, if [Table]
Revised: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Original: applicable? [N/A][Table] (b) Did you describe any potential participant risks, with links to Institutional Review 
Revised: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

Original: [Table]

Original: Board (IRB) approvals, if applicable?

Original: [N/A][Table] (c) Did you include the estimated hourly wage paid to participants and the total amount [Table]
Revised: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

Original: spent on participant compensation?

Original: [N/A]
Revised: [N/A]

Original: [Table]