Original: T AKING A S TEP B ACK WITH KC AL : M ULTI -C LASS K ERNEL -B ASED C ALIBRATION FOR D EEP N EURAL N ETWORKS
Revised: T AKING A S TEP B ACK WITH KC AL : M ULTI -C LASS K ERNEL -B ASED C ALIBRATION FOR D EEP N EURAL N ETWORKS

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Deep neural network (DNN) classifiers are often overconfident, producing miscalibrated class probabilities.
Revised: Deep neural network (DNN) classifiers are often overconfident, producing miscalibrated class probabilities.

Original: In high-risk applications like healthcare, practitioners require fully calibrated probability predictions for decision-making.
Revised: In high-risk applications like healthcare, practitioners require fully calibrated probability predictions for decision-making.

Original: That is, conditioned on the prediction vector , every class’ probability should be close to the predicted value.
Revised: That is, conditioned on the prediction vector , every class’ probability should be close to the predicted value.

Original: Most existing calibration methods either lack theoretical guarantees for producing calibrated outputs, reduce classification accuracy in the process, or only calibrate the predicted class.
Revised: Most existing calibration methods either lack theoretical guarantees for producing calibrated outputs, reduce classification accuracy in the process, or only calibrate the predicted class.

Original: This paper proposes a new Kernel-based calibration method called KCal.
Revised: This paper proposes a new Kernel-based calibration method called KCal.

Original: Unlike existing calibration procedures, KCal does not operate directly on the logits or softmax outputs of the DNN.
Revised: Unlike existing calibration procedures, KCal does not operate directly on the logits or softmax outputs of the DNN.

Original: Instead, KCal learns a metric space on the penultimate-layer latent embedding and generates predictions using kernel density estimates on a calibration set.
Revised: Instead, KCal learns a metric space on the penultimate-layer latent embedding and generates predictions using kernel density estimates on a calibration set.

Original: We first analyze KCal theoretically, showing that it enjoys a provable full calibration guarantee.
Revised: We first analyze KCal theoretically, showing that it enjoys a provable full calibration guarantee.

Original: Then, through extensive experiments across a variety of datasets, we show that KCal consistently outperforms baselines as measured by the calibration error and by proper scoring rules like the Brier Score.
Revised: Then, through extensive experiments across a variety of datasets, we show that KCal consistently outperforms baselines as measured by the calibration error and by proper scoring rules like the Brier Score.

Revised: Our code is available at https://github.com/zlin7/KCal .

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: The notable successes of Deep Neural Networks (DNNs) in complex classification tasks, such as object detection (Ouyang & Wang, 2013), speech recognition (Deng et al., 2013), and medical diagnosis (Qiao et al., 2020; Biswal et al., 2017), have made them essential ingredients within various critical decision-making pipelines.
Revised: The notable successes of Deep Neural Networks (DNNs) in complex classification tasks, such as object detection (Ouyang & Wang, 2013), speech recognition (Deng et al., 2013), and medical diagnosis (Qiao et al., 2020; Biswal et al., 2017), have made them essential ingredients within various critical decision-making pipelines.

Original: In addition to the classification accuracy, a classifier should ideally also generate reliable uncertainty estimates represented in the predicted probability vector.
Revised: In addition to the classification accuracy, a classifier should ideally also generate reliable uncertainty estimates represented in the predicted probability vector.

Original: An influential study (Guo et al., 2017) reported that modern DNNs are often overconfident or miscalibrated , which could lead to severe consequences in high-stakes applications such as healthcare (Jiang et al., 2012).
Revised: An influential study (Guo et al., 2017) reported that modern DNNs are often overconfident or miscalibrated , which could lead to severe consequences in high-stakes applications such as healthcare (Jiang et al., 2012).

Original: Calibration is the process of closing the gap between the prediction and the ground truth distribution given this prediction.
Revised: Calibration is the process of closing the gap between the prediction and the ground truth distribution given this prediction.

Original: For a K -class classification problem, with covariates X ∈ X and the label Y ∈ Y =
Revised: For a K -classification problem, with covariates X ∈ X and the label Y ∈ Y =

Original: [ K ] , denote our classifier X (cid:55)→ ∆ K − 1 as ˆ p =
Revised: [ K ] , denote our classifier X (cid:55)→ ∆ K − 1 as ˆ p =

Original: [ˆ p 1 , . . .
Revised: [ˆ p 1 , . . .

Original: , ˆ p K ] , where ∆ K − 1 is the ( K − 1 )- simplex.
Revised: , ˆ p K ] , with ∆ K − 1 being ( K -1)-simplex.

Original: Then,
Revised: Then,

Original: Definition 1.
Revised: Definition 1.

Original: (Full Calibration (Kull et al., 2019; Vaicenavicius et al., 2019))
Revised: (Full Calibration (Vaicenavicius et al., 2019))

Original: ˆ p is fully-calibrated if ∀ k ∈
Revised: ˆ p is fully-calibrated if ∀ k ∈

Original: [ K ] :
Revised: [ K ] :

Original: [Equation]
Revised: [Equation]

Original: It is worth noting that Def.
Revised: It is worth noting that Def.

Original: (1) implies nothing about accuracy .
Revised: (1) implies nothing about accuracy .

Original: In fact, ignoring X and simply predicting π , the class frequency vector, results in a fully calibrated but inaccurate classifier.
Revised: In fact, ignoring X and simply predicting π , the class frequency vector, results in a fully calibrated but inaccurate classifier.

Original: As a result, our goal is always to improve calibration while maintaining accuracy .
Revised: As a result, our goal is always to improve calibration while maintaining accuracy .

Original: Another important requirement is that ˆ p ∈ ∆ K − 1 .
Revised: Another important requirement is that ˆ p ∈ ∆ K − 1 .

Original: Many binary calibration methods such as Zadrozny & Elkan (2001; 2002) result in vectors that are not interpretable as probabilities, and have to be normalized.
Revised: Many binary calibration methods such as Zadrozny & Elkan (2001; 2002) result in vectors that are not interpretable as probabilities, and have to be normalized.

Original: Many existing works only consider confidence calibration (Guo et al., 2017; Zhang et al., 2020;
Revised: Many existing works only consider confidence calibration (Guo et al., 2017; Zhang et al., 2020;

Original: Wenger et al., 2020; Ma & Blaschko, 2021), a much weaker notion than that encapsulated by Def.
Revised: Wenger et al., 2020; Ma & Blaschko, 2021), a much weaker notion than that encapsulated by Def.

Original: (1) and only calibrates the predicted class (Kull et al., 2019; Vaicenavicius et al., 2019).
Revised: (1) and only calibrates the predicted class (Kull et al., 2019; Vaicenavicius et al., 2019).

Original: Definition 2.
Revised: Definition 2.

Original: (Confidence Calibration)
Revised: (Confidence Calibration)

Original: ˆ p is confidence-calibrated if:
Revised: ˆ p is confidence-calibrated if:

Original: [Equation]
Revised: [Equation]

Original: III C : C o n f i d e n c e Uncalibrated 0.0 0.2 0.4 0.6 0.8 1.0 0. 1.0 Temperature Scaling 0.0 0.2 0.4 0.6 0.8 1.0 prediction 0.
Revised: [Equation]

Original: III C : S e i z u
Revised: [Figure]

Original: r e 0.0 0.2 0.4 0.6 0.8 1.0 prediction 0.
Revised: [Figure] [Figure]

Original: GapActual Frequency and Seizure (bottom).
Revised: [Figure] n [Figure] and Seizure (bottom).

Original: The popular temperaturescaling (right) only calibrates the confidence, leaving Seizure poorly calibrated.
Revised: The pop- temperature only calibrates the confidence, leaving Seizure poorly calibrated.

Original: See Figure 2 and the Appendix for complete reliability diagrams.
Revised: See Figure 2 and the Appendix for complete reliability diagrams.

Original: However, confidence calibration is far from sufficient.
Revised: However, confidence calibration is far from sufficient.

Original: Doctors need to perform differential diagnoses on a patient, where multiple possible diseases should be considered with proper probabilities for all of them, not only the most likely diagnosis.
Revised: Doctors need to perform differential diagnoses on a patient, where multiple possible diseases should be considered with proper probabilities for all of them, not only the most likely diagnosis.

Original: Figureshows an example where the confidence is calibrated, but prediction for important classes like Seizure is poorly calibrated.
Revised: Figureshows an example where the confidence is calibrated, but prediction for important classes like Seizure is poorly calibrated.

Original: A classifier can be confidence-calibrated but not useful for differential diagnoses if the probability assignments for most diseases are inaccurate.
Revised: A classifier can be confidence-calibrated but not useful for such tasks if the probabilities assigned to most diseases are inaccurate.

Original: Recent research effort has started to focus on full calibration, for example, in Vaicenavicius et al.
Revised: Recent research effort has started to focus on full calibration, for example, in Vaicenavicius et al.

Original: Kull et al.
Revised: Kull et al.

Original: ; Widmann et al.
Revised: ; Widmann et al.

Original: ; Karandikar et al.

Original: ; Mukhoti et al.

Original: Patel et al.

Original: We approach this problem by leveraging the latent neural network embedding in a nonparametric manner.
Revised: We approach this problem by leveraging the latent neural network embedding in a nonparametric manner.

Original: Nonparametric methods such as histogram binning (HB) (Zadrozny & Elkan, 2001) and isotonic regression (IR) (Zadrozny & Elkan, 2002), are natural for calibration and have become popular.
Revised: Nonparametric methods such as histogram binning (HB) (Zadrozny & Elkan, 2001) and isotonic regression (IR) (Zadrozny & Elkan, 2002), are natural for calibration and have become popular.

Original: Gupta & Ramdas (2021) recently showed a calibration guarantee for HB.
Revised: Gupta & Ramdas (2021) recently showed a calibration guarantee for HB.

Original: However, HB usually leads to noticeable drops in accuracy (Patel et al., 2021), and IR is prone to overfitting (Niculescu-Mizil &
Revised: However, HB usually leads to noticeable drops in accuracy (Patel et al., 2021), and IR is prone to overfitting (Niculescu-Mizil &

Original: Caruana, 2005).
Revised: Caruana, 2005).

Original: Unlike existing methods, we take one step back and train a new low-dimensional metric space on the penultimate-layer embeddings of DNNs.
Revised: Unlike existing methods, we take one step back and train a new low-dimensional metric space on the penultimatelayer embeddings of DNNs.

Original: Then, we use a kernel density estimationbased classifier to predict the class probabilities directly.
Revised: Then, we use a kernel density estimation-based classifier to predict the class probabilities directly.

Original: We refer to our K ernel-based Cal ibration method as KCal.
Revised: We refer to our K ernel-based Cal ibration method as KCal.

Original: Unlike most calibration methods, KCal provides high probability error bounds for full calibration under standard assumptions.
Revised: Unlike most calibration methods, KCal provides high probability error bounds for full calibration under standard assumptions.

Original: Empirically, we show that with little overhead, KCal outperforms all existing calibration methods in terms of calibration quality, across multiple tasks and DNN architectures, while maintaining and sometimes improving the classification accuracy.
Revised: Empirically, we show that with little overhead, KCal outperforms all existing calibration methods in terms of calibration quality, across multiple tasks and DNN architectures, while maintaining and sometimes improving the classification accuracy.

Original: Summary of Contributions :
Revised: Summary of Contributions :

Original: • We propose KCal, a principled method that calibrates DNNs using kernel density estimation on the latent embeddings .
Revised: • We propose KCal, a principled method that calibrates DNNs using kernel density estimation on the latent embeddings .

Original: • We present an efficient pipeline to train KCal, including a dimension-reducing projection and a stratified sampling method to facilitate efficient training.
Revised: • We present an efficient pipeline to train KCal, including a dimension-reducing projection and a stratified sampling method to facilitate efficient training.

Original: • We provide finite sample bounds for the calibration error of KCal-calibrated output under standard assumptions.
Revised: • We provide finite sample bounds for the calibration error of KCal-calibrated output under standard assumptions.

Original: To the best of our knowledge, this is the first method with a full calibration guarantee.
Revised: To the best of our knowledge, this is the first method with a full calibration guarantee, especially for neural networks.

Original: • In extensive experiments on multiple datasets and state-of-the-art models, we found that KCal outperforms existing calibration methods in commonly used evaluation metrics.
Revised: • In extensive experiments on multiple datasets and state-of-the-art models, we found that KCal outperforms existing calibration methods in commonly used evaluation metrics.

Original: We also show that KCal provides more reliable predictions for important classes in the healthcare datasets.
Revised: We also show that KCal provides more reliable predictions for important classes in the healthcare datasets.

Original: The code to replicate all our experimental results is submitted along with supplementary materials.
Revised: The code to replicate all our experimental results is submitted along with supplementary materials.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Research on calibration originated in the context of meteorology and weather forecasting (see Murphy & Winkler (1984) for an overview) and has a long history, much older than the field of machine learning (Brier, 1950; Murphy & Winkler, 1977; Degroot & Fienberg, 1983).
Revised: Research on calibration originated in the context of meteorology and weather forecasting (see Murphy & Winkler (1984) for an overview) and has a long history, much older than the field of machine learning (Brier, 1950; Murphy & Winkler, 1977; Degroot & Fienberg, 1983).

Original: We refer to Filho et al.
Revised: We refer to Filho et al.

Original: (2021) for a holistic overview and focus below on methods proposed in the context of modern neural networks.
Revised: (2021) for a holistic overview and focus below on methods proposed in the context of modern neural networks.

Original: Based on underlying methodological similarities, we cluster them into distinct categories.
Revised: Based on underlying methodological similarities, we cluster them into distinct categories.

Original: Scaling: A popular family of calibration methods is based on scaling, in which a mapping is learned from the predicted logits to probability vectors.
Revised: Scaling: A popular family of calibration methods is based on scaling, in which a mapping is learned from the predicted logits to probability vectors.

Original: Confidence calibration scaling methods include temperature scaling ( TS ) (Guo et al., 2017) and its antecedent Platt scaling (Platt, 1999), an ensemble of TS (Zhang et al., 2020), Gaussian-Process scaling (Wenger et al., 2020), combining a base calibrator ( TS ) with a rejection option (Ma & Blaschko, 2021).
Revised: Confidence calibration scaling methods include temperature scaling ( TS ) (Guo et al., 2017) and its antecedent Platt scaling (Platt, 1999), an ensemble of TS (Zhang et al., 2020), Gaussian-Process scaling (Wenger et al., 2020), combining a base calibrator ( TS ) with a rejection option (Ma & Blaschko, 2021).

Original: Matrix scaling with regularization was also used to perform full calibration (Kull et al., 2019).
Revised: Matrix scaling with regularization was also used to perform full calibration (Kull et al., 2019).

Original: While some scaling-based methods can be data-efficient, there are no known theoretical guarantees for them to the best of our knowledge.
Revised: While some scaling-based methods can be data-efficient, there are no known theoretical guarantees for them to the best of our knowledge.

Original: Binning: Another cluster of solutions relies on binning and its variants, and includes uniformmass binning (Zadrozny & Elkan, 2001), scaling before binning (Kumar et al., 2019), and mutual- information-maximization-based binning (Patel et al., 2021).
Revised: Binning: Another cluster of solutions relies on binning and its variants, and includes uniformmass binning (Zadrozny & Elkan, 2001), scaling before binning (Kumar et al., 2019), and mutualinformation-maximization-based binning (Patel et al., 2021).

Original: Isotonic regression (Zadrozny & Elkan, 2002) is also often interpreted as binning.
Revised: Isotonic regression (Zadrozny & Elkan, 2002) is also often interpreted as binning.

Original: Uniform-mass binning (Zadrozny & Elkan, 2001) has a distribution-free finite sample calibration guarantee (Gupta & Ramdas, 2021) and asymptotic convergent ECE estimation (Vaicenavicius et al., 2019).
Revised: Uniform-mass binning (Zadrozny & Elkan, 2001) has a distribution-free finite sample calibration guarantee (Gupta & Ramdas, 2021) and asymptotic convergent ECE estimation (Vaicenavicius et al., 2019).

Original: However, in practice, binning tends to decrease accuracy (Patel et al., 2021; Guo et al., 2017).
Revised: However, in practice, binning tends to decrease accuracy (Patel et al., 2021; Guo et al., 2017).

Original: Binning can also be considered a member of the broader nonparametric calibration family of methods.
Revised: Binning can also be considered a member of the broader nonparametric calibration family of methods.

Original: Such methods also include Gaussian
Revised: Such methods also include Gaussian

Original: Process Calibration (Wenger et al., 2020), which however also only considers confidence calibration.
Revised: Process Calibration (Wenger et al., 2020), which however also only considers confidence calibration.

Original: Loss regularization: There are also attempts to train a calibrated DNN to begin with.
Revised: Loss regularization: There are also attempts to train a calibrated DNN to begin with.

Original: Such methods typically add a suitable regularizer to the loss function (Karandikar et al., 2021;
Revised: Such methods typically add a suitable regularizer to the loss function (Karandikar et al., 2021;

Original: Mukhoti et al., 2020;
Revised: Mukhoti et al., 2020;

Original: Kumar et al., 2018), which can sometimes result in expensive optimization and reduction in accuracy.
Revised: Kumar et al., 2018), which can sometimes result in expensive optimization and reduction in accuracy.

Original: Use of Kernels: Although not directly used for calibration, kernels have also been used for uncertainty quantification for deep learning classification.
Revised: Use of Kernels: Although not directly used for calibration, kernels have also been used for uncertainty quantification for deep learning classification.

Original: In classification with rejection, the k-nearest-neighbors algorithm (kNN), closely related to kernel-based methods, has been used to provide a “confidence measure” which is used to make a binary decision (i.e., whether to reject or to predict) (Papernot &
Revised: In classification with rejection, the k-nearest-neighbors algorithm (kNN), closely related to kernel-based methods, has been used to provide a “confidence measure” which is used to make a binary decision (i.e., whether to reject or to predict) (Papernot &

Original: McDaniel, 2018; Jiang et al., 2018).
Revised: McDaniel, 2018; Jiang et al., 2018).

Original: Recently, continuous kernels have also been used to measure calibration quality or used as regularization during training (Widmann et al., 2019; Kumar et al., 2018).
Revised: Recently, continuous kernels have also been used to measure calibration quality or used as regularization during training (Widmann et al., 2019; Kumar et al., 2018).

Original: Zhang et al.
Revised: Zhang et al.

Original: (2020) introduced a kernel density estimation (KDE) proxy estimator for estimating
Revised: (2020) introduced a kernel density estimation (KDE) proxy estimator for estimating

Original: ECE.
Revised: ECE.

Original: However, it uses a un-optimized kernel over ∆ K − 1 , and shows the KDE-ECE estimator (but not the calibration map) is consistent.
Revised: However, it uses a un-optimized kernel over ∆ K − 1 , and shows the KDE-ECE estimator (but not the calibration map) is consistent.

Original: To the best of our knowledge, use of trained KDE to calibrate predictions hasn’t been proposed before.
Revised: To the best of our knowledge, use of trained KDE to calibrate predictions hasn’t been proposed before.

Original: Further, we also provide a bound on the calibration error.
Revised: Further, we also provide a bound on the calibration error.

Original:  KC AL : K ERNEL - BASED C ALIBRATION
Revised:  KC AL : K ERNEL - BASED C ALIBRATION

Original: In this section, we formally introduce KCal, study its calibration properties theoretically, and present crucial implementation details and comparisons with other methods.
Revised: In this section, we formally introduce KCal, study its calibration properties theoretically, and present crucial implementation details and comparisons with other methods.

Original: Specifically, in Section 3.1, we discuss how to construct (automatically) calibrated predictions for test data using a calibration set S cal .
Revised: Specifically, in Section 3.1, we discuss how to construct (automatically) calibrated predictions for test data using a calibration set S cal .

Original: Doing so requires a well-trained kernel and metric space, and we describe a procedure to train such a kernel in Section 3.2.
Revised: Doing so requires a well-trained kernel and metric space, and we describe a procedure to train such a kernel in Section 3.2.

Original: In Section 3.3, we show that an appropriate shrinkage rate of the bandwidth ensures that the KCal prediction is automatically calibrated.
Revised: In Section 3.3, we show that an appropriate shrinkage rate of the bandwidth ensures that the KCal prediction is automatically calibrated.

Original: Sections 3.4 provides implementation details.
Revised: Sections 3.4 provides implementation details.

Original: Finally, in Section 3.5, we compare and contrast KCal with existing methods.
Revised: Finally, in Section 3.5, we compare and contrast KCal with existing methods.

Original: 3. C LASSIFICATION WITH K ERNEL D ENSITY E STIMATION
Revised: 3. C LASSIFICATION WITH K ERNEL D ENSITY E STIMATION

Original: Following the calibration literature, we first require a holdout calibration set S cal
Revised: Following the calibration literature, we first require a holdout calibration set S cal

Original: = { X i , Y i } Ni =1 .
Revised: = { X i , Y i } Ni =1 .

Original: In
Revised: In

Original: KCal, we fix a kernel function ˆ ϕ which is learned (the learning procedure is described in Section 3.2).
Revised: KCal, we fix a kernel function ˆ ϕ which is learned (the learning procedure is described in Section 3.2).

Original: For a new datum X N +1 , the class probability ˆ p k
Revised: For a new datum X N +1 , the class probability ˆ p k

Original: ( X N +1 ) takes the following form:
Revised: ( X N +1 ) takes the following form:

Original: [Equation]
Revised: [Equation]

Original: where S k cal :
Revised: where S k cal :

Original: = { ( x, y ) ∈ S cal | y
Revised: = { ( x, y ) ∈ S cal | y

Original: = k } .
Revised: = k } .

Original: The notation ˆ p k
Revised: The notation ˆ p k

Original: ( X N +1 ; ˆ ϕ, S cal ) emphasizes the dependence on ˆ ϕ and S cal .
Revised: ( X N +1 ; ˆ ϕ, S cal ) emphasizes the dependence on ˆ ϕ and S cal .

Original: However, we will use ˆ p k
Revised: However, we will use ˆ p k

Original: ( X N +1 ) when the dependence is clear from context.
Revised: ( X N +1 ) when the dependence is clear from context.

Original: Remarks : What we have described is essentially the classical nonparametric procedure of applying kernel density estimation for classification.
Revised: Remarks : What we have described is essentially the classical nonparametric procedure of applying kernel density estimation for classification.

Original: For a moment, suppose we know the true density function f k of P k (the distribution of all the data in class k ), and the proportion of class k , denoted π k (such that (cid:80)
Revised: For a moment, suppose we know the true density function f k of P k (the distribution of all the data in class k ), and the proportion of class k , denoted π k (such that (cid:80)

Original: k ∈
Revised: k ∈

Original: [ K ] π k = 1 ).
Revised: [ K ] π k = 1 ).

Original: Then, for any particular x 0 , using the Bayes rule we get:
Revised: Then, for any particular x 0 , using the Bayes rule we get:

Original: [Equation]
Revised: [Equation]

Original: Now, replacing f k with the kernel density estimate ˆ f k ( x 0 ) :
Revised: Now, replacing f k with the kernel density estimate ˆ f k ( x 0 ) :

Original: = ( (cid:80) ( x,y ) ∈S k cal ˆ ϕ b ( x, x 0 )) /
Revised: = ( (cid:80) ( x,y ) ∈S k cal ˆ ϕ b ( x, x 0 )) /

Original: |S k cal | , and the class proportion π k with ˆ π k :=
Revised: |S k cal | , and the class proportion π k with ˆ π k : =

Original: |S k cal | / |S cal | we get back Eq.
Revised: |S k cal | / |S cal | we get back Eq.

Original: 3. T RAINING
Revised: 3. T RAINING

Original: Employing an appropriate kernel function ˆ ϕ is crucial for good performance under the kernel density framework. The kernel  in turn has a critical reliance on the choice of the underlying metric.
Revised: For good performance under the kernel density framework, it is crucial to employ an appropriate kernel function ˆ ϕ , which in turn relies on the choice of the underlying metric.

Original: To obtain good performance using deep learning learning models, we train a metric space on top of the penultimate layer embeddings.
Revised: Therefore, we train a metric space on top of the penultimate layer embeddings of deep learning models.

Original: To begin, we assume a deep neural network is already trained on S train =
Revised: To begin, we assume a deep neural network is already trained on S train =

Original: { X traini , Y traini } Mi =1 .
Revised: { X traini , Y traini } Mi =1 .

Original: We place no limitations on the form of loss function, optimizer, or the model architecture.
Revised: We place no limitations on the form of loss function, optimizer, or the model architecture.

Original: However, we do require the neural net to compute an embedding before a final prediction layer, which is always the case in modern classification models.
Revised: However, we do require the neural net to compute an embedding before a final prediction layer, which is always the case in modern classification models.

Original: We denote the embedding function from X (cid:55)→ R h as f .
Revised: We denote the embedding function from X (cid:55)→ R h as f .

Original: Given a base “mother kernel” function ϕ , such as the Radial Basis Function (RBF) kernel, we denote the kernel with bandwidth b as ϕ b := 1 b ϕ ( · b ) .
Revised: Given a base “mother kernel” function ϕ , such as the Radial Basis Function (RBF) kernel, we denote the kernel with bandwidth b as ϕ b := 1 b ϕ ( · b ) .

Original: We parameterize the learnable kernel as:
Revised: We parameterize the learnable kernel as:

Original: [Equation]
Revised: [Equation]

Original: where Π : R h (cid:55)→ R d is a dimension-reducing projection parameterized by a shallow MLP (Section 3.4).
Revised: where Π : R h (cid:55)→ R d is a dimension-reducing projection parameterized by a shallow MLP (Section 3.4).

Original: Since the inference time is linear in d , letting d < h also affords computational benefits.
Revised: Since the inference time is linear in d , letting d < h also affords computational benefits.

Original: Given that the embedding function f ( x ) from the neural network is fixed, the only learnable entities are b and Π .
Revised: Given that the embedding function f ( x ) from the neural network is fixed, the only learnable entities are b and Π .

Original: In the training phase, we fix b = 1 , and train Π using (stochastic) gradient descent and log-loss.
Revised: In the training phase, we fix b = 1 , and train Π using (stochastic) gradient descent and log-loss.

Original: The specific value of b does not matter since it can be folded into Π .
Revised: The specific value of b does not matter since it can be folded into Π .

Original: Let us denote
Revised: Let us denote

Original: S k train = { ( x, y ) ∈ S train :
Revised: S k train = { ( x, y ) ∈ S train :

Original: y
Revised: y

Original: = k } .
Revised: = k } .

Original: In each iteration, we randomly sample two batches of data from
Revised: In each iteration, we randomly sample two batches of data from

Original: S train - the prediction data, denoted as S B train , to evaluate Π , and “background” data for each k , denoted as B k , from S k train \ S B train to construct the KDE classifier.
Revised: S train - the prediction data, denoted as S B train , to evaluate Π , and “background” data for each k , denoted as B k , from S k train \ S B train to construct the KDE classifier.

Original: Then, the prediction for any x j is given by 
Revised: Then, the prediction for any x j is given by ˆ p k ( x j ; ˆ ϕ, S train \ S B train ) : = (cid:16) (cid:88) ( x,y ) ∈B k |S k train \ S B train | |B k | ˆ ϕ ( x, x j ) (cid:17) / (cid:16) (cid:88) k ′ , ( x,y ) ∈B k ′ |S k ′ train \ S B train | |B k ′ | ˆ ϕ ( x, x j ) (cid:17) (6) where ˆ ϕ is shorthand for ˆ ϕ Π , f ,b =1 defined in Eq. Algorithm Input :

Original: [Equation]
Revised: S train : { ( X traini , Y traini ) } Mi =1 used to train the NN S cal : { ( X i , Y i ) } Ni =1 calibration set f : Embedding function X → R h (trained NN) X N +1 :

Revised: Unseen datum for prediction

Revised: Training (of the projection Π ) :

Revised: Denote S k train :=

Revised: { ( x, y ) ∈ S train | y

Revised: = k } .

Revised: Denote ϕ b as a base kernel function (e.g. RBF) with bandwidth b .

Revised: repeat

Revised: Sample S B train = { ( x j , y j ) } B j =1 from S train .

Revised: Compute ˆ p ( x j ) via Eq.

Revised: Loss l ← 1 B (cid:80)

Revised: Bj =1 LogLoss (ˆ p ( x j ) , y j ) .

Revised: Update Π with (stochastic) gradient descent.

Revised: until the loss l does not improve.

Original: where ˆ ϕ is shorthand for ˆ ϕ Π , f ,b =1 defined in Eq.
Revised: Set ˆ ϕ b ← ˆ ϕ Π , f ,b for inference. Inference : Denote S k cal := { ( x, y ) ∈ S cal | y = k } .

Revised: Tune b ∗ on S cal by minimizing log loss.

Revised: ˆ p k

Revised: ( X N

Revised: [Equation]

Revised: ) ←

Revised: [Equation]

Original: The log-loss is given formally by
Revised: The log-loss is given formally by

Original: [Equation]
Revised: [Equation]

Original: Finally, we pick a b = b ∗ on the calibration set S cal using cross-validation.
Revised: Finally, we pick a b = b ∗ on the calibration set S cal using cross-validation.

Original: This is because b should be chosen contingent on the sample size (Section 3.3).
Revised: This is because b should be chosen contingent on the sample size (Section 3.3).

Original: Choosing b can be done efficiently (Section 3.4).
Revised: Choosing b can be done efficiently (Section 3.4).

Original: Algorithm 1 summarizes the steps we explicated upon so far.
Revised: Algorithm 1 summarizes the steps we explicated upon so far.

Original: 3. T HEORETICAL A NALYSIS : C ALIBRATION C OMES F REE
Revised: 3. T HEORETICAL A NALYSIS : C ALIBRATION C OMES F REE

Original: In the previous section, we have only described a procedure to improve the prediction accuracy for ˆ p on S train .
Revised: In the previous section, we have only described a procedure to improve the prediction accuracy for ˆ p on S train .

Original: This section will show that calibration comes free with the ˆ p obtained using Algorithm 1.
Revised: This section will show that calibration comes free with the ˆ p obtained using Algorithm 1.

Original: In particular, we show that as the sample-size for each class in S cal increases, ˆ p converges to the true frequency vector of Y given the input.
Revised: In particular, we show that as the sample-size for each class in S cal increases, ˆ p converges to the true frequency vector of Y given the prediction.

Original: In interest of smoother presentation, we only state the relevant claims in what follows.
Revised: For smoother presentation, we only state the relevant claims in what follows.

Original: Detailed proofs are presented in the Appendix.
Revised: Detailed proofs are presented in the Appendix.

Original: To begin, we make a few standard assumptions, such as in Chacón & Duong (2018), including:
Revised: To begin, we make a few standard assumptions, such as in Chacón & Duong (2018), including:

Original: • ( ∀ k )
Revised: • ( ∀ k )

Original: The density on the embedded space, Π ( f ( X| Y = k )) , denoted as f Π ◦ f ,k , is square integrable and twice differentiable, with all second order partials bounded, continuous, and square integrable.
Revised: The density on the embedded space, Π ( f ( X| Y = k )) , denoted as f Π ◦ f ,k , is square integrable and twice differentiable, with all second order partials bounded, continuous, and square integrable.

Original: • ϕ is spherically symmetric, with a finite second moment.
Revised: • ϕ is spherically symmetric, with a finite second moment.

Original: Lemma 3.1 and 3.2 focus on an arbitrary class k and ignore the subscript k to the density f for readability.
Revised: Lemma 3.1 and 3.2 focus on an arbitrary class k and ignore the subscript k to the density f for readability.

Original: We denote the size |S k cal | = m .
Revised: We denote the size |S k cal | = m .

Original: Intuitively, due to the bias-variance trade-off, a suitable bandwidth b will depend on m : A small b reduces bias, but with the finite m , a smaller b also leads to increased variance.
Revised: Intuitively, due to the bias-variance trade-off, a suitable bandwidth b will depend on m : A small b reduces bias, but with the finite m , a smaller b also leads to increased variance.

Original: Thus, b should go to 0 “slowly”, which is formally stated below:
Revised: Thus, b should go to 0 “slowly”, which is formally stated below:

Original: Lemma 3.1.
Revised: Lemma 3.1.

Original: For almost all x , if b d m → ∞ and b → 0 as m → ∞ , then we have
Revised: For almost all x , if b d m → ∞ and b → 0 as m → ∞ , then we have

Original: [Equation]
Revised: [Equation]

Original: Here ˆ f Π ◦ f ,k is the estimated f Π ◦ f ,k using S cal .
Revised: Here ˆ f Π ◦ f ,k is the estimated f Π ◦ f ,k using S cal .

Original: Recall that d is the dimension of Π ( f ( X )) .
Revised: Recall that d is the dimension of Π ( f ( X )) .

Original: We will call such a bandwidth b admissible , and we sometimes write b ( m ) to emphasize the dependence on m .
Revised: We will call such a bandwidth b admissible , and we sometimes write b ( m ) to emphasize the dependence on m .

Original: The following lemma gives the optimal admissible bandwidth:
Revised: The following lemma gives the optimal admissible bandwidth:

Original: Lemma 3.2.
Revised: Lemma 3.2.

Original: The optimal bandwidth is b =
Revised: The optimal bandwidth is b =

Original: Θ( m − 1 d +4 ) , which leads to the fastest decreasing MSE (i.e. E [ ∥ ˆ f Π ◦ f ,k ( x ) − f Π ◦ f ,k ( x ) ∥ 2 ] ) of O ( m − 4 d +4 ) .
Revised: Θ( m − 1 d +4 ) , which leads to the fastest decreasing MSE (i.e. E [ ∥ ˆ f Π ◦ f ,k ( x ) − f Π ◦ f ,k ( x ) ∥ 2 ] ) of O ( m − 4 d +4 ) .

Original: Now we are in a position to present the main theoretical results.
Revised: Now we are in a position to present the main theoretical results.

Original: In the following, m denotes the rarest class’s count ( m :=
Revised: In the following, m denotes the rarest class’s count ( m :=

Original: min k {|S k cal |} .
Revised: min k {|S k cal |} .

Original: Theorem 3.3 provides a bound between ˆ p and the true conditional probability vector on the embedded space p ( Π ( f ( X ))) :
Revised: Theorem 3.3 provides a bound between ˆ p and the true conditional probability vector on the embedded space p ( Π ( f ( X ))) :

Original: Theorem 3.3.
Revised: Theorem 3.3.

Original: Fixing x such that the density of Π ( f ( x )) is positive, with b ( m ) =
Revised: Fixing x such that the density of Π ( f ( x )) is positive, with b ( m ) =

Original: Θ( m − 1 d +4 ) , for any λ ∈ (0 , 2) : 
Revised: Θ( m − 1 d +4 ) , for any λ ∈ (0 , 2) : P {| ˆ p k

Original: P
Revised: ( x ) − p k ( Π ( f ( x ))) | > (3 K + 1) Cm − λ d

Original: [Equation]
Revised: [Equation]

Original: 9) where p k ( Π ( f ( x ))) :
Revised: 4 } ≤ Ke − Bm [Equation]

Original: = P { Y = k
Revised: [Equation]

Original: | Π ( f ( X )) =
Revised: where p [Equation]

Original: Π ( f ( x )) } (10) for some constant B and C .
Revised: for some constant B and C .

Original: As a corollary, ˆ p ( x ) almost surely → p ( Π ( f ( x ))) as m → ∞ .
Revised: As a corollary, ˆ p ( x ) almost surely → p [Equation] m → ∞ .

Original: Next, we bound the full calibration error with additional standard assumptions.
Revised: Next, we bound the full calibration error with additional standard assumptions.

Original: More specifically, we use and build upon the main uniform convergence result for classical KDE presented in Jiang (2017), to obtain Theorem 3.4:
Revised: More specifically, we use and build upon the main uniform convergence result for classical KDE presented in Jiang (2017), to obtain Theorem 3.4:

Original: Theorem 3.4.
Revised: Theorem 3.4.

Original: Assume f
Revised: Assume f

Original: Π ◦ f ,k is α -Hölder continuous and bounded away from 0 for any k .
Revised: Π ◦ f ,k is α -Hölder continuous and bounded away from 0 for any k .

Original: For an admissible b ( m ) with shrinkage rate Θ(( log mm ) 1 d +2 α ) , for some constants B and C we have:
Revised: For an admissible b ( m ) with shrinkage rate Θ(( log mm ) 1 d +2 α ) , for some constants B and C we have:

Original: [Equation]
Revised: [Equation]

Original: Input f : Embedding function X → R h (trained NN)

Original: X N +1 :

Original: Unseen datum for prediction

Original: Training (of the projection Π ) :

Original: Denote S k train :=

Original: { ( x, y ) ∈ S train | y

Original: = k } .

Original: Denote ϕ b as a base kernel function (e.g. RBF) with bandwidth b .

Original: repeat

Original: Sample S B train = { ( x j , y j ) }

Original: Bj =1 from S train .

Original: Compute ˆ p ( x j ) via Eq.

Original: Loss l ← 1 B (cid:80)

Original: Bj =1 LogLoss (ˆ p ( x j ) , y j ) .

Original: Update Π with (stochastic) gradient descent.

Original: until the loss l does not improve.

Original: Set ˆ ϕ b ← ˆ ϕ Π , f ,b for inference.

Original: Inference :

Original: Denote S k cal :=

Original: { ( x, y ) ∈ S cal | y = k } .

Original: Tune b ∗ on S cal by minimizing log loss.

Original: ˆ

Original: [Equation] ← [Equation] We now proceed to present details pertaining to the efficient implementation of KCal.
Revised: We now proceed to present details pertaining to the efficient implementation of KCal.

Original: 3. I MPLEMENTATION T ECHNIQUES
Revised: 3. I MPLEMENTATION T ECHNIQUES

Original: Efficient Training : As might be immediately apparent, utilizing algorithm 1 for prediction using full S train \ S B train can be an expensive ex- ercise.
Revised: Efficient Training : As might be immediately apparent, utilizing algorithm 1 for prediction using full S train \ S B train can be an expensive exercise.

Original: In order to afford a training speedup, we consider a random subset from S train \ S B train us- ing a modified stratified sampling.
Revised: In order to afford a training speedup, we consider a random subset from S train \ S B train using a modified stratified sampling.

Original: Specifically, we take m random samples from each S k train , denoted as S k,m train , and replace the right-hand side of Eq.
Revised: Specifically, we take m random samples from each S k train , denoted as S k,m train , and replace the right-hand side of Eq.

Original: 6 with:
Revised: 6 with:

Original: [Equation]
Revised: [Equation]

Original: The re-scaling term |S k train | m is crucial to get an unbiased estimate of ˆ f k ˆ π k .
Revised: The re-scaling term |S k train | m is crucial to get an unbiased estimate of ˆ f k ˆ π k .

Original: The stratification employed makes the training more stable while also reducing the estimation variance for 
Revised: The stratification employed makes the training more stable, while also reducing the estimation variance for rarer classes (more details in Appendix B).

Original: rarer

Original: classes (more details in Appendix B).

Original: The overall complexity is now O ( KmdhB ) per batch.
Revised: The overall complexity is now O ( KmdhB ) per batch.

Original: In all experiments, we used m = 20 and B = 64 .
Revised: In all experiments, we used m = 20 and B = 64 .

Original: Form of Π : While there is considerable freedom in choosing a suitable form for Π , we parameterize Π with a two layer MLP with a skip connection.
Revised: Form of Π : While there is considerable freedom in choosing a suitable form for Π , we parameterize Π with a two layer MLP with a skip connection.

Original: Consequently, Π can reduce to linear projection when sufficient, and be more expressive when necessary.
Revised: Consequently, Π can reduce to linear projection when sufficient, and be more expressive when necessary.

Original: We also experimented with using only a linear projection, the results for which are included in the appendix.
Revised: We also experimented with using only a linear projection, the results for which are included in the appendix.

Original: We fix the output dimension to d = min { dim ( f ) , 32 } , except for ImageNet ( d = 128 ).
Revised: We fix the output dimension to d = min { dim ( f ) , 32 } , except for ImageNet ( d = 128 ).

Original: Bandwidth Selection : Finally, to find the optimal bandwidth using S cal , we use Golden-Section search (Kiefer, 1953) to find the log-loss-minimizing b ∗ .
Revised: Bandwidth Selection : Finally, to find the optimal bandwidth using S cal , we use Golden-Section search (Kiefer, 1953) to find the log-loss-minimizing b ∗ .

Original: This takes O (log ub − lb tol ) steps where [ lb, ub ] is the search space, and tol is the tolerance.
Revised: This takes O (log ub − lb tol ) steps where [ lb, ub ] is the search space, and tol is the tolerance.

Original: Essentially, we assume that the loss is a convex function with respect to b , permitting an efficient search (see Appendix H, which presents empirical evidence that the convexity assumption is valid across datasets).
Revised: Essentially, we assume that the loss is a convex function with respect to b , permitting an efficient search (see Appendix H, which presents empirical evidence that the convexity assumption is valid across datasets).

Original: 3. C OMPARISONS WITH E XISTING C ALIBRATION M ETHODS
Revised: 3. C OMPARISONS WITH E XISTING C ALIBRATION M ETHODS

Original: Most existing calibration methods discussed in Section 2 and KCal all utilize a holdout calibration set.
Revised: Most existing calibration methods discussed in Section 2 and KCal all utilize a holdout calibration set.

Original: However, unlike KCal, existing works usually fix the last neural network layer.
Revised: However, unlike KCal, existing works usually fix the last neural network layer.

Original: KCal, on the other hand, “takes a step back”, and replaces the last prediction layer with a kernel density estimation based classifier.
Revised: KCal, on the other hand, “takes a step back”, and replaces the last prediction layer with a kernel density estimation based classifier.

Original: Since the DNN f is fixed regardless of whether we use the original last layer or not, we are really comparing a KDE classifier (KCal) with linear models trained in various ways, after mapping all the data with f .
Revised: Since the DNN f is fixed regardless of whether we use the original last layer or not, we are really comparing a KDE classifier (KCal) with linear models trained in various ways, after mapping all the data with f .

Original: Note that this characterization is true for most existing methods, with a few exceptions (e.g., those summarized under “loss regularization” in Section 2).
Revised: Note that this characterization is true for most existing methods, with a few exceptions (e.g., those summarized under “loss regularization” in Section 2).

Original: Employing a KDE classifier affords some clear advantages such as a straightforward convergence guarantee and some interpretability 1 .
Revised: Employing a KDE classifier affords some clear advantages such as a straightforward convergence guarantee and some interpretability 1 .

Original: Furthermore, KCal can also be improved in an online fashion, a benefit especially desirable in certain high-stakes applications such as in healthcare.
Revised: Furthermore, KCal can also be improved in an online fashion, a benefit especially desirable in certain high-stakes applications such as in healthcare.

Original: For example, a hospital can calibrate a trained model prior to deployment using its own patient data (which is usually not available to train the original model) as it becomes available.
Revised: For example, a hospital can calibrate a trained model prior to deployment using its own patient data (which is usually not available to train the original model) as it becomes available.

Original: Another important advantage of KCal is concerning normalization.
Revised: Another important advantage of KCal is concerning normalization.

Original: In fact, simultaneously calibrating all classes while satisfying the constraint that ˆ
Revised: In fact, simultaneously calibrating all classes while satisfying the constraint that ˆ

Original: p ∈ ∆ K − 1 is a distinguishing challenge for multi-class calibration.
Revised: p ∈ ∆ K − 1 is a distinguishing challenge for multi-class calibration.

Original: Many calibration methods perform one-vs-rest calibration for each class, and require a separate normalization step at test time (Zadrozny & Elkan, 2001; 2002; Patel et al., 2021; Gupta et al., 2021).
Revised: Many calibration methods perform one-vs-rest calibration for each class, and require a separate normalization step at test time (Zadrozny & Elkan, 2001; 2002; Patel et al., 2021; Gupta et al., 2021).

Original: This creates a gap between training and testing and could lead to drastic drop in performance (Section 4).
Revised: This creates a gap between training and testing and could lead to drastic drop in performance (Section 4).

Original: On the other hand, KCal automatically satisfies ˆ
Revised: On the other hand, KCal automatically satisfies ˆ

Original: p ∈
Revised: p ∈

Original: ∆ K − 1 , and the normalization is consistent during training and testing.
Revised: ∆ K − 1 , and the normalization is consistent during training and testing.

Original: A disadvantage of KCal is the need to remember the Π ( f ( S cal )) used to generate the KDE prediction.
Revised: A disadvantage of KCal is the need to remember the Π ( f ( S cal )) used to generate the KDE prediction.

Original: This is however mitigated to a large extent by the dimension reduction step, which already reduces the computational overhead significantly 2 .
Revised: This is however mitigated to a large extent by the dimension reduction step, which already reduces the computational overhead significantly 2 .

Original: For example, in one of our experiments on CIFAR-100, there are 160K (5K images, d = 32 ) scalars to remember, which is only 0.2% of the parameters (85M+) of the original DNN (ViT-base-patch16).
Revised: For example, in one of our experiments on CIFAR-100, there are 160K (5K images, d = 32 ) scalars to remember, which is only 0.2% of the parameters (85M+) of the original DNN (ViT-base-patch16).

Original: Moreover, KDE inference is trivial to parallelize on
Revised: Moreover, KDE inference is trivial to parallelize on

Original: GPUs.
Revised: GPUs.

Original: There is also a rich, under-explored, literature to further speed up the inference.
Revised: There is also a rich, under-explored, literature to further speed up the inference.

Original: Examples include, KDE merging (Sodkomkham et al., 2016),
Revised: Examples include, KDE merging (Sodkomkham et al., 2016),

Original: Dual-Tree (Gray & Moore, 2003), and Kernel
Revised: Dual-Tree (Gray & Moore, 2003), and Kernel

Original: Herding (Chen et al., 2010).
Revised: Herding (Chen et al., 2010).

Original: These methods can easily be used in conjunction with KCal.
Revised: These methods can easily be used in conjunction with KCal.

Original:  E XPERIMENTS 4. D ATA AND N EURAL N ETWORKS
Revised:  E XPERIMENTS 4. D ATA AND N EURAL N ETWORKS

Original: We utilize two sets of data: computer vision benchmarks on which previous calibration methods were tested, and health monitoring datasets where full calibration is crucial for diagnostic applications.
Revised: We utilize two sets of data: computer vision benchmarks on which previous calibration methods were tested, and health monitoring datasets where full calibration is crucial for diagnostic applications.

Original: Table 1 summarizes the datasets and their splits.
Revised: Table 1 summarizes the datasets and their splits.

Original: [Table]
Revised: [Table]

Original: Benchmark data Following Kull et al.
Revised: Benchmark data Following Kull et al.

Original: (2019), we use multiple image benchmark datasets, including
Revised: (2019), we use multiple image benchmark datasets, including

Original: CIFAR-10, CIFAR-100, and SVHN (Krizhevsky, 2009; Netzer et al., 2011).
Revised: CIFAR-10, CIFAR-100, and SVHN (Krizhevsky, 2009; Netzer et al., 2011).

Original: We reserve 10% of the training data as the calibration set.
Revised: We reserve 10% of the training data as the calibration set.

Original: We fine-tune pretrained ViT (Dosovitskiy et al., 2021) and
Revised: We fine-tune pretrained ViT (Dosovitskiy et al., 2021) and

Original: MLP-Mixer (Mixer)
Revised: MLP-Mixer (Mixer)

Original: (Tolstikhin et al., 2021) from the timm library (Wightman, 2019).
Revised: (Tolstikhin et al., 2021) from the timm library (Wightman, 2019).

Original: We chose ViT and Mixer because they are the state-of-the-art neural architectures in computer vision, and accuracy should come before calibration quality.
Revised: We chose ViT and Mixer because they are the state-of-the-art neural architectures in computer vision, and accuracy should come before calibration quality.

Original: We also included the ImageNet dataset (Deng et al., 2009) and use the pretrained Inception ResNet V2 (Szegedy et al., 2017) following Patel et al.
Revised: We also included the ImageNet dataset (Deng et al., 2009) and use the pretrained Inception ResNet V2 (Szegedy et al., 2017) following Patel et al.

Original: Health monitoring data We also use three health monitoring datasets for diagnostic tasks: IIIC (Jing et al., 2021), an ictal-interictal-injury-continuum (IIIC) patterns classification dataset; ISRUC (Khalighi et al., 2016), a sleep staging (classification) dataset using polysomnographic (PSG) recordings;
Revised: Health monitoring data We also use three health monitoring datasets for diagnostic tasks: IIIC (Jing et al., 2021), an ictal-interictal-injury-continuum (IIIC) patterns classification dataset; ISRUC (Khalighi et al., 2016), a sleep staging (classification) dataset using polysomnographic (PSG) recordings;

Original: PN2017 (2017 PhysioNet Challenge) (Clifford et al., 2017; Goldberger et al., 2000), a public electrocardiogram (ECG) dataset for rhythm (particularly Atrial Fibrillation) classification.
Revised: PN2017 (2017 PhysioNet Challenge) (Clifford et al., 2017; Goldberger et al., 2000), a public electrocardiogram (ECG) dataset for rhythm (particularly Atrial Fibrillation) classification.

Original: For the training set, we follow Hong et al.
Revised: For the training set, we follow Hong et al.

Original: ; Jing et al.
Revised: ; Jing et al.

Original: (2021) for PN2017 and IIIC, and used 69 patients’ data for ISRUC.
Revised: (2021) for PN2017 and IIIC, and used 69 patients’ data for ISRUC.

Original: For the remaining data, 5% is used as the calibration set and 95% for testing.
Revised: For the remaining data, 5% is used as the calibration set and 95% for testing.

Original: We perform additional experiments after splitting into training/calibration/test sets by patients for IIIC and ISRUC 3 , marked as the “pat” version in tables.
Revised: We perform additional experiments after splitting into training/calibration/test sets by patients for IIIC and ISRUC 3 , marked as the “pat” version in tables.

Original: The calibration/test split is 20/80 in “IIIC (pat)” and “ISRUC (pat)” because the number of patients is small.
Revised: The calibration/test split is 20/80 in “IIIC (pat)” and “ISRUC (pat)” because the number of patients is small.

Original: For IIIC and ISRUC, we follow the standard practice and train a CNN (ResNet) on the spectrogram (Biswal et al., 2017; Ruffini et al., 2019; Yuan et al., 2019).
Revised: For IIIC and ISRUC, we follow the standard practice and train a CNN (ResNet) on the spectrogram (Biswal et al., 2017; Ruffini et al., 2019; Yuan et al., 2019; Yang et al., 2022).

Original: For PN2017, we used a top-performing model from the 2017 PhysioNet 
Revised: For PN2017, we used a top-performing model from the 2017 PhysioNet Challenge, MINA

Original: Challenge, MINA (Hong et al., 2019).
Revised: (Hong et al., 2019).

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table]
Revised: [Table]

Original: 4. B ASELINES M ETHODS
Revised: 4. B ASELINES M ETHODS

Original: We compare KCal with the multiple state-of-the-art calibration methods, including Temperature
Revised: We compare KCal with the multiple state-of-the-art calibration methods, including Temperature

Original: Scaling ( TS ) (Guo et al., 2017), Dirichlet Calibration ( DirCal ) (Kull et al., 2019), Mutualinformation-maximization-based Binning ( I-Max )
Revised: Scaling ( TS ) (Guo et al., 2017), Dirichlet Calibration ( DirCal ) (Kull et al., 2019), Mutualinformation-maximization-based Binning ( I-Max )

Original: (Patel et al., 2021), Gaussian Process Calibration
Revised: (Patel et al., 2021), Gaussian Process Calibration

Original: [Table]
Revised: [Table]

Original: baselines on Accuracy, CECE and Brier, and the difference between most methods on ECE is small.
Revised: on Accuracy, CECE and Brier, and the difference between most methods on ECE is small.

Original: [Table]
Revised: [Table]

Original: ( GP ) (Wenger et al., 2020), Intra Order-preserving Calibration ( IOP ) (Rahimi et al., 2020), Splinesbased Calibration ( Spline ) (Gupta et al., 2021), Focal-loss-based calibration ( Focal ) (Mukhoti et al., 2020), MMCE-based calibration ( MMCE ) (Kumar et al., 2018).
Revised: ( GP ) (Wenger et al., 2020), Intra Order-preserving Calibration ( IOP ) (Rahimi et al., 2020), Splinesbased Calibration ( Spline ) (Gupta et al., 2021), Focal-loss-based calibration ( Focal ) (Mukhoti et al., 2020), MMCE-based calibration ( MMCE ) (Kumar et al., 2018).

Original: 4. E VALUATION M ETRICS
Revised: 4. E VALUATION M ETRICS

Original: We report standard evaluation metrics: Accuracy, class-wise expected calibration error (CECE) (Kull et al., 2019; Patel et al., 2021; Nixon et al., 2019), expected calibration error (ECE) (Guo et al., 2017), and Brier score (Brier, 1950).
Revised: We report standard evaluation metrics: Accuracy, class-wise expected calibration error (CECE) (Kull et al., 2019; Patel et al., 2021; Nixon et al., 2019), expected calibration error (ECE) (Guo et al., 2017), and Brier score (Brier, 1950).

Original: CECE is typically used as a proxy to evaluate full calibration quality, because directly binning basing on the entire vector ˆ p requires exponentially (in K ) many bins.
Revised: CECE is typically used as a proxy to evaluate full calibration quality, because directly binning basing on the entire vector ˆ p requires exponentially (in K ) many bins.

Original: Similar to Patel et al.
Revised: Similar to Patel et al.

Original: ; Nixon et al.
Revised: ; Nixon et al.

Original: (2019), we ignore all predictions with very small probabilities (less than max { 0 . 01 , 1 K } ).
Revised: (2019), we ignore all predictions with very small probabilities (less than max { 0 . 01 , 1 K } ).

Original: ECE, on the other hand, only measures confidence calibration (Def 2).
Revised: ECE, on the other hand, only measures confidence calibration (Def 2).

Original: For both ECE and CECE, we use the “adaptive” version with equal number of samples in each bin (with 20 bins), because this is shown to measure the calibration quality better than the equal-width version (Nixon et al., 2019).
Revised: For both ECE and CECE, we use the “adaptive” version with equal number of samples in each bin (with 20 bins), because this is shown to measure the calibration quality better than the equal-width version (Nixon et al., 2019).

Original: Brier score can be viewed as the sum of a “calibration” term, and a “refinement” term measuring how discriminative a model is (Kull & Flach, 2015).
Revised: Brier score can be viewed as the sum of a “calibration” term, and a “refinement” term measuring how discriminative a model is (Kull & Flach, 2015).

Original: Here we focus on the brier score of the top class.
Revised: Here we focus on the brier score of the top class.

Original: We refer to (Guo et al., 2017; Kull et al., 2019; Nixon et al., 2019) for further discussion on these metrics.
Revised: We refer to (Guo et al., 2017; Kull et al., 2019; Nixon et al., 2019) for further discussion on these metrics.

Original: 4. R ESULTS
Revised: 4. R ESULTS

Original: The results are presented in Tables 2, 3, 4 and 5.
Revised: The results are presented in Tables 2, 3, 4 and 5.

Original: All experiments are repeated 10 times by reshuffling calibration and test sets, and the standard deviations are reported.
Revised: All experiments are repeated 10 times by reshuffling calibration and test sets, and the standard deviations are reported.

Original: For ImageNet, we skipped Focal and MMCE because the base NN is given and these methods require training from scratch.
Revised: For ImageNet, we skipped Focal and MMCE because the base NN is given and these methods require training from scratch.

Original: Due to space constraints, we include ablation studies in the Appendix.
Revised: Due to space constraints, we include ablation studies in the Appendix.

Original: In general, KCal has the best CECE, accuracy and Brier score, and is highly competitive in terms of ECE as well.
Revised: In general, KCal has the best CECE, accuracy and Brier score, and is highly competitive in terms of ECE as well.

Original: Note that KCal is also the only method with provable calibration guarantee.
Revised: Note that KCal is also the only method with provable calibration guarantee.

Original: TS is effective in controlling overall ECE but shows little improvement on CECE over UnCal .
Revised: TS is effective in controlling overall ECE but shows little improvement on CECE over UnCal .

Original: DirCal often ranks high for the calibration quality but tends to decrease accuracy as K increases.
Revised: DirCal often ranks high for the calibration quality but tends to decrease accuracy as K increases.

Original: DirCal ’s performance also has a higher cost: Every experiment requires training over hundreds of models with SGD and taking the best ensemble, accounting for most of the experiment computation cost in this paper.
Revised: DirCal ’s performance also has a higher cost: Every experiment requires training over hundreds of models with SGD and taking the best ensemble, accounting for most of the experiment computation cost in this paper.

Original: The amount of tuning suggested for good performance indicates sensitivity to the choice of hyper-parameters, which we have indeed observed to be the case.
Revised: The amount of tuning suggested for good performance indicates sensitivity to the choice of hyper-parameters, which we have indeed observed to be the case.

Original: Spline , IOP and GP are similar to DirCal on vision datasets, but generally perform worse on the healthcare datasets.
Revised: Spline , IOP and GP are similar to DirCal on vision datasets, but generally perform worse on the healthcare datasets.

Original: In Patel et al.
Revised: In Patel et al.

Original: (2021), I-Max lowers ECE and CECE significantly.
Revised: (2021), I-Max lowers ECE and CECE significantly.

Original: However, it has a critical issue - it does not produce a valid probability vector 4 .
Revised: However, it has a critical issue - it does not produce a valid probability vector 4 .

Original: Once normalized, as reported in our experiments, the performance worsens.
Revised: Once normalized, as reported in our experiments, the performance worsens.

Original: Since calibrating all the classes simultaneously is the distinguishing challenge in multiclass classification, we interpret the observation as: If this normalization constraint is removed, the “optimization problem” (to lower calibration error) is much simpler, but the results are invalid hence unusable probability vectors.
Revised: Since calibrating all the classes simultaneously is the distinguishing challenge in multiclass classification, we interpret the observation as: If this normalization constraint is removed, the “optimization problem” (to lower calibration error) is much simpler, but the results are invalid hence unusable probability vectors.

Original: Spline also requires a re-normalization step, but its performance stays consistent.
Revised: Spline also requires a re-normalization step, but its performance stays consistent.

Original: Focal is worse than the UnCal in many experiments.
Revised: Focal is worse than the UnCal in many experiments.

Original: While calibration performance may improve by combing Focal with other methods, the drop in accuracy is harder to overcome 5 .
Revised: While calibration performance may improve by combing Focal with other methods, the drop in accuracy is harder to overcome 5 .

Original: We also observed that for healthcare datasets, being able to tune on a different set of patients boosts the performance significantly.
Revised: We also observed that for healthcare datasets, being able to tune on a different set of patients boosts the performance significantly.

Original: This is reflected in the accuracy gain for DirCal and KCal, and suggests that the embeddings/logits are quite transferable, but the prediction criteria itself can vary from patient to patient.
Revised: This is reflected in the accuracy gain for DirCal and KCal, and suggests that the embeddings/logits are quite transferable, but the prediction criteria itself can vary from patient to patient.

Original: Finally, we summarize the rankings of all datasets in Table 6.
Revised: Finally, we summarize the rankings of all datasets in Table 6.

Original: It is clear that KCal consistently improves calibration quality for all classes and maintains or improves accuracy.
Revised: It is clear that KCal consistently improves calibration quality for all classes and maintains or improves accuracy.

Original: And if we look at only the confidence prediction (Brier or ECE), KCal is still highly competitive.
Revised: And if we look at only the confidence prediction (Brier or ECE), KCal is still highly competitive.

Original: [Figure]
Revised: [Figure]

Original:  C : C 
Revised: III C : C o n

Revised: [Figure]

Original: o n f i d e n c 
Revised: i d e n c e

Original: [Figure]
Revised: [Figure]

Original: 4. C ASE S TUDY FOR S EIZURE P REDICTION
Revised: 4. C ASE S TUDY FOR S EIZURE P REDICTION

Original: We show the reliability diagrams (Kull et al., 2019; Guo et al., 2017) on the IIIC dataset to illustrate the importance of full calibration in Figure 2.
Revised: We show the reliability diagrams (Kull et al., 2019; Guo et al., 2017) on the IIIC dataset to illustrate the importance of full calibration in Figure 2.

Original: We include both the the predicted class (confidence calibration) and Seizure.
Revised: We include both the the predicted class (confidence calibration) and Seizure.

Original: More reliability diagrams can be found in the Appendix, and the results are consistent for all classes.
Revised: More reliability diagrams can be found in the Appendix, and the results are consistent for all classes.

Original: The un-calibrated predictions have large gaps for both confidence and
Revised: The un-calibrated predictions have large gaps for both confidence and

Original: Seizure.
Revised: Seizure.

Original: Most baselines provide calibrated confidence calibration, but fail to calibrated the output for the rare class Seizure.
Revised: Most baselines provide calibrated confidence calibration, but fail to calibrated the output for the rare class Seizure.

Original: KCal, on the other hand, achieves the most consistent results.
Revised: KCal, on the other hand, achieves the most consistent results.

Original: We note again that since all competing classes must be considered together for any clinical decision, full calibration is indispensable in medical applications.
Revised: We note again that since all competing classes must be considered together for any clinical decision, full calibration is indispensable in medical applications.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: This paper proposed KCal, a learned-kernel-based calibration method for deep learning models.
Revised: This paper proposed KCal, a learned-kernel-based calibration method for deep learning models.

Original: KCal consists of a supervised dimensionality reduction step on the penultimate layer neural network embedding to improve efficiency.
Revised: KCal consists of a supervised dimensionality reduction step on the penultimate layer neural network embedding to improve efficiency.

Original: A KDE classifier using the calibration set is employed in this new metric space.
Revised: A KDE classifier using the calibration set is employed in this new metric space.

Original: As a natural consequence of the construction, KCal provides a calibrated probability vector prediction for all classes.
Revised: As a natural consequence of the construction, KCal provides a calibrated probability vector prediction for all classes.

Original: Unlike most existing calibration methods, KCal is also provably asymptotically fully calibrated with finite sample error bounds.
Revised: Unlike most existing calibration methods, KCal is also provably asymptotically fully calibrated with finite sample error bounds.

Original: We also showed that empirically, it outperforms existing state-of-the-art calibration methods in terms of accuracy and calibration quality.
Revised: We also showed that empirically, it outperforms existing state-of-the-art calibration methods in terms of accuracy and calibration quality.

Original: Moreover, KCal is more robust to distributional shift, which is common in high-risk applications such as healthcare, where calibration is far more crucial.
Revised: Moreover, KCal is more robust to distributional shift, which is common in high-risk applications such as healthcare, where calibration is far more crucial.

Original: The major limitation of KCal is the need to store the entire calibration set, which is a small overhead with the dimension reduction step and potential improvements. [Table]
Revised: The major limitation of KCal is the need to store the entire calibration set, which is a small overhead with the dimension reduction step and potential improvements. A CKNOWLEDGMENTS

Original: [Table]
Revised: This work was supported by NSF award SCH-2205289, SCH-2014438, IIS-1838042, NIH award

Original: [Table] [Table]
Revised: R01 1R01NS107291-01.