A Reproducibility Study of Decoupling Feature Extraction and Classification Layers for Calibrated Neural Networks
Authors: Johanna D'Ciofalo Khodaverdian
Eric Banzuzi
Katharina Deckenbach
Venue: Accepted by TMLR ()

Abstract: Many neural networks, especially over-parameterized ones, suffer from poor calibration and overconfidence. To address this, Jordahn & Olmos (2024) recently proposed a Two-Stage Training (TST) procedure that decouples the training of feature extraction and classification layers. In this study, we replicate their findings and extend their work through a series of ablation studies. We reproduce their main results and find that most of them replicate, with slight deviation for CIFAR100. Additionally, we extend the author's results by exploring the impact of different model architectures, Monte Carlo (MC) sample sizes, and classification head designs. We further compare the method with focal loss -- an implicit regularization technique known to improve calibration -- and investigate whether calibration can be improved further by combining the two methods. Beyond focal loss, we also evaluate the effect of incorporating other similar regularization techniques such as label smoothing and L2 regularization during two-stage training. We find that calibration can be improved even further by using focal loss in the first training stage of two-stage training. Similar improvements are observed when combining two-stage training with label smoothing and L2 regularization. Our experiments validate the claims made by Jordahn & Olmos (2024), and show the transferability of the two-stage training to different architectures.

Published in Transactions on Machine Learning Research (06/2025)
A Reproducibility Study of Decoupling Feature Extraction and Classification Layers for Calibrated Neural Networks
Eric Banzuzi
banauzi@kth.se
KTH Royal Institute of Technology
Johanna D'Ciofalo Khodaverdian
jdkookth.se
KTH Royal Institute of Technology
Katharina Deckenbach
kadec@kth.ss se
KTH Royal Institute of Technology
Reviewed on OpenReview: https: // opw.reviee.r........mmo. net/ forum? id= 6Huz f
Abstract
Many neural networks, especially over-parameterized ones, suffer from poor calibration and overconfidence. To address this, Jordahn & Olmos (2024) recently proposed a Two-Stage Training (TST) procedure that decouples the training of feature extraction and classifica- tion layers. In this study, we replicate their findings and extend their work through a series of ablation studies. We reproduce their main results and find that most of them replicate, with slight deviation for CIFAR100. Additionally, ploring the impact of different model architectures, Monte Carlo (MC) sample sizes, and classification head designs. We further compare the method with focal loss - an implicit regularization technique known to improve calibration and investigate whether calibration can be improved further by combining the two methods. Beyond focal loss, we also evaluate the effect of incorporating other similar regularization techniques such as label smoothing and L2 regularization during (wo-stage training. We find that calibration can be improved further by using focal loss in the first training stage of two-stage training. Similar improvements observed when combining two-stage training with label smoothing and L2 regularization. Our experiments validate the claims made by Jordahn & Olmos (2024). and show the transferability of the two-stage training to different architectures.
extend the author's results by ex-
1 Introduction
Calibration is crucial for trustworthy machine learning, especially in safety-critical fields like medicine and self-driving cars, where understanding model's uncertainty is key to assessing decision reliability. Neural networks, particularly paraperraeeer ones, tend to be poorly calibrated and have an overconfidence bias. Efforts to improve calibration include post-hoc methods and modifications to training procedures (Wang
2023
In Decoupling Feature Extraction and Classification Layers for Calibrated Neural Networks, Jordahn & Ol- mos (2024) propose Two-Stage Training (TST) and Variational Two-Stage Training (V-TST), for calibration of Deep Neural Networks (DNNs). The methods decouple the training of the feature extraction and classifi- cation layers. At first, all layers are trained jointly from scratch and in the second stage of training, the last layer is re-initialized and re-trained from scratch while the weights of the feature extraction layers are kept frozen.
We conducted several experiments to analyze the performance of TST and V-TST under different configu- rations, and extend their method further by combining with focal loss both its constant and adaptive variant. Our main contributions are:

---

Published in Transactions on Machine Learning Research (06/2025)
Reproducing the main results (first two rows of Tables 1 and 2 of Jordahn & Olmos (2244)
Reproducing two of their ablation studies, namely applying their method an under- cader-parameterized CNN and fine-tuned ViT (Tables and of Jordahn & Olmos (2024))
Conducting three ablation studies to investigate the dependence on model architecture, the sample size of a Monte Carlo (MC) estimation used in the second-stage of training and the architecture of the classification head.
Further extending the original paper by comparing and combining the method with three different regularization techniques, namely focal loss (Lin et al. 2017), label smoothing (Szegedy et al. 2016) and L2 regularization.
2 Related Work
Neural network calibration methods can be grouped into four different categories Waag 2023 Gawlikowski
ct al.
20230: post-hoc calibration, regularization methods, uncertainty estimation, and methods that combine the aforementioned categories. Within the category of post-hoc methods, Temperature Scaling (TS) (Guo 2017) is widely used, simple and computationally efficient approach with many variations (Wang, Post-hoc methods are generally characterized by their applicability after training. While the main goal of uncertainty cstimation approaches is to provide uncertainty estimates of model prodictions, they can also improve calibration by reducing epistemic uncertainly (Gawlikowski et al. 2023). Prominent examples of these techniques are and ensemble methods and Bayesian neural networks (BNNs). Regularization techniques
et al.
2023
(Szegedy et al.] 2016). L,2 regularization Guo et al.] (2017). exposure to OOD samples (Hendrycks et al. 2019), objective function
encompass data augmentation (Thulasidasan et al., 2019). label smoothing (IS)
modification (Pereyra al.) 2017], methods that. use calibration errors as differentiable proxy (Kumar al.
2018
Wang 2023) as well as the use of different loss functions (Hebbalaguppe et al. 2022 Shamsi et al. Lin ct al. 2017).
2023
For this study, implicit regularization techniques are the most relevant as TST and V-TST fall into this category. An example of such an implicit technique Focal Loss Lin et al. 2017 defined as
FL(p:) (1 - pc) log (pt)
(1)
Focal loss is a modification of cross-entropy loss designed to address class imbalance during training. It introduces a weighting factor, (1 (1- p,)", where pt is the predicted probability of the true class. This factor down-weights the loss of well-classifidd examples (pr 1) and amplifics the influence of hard examples (pr & 0), making training more effective. The focusing parameter 7 2 0 determines the degree of this effect. Focal loss has been shown to be an upper bound on the regularized KL-divergence and improves calibration of networks independent of their architecture (see Mukhoti et al. (2220) for a more elaborate discussion on why focal loss improves calibration). The success of focal loss has been connected to its ability to control the curvalure of the loss landscape (Kimura Naganuma) 2025). Several advancements of focal loss have been suggested, such as dual focal loss (Tao et al.) (023a), calibration-awar aware adaptive focal loss called AdaFocal combining focal loss with temperature scaling Komisarenko & Kull 2024
GGooh al. 2022
Another widely used implicit regularization technique is label smoothing (Szeeddy et al. 2016 which reduces the confidence of predictions by computing the cross-entropy loss with a smoothened target distribution that assigns a small portion of the probability mass to incorrect classes. Instead of assigning the true probability to the target class, label smoothing assigns a slightly lower probability 1-c to the target class and distributes the remaining e uniformly across the incorrect classes. Similarly to focal loss, label smoothing regularization reduces overconfidence and generally improves model calibration (Müller et al., 2019
3 Methods
Jordahn & Olmos 2024) propose a novel two-stage training procedure for improved calibration of DNNs in classification tasks. Typically, the training of feature extraction layers and classification layers occurs
2

---

Published in Transactions on Machine Learning Research (06/2025)
jointly. Their findings indicate that training these layers separately in two-stage fashion can improve model calibration. They argue that freezing the feature extraction layers limits the model's flexibility, preventing the classifier from artificially increasing label likelihood. Thus, their method falls in the category of regularization techniques, and the success of their method is connected to finding the perfect regularization for learned feature extractor (Jordahn 2024 In their paper, they present two variat ions of the two-stage training method: TST and V-TST.
3.1 TST
In TST, a model with initial parameters {8, Ø), denoting the parameters of the feature extraction layers and classification layers respectively, is trained from scratch using cross-entropy loss (Jordahn & Olmos 2024). The parameters of the feature extraction layers are then frozen and the classification layers are re-init ialzed and re-trained from scratch using the same datasct Drain (sce Fig Further calibration improvements can be obtained by adjusting the dimensionality Z of the last hidden layer z of the classification layers, as the calibration performance varies with different choices of Z.
Feature Extraction
Classification
Feature Extraction Classification
Input Image
Ioput Image
Stage 1:
Stage 2:
(a) First training stage: feature extraction and classification layers are trained jointly from scratch.
(b) Second training stage: feature extraction layers are frozen and classification layers are re-initialized and trained from scratch.
Figure Mlusttaation of training procedure for TST and V-TST. Blue indicates the layers are being trained while gray indicates frozen layers.
3.2 V-TST
V-TST poses additional regularization For the second stage, the parameters of the fully connected (FC) layers are optimized using the evidence lower-bound (ELBO):
log ply) Z ELBO,,- (ogesll)) -Dre Drr((((()))l(())))
(2)
They assume a standard normal Gaussian prior p(z) A'(0, I) and an approximate posterior of the form g(z v) = (x)
The predicted output label of the model is given by p(y|=) Cat(TT To approximate and sample from q<][z), they apply the reparameterization trick (e,0(x)) ( @c with I). This allows the expectation of p(y|z) to be determined by MC sampling using samples, described in Kingma & Welling (2022).
3.3 Reproducing results
The suggested two-stage training methods are notable because they can be applied to pretrained networks. As such, they neither strict post-hoe nor online-learning approach like data augmentation and potentially be combined with other calibration methods. Another advantage is that they impose less additional computational cost than training an ensemble. Replacing the last layer with MLP is applicable to a wide range of architectures, and during the second stage of training, models converge quile fast, as Jordahn Olmos (2024) pointed out. Furthermore, their method outperformed Temperature Scaling, which is another popular post-hoc calibration method. These results look promising, even though their tochnique was not derived theoretically. In fact, they discovered this method when conducting other experiments (Jordahn 2024). Therefore, further investigation of their observed phenomenon contribute to better
3

---

Published in Transactions on Machine Learning Research (06/2025)
understanding. While this paper focuses on the empirical side, we hope to encourage thorough theoretical analysis.
Using the code provided by the authors, we trained a WRN-28-10 architecture from scratch to reproduce the main results. We used the proposed methods TST and V-TST on the CIFAR10 and CIFAR100 datasets and included the evaluation on the shifted and out-of-distribution (OOD) data. We also reproduced two of the ablation studies in the original paper, The first ablation study evaluates the performance of two-stage param-paraeetrrized CNN architecture. Secondly, performed the ablation study regarding two-stage training for pre-trained models using ViT-B/16 (Dosovitskiy et al. 2021 on Tiny ImageNet (Le & Yang 2015a).
training
3.4 Our extensions
In addition to reproducing the original results, we introduced several extensions to deepen our understanding of TST and V-TST, and lo examine how various experimental setups influence their performance. While Jordahn & Olmos (2024) conducted a range of calibration studies, such as exploring the effects of end-to-end training, we aim to build on this by performing new ablation studics within a similar scope to gain new insights. For all experiments, we extend the results from Jordahn & Olmos (2024) by also including the Brier score to bring new insights. See Appendix top of ECE and MCE.
B for a detailed motivation including the Brier score on
First, we investigated the effectivencss of two-stage training using two additional CNN architectures: ResNet HH et al. 2015) and Efficient (Tan & Le. 2020). ResNet was chosen since it serves as canonical baseline for many subsequent CNN designs, making it a natural choice for comparison. Efficient was selected due to its superior performance on benchmarks like ImageNet, which makes it an appealing architecture for practical deployment Tan & Le (2020). Moreover, EfficientNet differs more substantially from WRN-28-10 by incorporating features such as compound scaling (Tan & Le. (020),, making it particularly interesting to assess the generalizability of two-stage training. Specifically, we used Nes.et (He al.) 2015) and EfficientNet-B5 (Tan & Le 2020), which are comparable in size to WRN-28-10 10 (36.5 million parameters). with 25.6 million and 30.4 million parameters, respectively.
Next, we investigated different MLP configurations in the TST and V-TST setting, and examined the effect of varying the number MC samples used in the second stage of V-TST. While the original paper used a three-layer MLP for the classification layer, we experimented with replacing it with a four-layer MLP and a single-layer MLP, the latter being equivalent to ro-initializing the final fully connected (FC) layer. Additionally, we explored whether increasing the number of MC samples in the second stage with V-TST from = 1 (original proposed approach) to = 5 could improve model calibration. These ablation studics were motivated by the observation that the choice of the latent dimension Z has a large effect on the calibration of the model and the effectiveness of two-stage training, as reported by Jordahn & Olmos (2024) in appendix F. For instance, using Z = 128 compared to Z = 8 in V-TST (m = 10) results in an ECE value of 1.270 instead of 12.88. We wanted to investigate whether the choice of other hyperparameters has a similarly strong impact on the performance of two-stage training.
Lastly, suggested in Jordahn & Olmos 2024), we investigated whether combining two-stage training with focal loss and other similar regularization methods can lead to further improvements in calibration perfor- mance. Jordahn & Olmos (2024) suggested that their method is complementary rather than an alternative to other implicit regularization techniques. For a thorough evaluation of their claim, we explored whether calibration improvements can be seen when combining two-stage training with two implicit regularization techniques (focal loss and label smoothing) and one explicit regularization method (L2 regularization). We leave the combination of their method with last-layer Laplace, OKO, and similar techniques for future studies.
4 Data
We used the same image datasets and preprocessing as the original paper, namely CIFAR10, CIFAR100 (Krizeessy et al. 2009) and Tiny TmageNet. (Le & Yangl 2015b) to train the models, and SVHN (Netzer
4

---

Published in Transactions on Machine Learning Research (06/2025)
et al. 2011 (OOD data) well CIFAR1O-C and CIFAR100-C (Hendrycks & Dietterich 2019 (shift data)
to eva palutte them. We randomly split the validation set of Tiny ImageNet into two halves to obtain validation and test sets, as the test set of the dataset is unlabeled. More details about the datasets and preprocessing steps can be found in Appendix A
For the ablation studies and the extensions, we used CIFAR10.
5 Experiments and Findings
Across all experiments, we evaluated the model performance using the metrics of the original paper, the expected calibration error (ECE) and the maximum calibration error (MCE) (Nacini et al. 2015). computed in percentage with 10 bins, on the regular and corrupted test sets. Furthermore, we measured negative log-likelihood (NLL) on both the training set and the test sct, and included metrics AUROC for the shiftod datasets and false-positive rale at 95% (FPR95) for the OOD datasels. In addition by the authors of the original paper, we included the Brier score, which is the mean squared error (MSE) applied to predicted probabilities, to get more comprehensive understanding. As in the original paper, we used the softmax outputs of the trained models as inputs for computing the evaluation metrics. Note, that treating the softmax output of networks as their confidence has been criticized (Gal & Ghahramani 2016). The softmax operation squeezes the prediction probability towards extreme values such as zero or one (He ct al. 2025) which could cause the observed miscalibration (Gawlikowski et al. 2023 Vasudevan et al. 2019: Hendrycks & Gimpel 2018). However, as our main objective is to reproduce the resulls of Jordahn & Olmos 2004)) we decided to follow their experimental settings closely.
the metrics used
In all experiments, the models were trained using the Adam optimizer with learning rate 10-4 and cross- entropy (CE) loss to align with the original paper, except in the extensions where we used focal loss (FL). In the first stage, the models were trained from scratch for 600 epochs with seed 1 unless otherwise specified. Throughout the paper, we refer to models after the first stage of training as base models. The TST and V- TST was done over 10 seeds (0-9) using 40 epochs and by re-initializing the final FC layer with a three-layer MLP head having dimensions [output, 32), [3Z, 2] and [Z, #classes), where the output corresponds to the dimension of the model architecture's output before the classifier. For each experiment, we followed the approach of the original code and stored the model with with the lowest validation loss. A single NVIDIA T4 16GB GPU was used for all experiments.
5.1 Replication of Main Results
The replication of the main results is reported in Tables 1 and 2 For TST and V-TST, we used the same latent dimensions Z as the original paper. The CIFARIO results for WRN-28-10 10 closely align with those reported in Jordahn & Olmos| (2024). Especially the NLL results are very similar and are mostly consistent with the Brier results. No single model performs best on all metrics, however, V-TST (m = 10) has the best overall performance. The numbers are generally slightly worse for almost all metries but lie within a close range to those found in the original paper. We believe that these slight differences can be attributed to the randomness involved in training the base model, leading to different base models to start from. While the authors specify running TST and V-TST over seeds 0-9, they do not specify the seed used for training the base model.
For CIFAR100, the results we obtained differ slightly from those reported by Jordahn & Olmos 2024).. First of all, we observed discrepancies in the base WRN-28-10 models. Despite trying eight different seeds, we were unable to replicate comparablc results. Our base model consistently showed lower accuracy, which, contrast to the original results, increased by a few percentage points during the second stage of training. A direct comparison of ECE and MCE for our and their base model is however difficult, as the bias of these metries depends on the accuracy (Minderer et al. 2021). Please refer to Appendix B for more details about our replication attempts and a short discussion of the dependence of ECE and MCE on the accuracy.
As the accuracy of our WRN-28-10 10 model trained on CIFAR100 differed by up to several percentage points, one should be careful when drawing conclusions about the ECE and MCE differences obtained. We considered the Brier score and NLL in addition and found that for our models, the Brier scores aligned with the NLL results. Similar to the findings in Jordahn & Olmos 2024), we observed improvements in both calibration and
5

---

Published in Transactions on Machine Learning Research (06/2025)
accuracy when applying TST and V-TSI, although the magnitude of these improvements differed compared to the ones of the original paper. For instance, V-TST (m 10) improved the ECE by 60.85%, whereas the original paper reported 73.0% improvement. A similar trend was observed for TST. Taking into account all metrics, V-TST (m = 10) secms best calibrated.
Consistent with Jordahn & Olmos (2024), we showed that V-TST and TST improve ECE and MCE across both datasets (see Table D Under distribution shifts and OOD settings, we similarly found that TST and V-TST significantly improve the calibration of the WRN-28-10 18-10 and outperform temperature scaling for CIFAR1O and CIFAR100.
Table 1: Replication of main results. Test metries in-distribution test datasets. For V-TST, and 10 indicate using 1 and 10 samples in the MC approximation, respectively. Temp. WRN is the temperature-scaled WRN. For CIFAR10, we used Z = 128 for V-TST (m = 10) and TST, and Z 32 for V-TST (m = 1). For CIFAR100, Z = 512 was used for V-TST (m = 10) and TST, and Z = 128 for V-TST (m = 1). Bold indicates the best-performing model.
Dataset
Model
Accuracy
ECE
MCE
Train NLL
Test NLL
Brier
V-TST, m
92.44=0.02
1.38::0.100
21.26==6.44
0.0655++0.0004
0.2579+4000117
0.0114::.0000
V-TST,
90.9++0.03
1.78+00148
24.80+8.981
0.1116+0.0012
0.3732++0.0041
0.0145=0.00007
CIFARIO
TST
92.55:10.01
3.341000016
24.68
0.0532:1000002
0.251611.0007
0.0115=0000001
Temp. WRN
92.52
4.70
25.2
0.0616
0.3024
0.0122
WRN
92.52
6.06
36.28
0.1031
0.5124
0.0133
V-TST, m 68.29++01
9.33++0.277
19.75+00.611
0.140+0.0022
1.3551
0.0045 0.0000
V-TST,
60.30=00000 05
00.6010.391
21.12:0.88 820
0.187:100027
L5510:10.0108
0.0048=0.0000
CIFAR100
TST
66.91=0.03
9.96::0002
18.81=0.305
0.1921:0.0005
L3530=10.0014
0.006610.00005
Temp. WRN WRN
62.63
12.07
22.42
0.3272
1.5512
0.0052
62.63
23.83
46.73
0.3309
2.1619
0.0059
V-TST, m 10
92.58 12.58+0..02
1.27-+0.
28.48-99.006
LL0671++0.0044
0,2617++0.011
V-TST, m 1
91 01.09=00006
1.52:+0.144
14.77: 1.781
0.1125:40.001
0.3732-+0.0043
CIFARIO original
TST
92.59:t0.03 0.18:
14.89:22.248
10573+000002
0.2482=0.004
Temp. WRN WRN
92.53
444
26.77
0.063
0,2965
92.53
5.88
35.84
0.1038
0,4944
V-TST,
10
71.93:10.04
5.83:10.143
14.03:10.59
0.078:0.001
1.202:10.0052
TST m
69.18++0.09
0.34+0.291
16.85++0.694
0.1 0.1094+0.0000
1.4815++0.0100
CIFAR100 original
TST
71.26+0..6
7.07++0.084
16.744+1.943
1.0887+0.000
1.1331 +0.012
Temp. WRN WRN
71.55
15.24
33.21
0.0739
L.3708
71.56
21.74
82.08
0.1161
2,2588
5.1.1 Replication of CNN ablation study
To reproduce Table 3, we implemented a simple CNN, following the architecture detailed in the Appendix of the original paper. After training the base model, we applicd TST and V-TST and matched the latent dimension with their results. Consistent with Jordahn & Olmos (2024), we found that two-stage training an under- dardttraaaaeteeizeee model worsens ECE but improves MCE, while TST can improve both as shown in Table 17 Appendix Eb. The lower proportion of frozen to trainable parameters in the second stage of an under-parameterizes model increases flexibility, possibly explaining the ineffectiveness of two-stage training. Additionally, we observed that for TST, the authors report improved ECF. with Z = 32, the only latent dimension where ECE improved. With Z = 128, we found the ECE to be worse than the base model.
5.1.2 Replication of ViT ablation study
The other ablation study concerned applying the two-stage training after fino-tuning a pre-trained VIT. Implementation details are described in Appendix Again, our results were generally consistent. with their findings. It can be seen from Table 18 (Appendix E that using V-TST can improve the ECE and MCE while TST worsens both ECE and MCE. They stress the good regularization propertics of V-TST and also conclude that training feature extraction layers and classification layers together from scratch is important for the two-stage training to be effective.
6

---

Published in Transactions on Machine Learning Research (06/2025)
Table 2: Replication of main results. Evaluation metrics on shifted data and OOD data.
Dataset
Model
SHIFT ECE
SHIFT MCE
OOD AUROC OOD FPR95
V-TST, m 10
9.41=0.3
20.33=00666
0.851+0.003
0.702=0.002
V-TST, m
12.53:10.29
20.19:10.56
0.766 0.003
0.008
CIFARIO
TST
13.09:0.04
29 20.32:10.10
0.88340.000 0.00
0.688=0.003
Temp. WRN WRN
16.44
37.93
0.885
0.698
20.06
45.62
0.874
0.708
V-TST, 10
18.99=0.25
32.52+ 0.46
0.784+0.00
0.783=0.003
V-TST, m
20.8110.44
35 0.
0.770.10.001
0.803-0.003
CIFARIOO
TST
21.13:10.76
30.82:1..13
0.692-10.002
0.846=0.004
Temp. WRN WRN
23.67
37.71
0.744
0.840
39.41
56.06
0.722
0.854
V-TST, m 10
10.41 10.41-0.22
22 02.922+0.57
0.821 0+0000
0.751=0.002
V-TST, m
12.89.10.3
21.01 10.58
0.722.10 0.0005
0.006
CIFARIO original
TST
11.62:00775
25.39 1.74
0.874+0.000 000
0.699=0.009
Temp. WRN WRN
16.45
36.77
0.872
0.746
20.27
45,06
0.891
0.653
V-TST, V-TST, M2
10
14.3::0.23
26.96:10.48
0.791
0.809=0.004
16.810.4
30.72.10.63
0.783.10.004
0.006
CIFARIO0 original
TST
17.44+0.00
28 28.11
0.823+0.000
0.764=0.006
Temp. WRN WRN
29.58
47.67
0.778
0.816
40.47
60.38
0.785
0,892
5.2 Ablation Study: Dependence on Network Architecture
As part our ablation studies, we investigated whether TST and V-TST generalize to other CNN architectures,
specifically ResNet-50 (He et al. 2015) and EfficientNet-B5 (Tan & Le 2020) (sce Section 3.4 for the motivation behind these choices). To accommodate the smaller 32 x 32 input size of CIFARIO, we adapted
the architectures of ResNet-50 and EfficientNet-BB (see Appendix C for details). All models were trained from scratch before applying TST and V-TST.
Both TST and V-TST improved calibration for ResNet-50 and EfficientNet-B5 (see Table 3 The best- performing model for ResNet-50 WRN-28-10. V-TST (m 10) improved ECE and MCE by 88.14% and 62.17%, respectively. This model also reached the highest accuracy of 89.80%. For EfficientNet-Ba the best calibration was obtained for tem- perature scaling, which improved ECE and MCE by 90.74% and 92.50% respectively. The best-performing two-stage procedure for EfficientNet-Bi. was V-TST (m = 1), for which ECE and MCE was improved by 40.93% and 81.42%.
obtained with V-TST (m = 10), which aligns with the results for
Calibration also improved when using TST and V-TST on shifted data. For ResNet-50, the best performance obtained for V-TST (m 10), while temperature scaling outperforms two-stage training for EfficientNet- B5. For both architectures, the model with the lowest Brier score was also the best-performing model across all of the calibration metrics for in-distribution and shifted data. However, this tendency did not hold for OOD data. Furthermore, the OOD data detection performance decreased for ResNet-50 when applying two-stage training.
The similarity in performance trends between WRN-28-10 and ResNet-50 may stem from their architectural resemblance, as WRN builds upon the ResNet design by increasing the width of residual blocks while reducing depth (Zagoruyko & Komodakis 2017 In contrast, EfficientNet differs more from WRN-28-100 10, by using features like compound scaling (Tan & Le 2020). These architectural differences may explain why the results achieved with WRN-28-10 did not directly translate to EfficientNet. Exploring alternative latent dimensions could further enhance the performance of TST and V-TST for models like Efficient
In summary, TST and V-TST seem to improve calibration across multiple CNN architectures. However, the extent of improvement appears to be architecture-dependen.. For ResNet-50, that architecturally closely relates to WRN-28-100 V-TST remains the most effective approach. In contrast, for EfficientNet-BB5 whose
7

---

Published in Transactions on Machine Learning Research (06/2025)
architecture diverge more significantly, simpler post-hoc methods such as temperature scaling outperform two-stage training. This highlights the importance of considering architectural characteristios when selecting calibration strategies.
Table 3: Dependence of TST and V-TST network architecture. The latent dimensions used for TST, V-TST (m = 1), and V-TST (m 10) were the same as those used for WRN-28-10, namely Z = 128, Z = 32 and Z = 128, respectively.
Model
Accuracy
ECE
MCE
Brier
SHIFT ECE
SHIFT MCE
OOD AUROC
OOD FPR95
V-IST,
89.8010.03
0.95=0.081
13.15-2.00
0,0154110
6.3611.24
0.77610.001
V-IST, m=1
86.77:00099
16010.075
30.52:6.821
0.020 0001
10.95:111:10
15. 05.40100.63
0.719100001
0.85210
TST
4.19+0.016
18.25+14677
0.005-00.0003
13.38+0.06
25.94+0.10
0284++000001
0.763+00004
Temp. ResNer-50
89.74
5.82
22.00
0.0164
16.80
33.82
0.861
0.744
ResNet-50
89.74
8.01
34.75
0.0236
21.58
44.94
0.832
0,782
V-TT, m 10
76.9310.004
1.9740.207
16, 10.400::1990
n.0827 10002750000005 00005
13.90.10.34
20.73:1..55
0.68310.001
0. 0.905.0 .011
V-TST,
71.87:00008
1.85:00.270
H.99116615
0.0364100 0001
14.54:10266
19.91:10.36
0.6631000001
0,906.00000 0012
TST
77.10.0..04
7.201:0225
43.6828.555
0.003010.00004
16.54.00.30
25.43.0047
0.0891100001
0.900.10.001
Temp. Net-Be
76.48
0.76
6.05
0.0502
8.95
12.80
0.681
0.920
Net-e
76.48
8:21
90.68
0.8315
18.58
28:35
0.678
0.917
5.3 Ablation Study: Second Stage Network Architecture
In order to investigate how the complexity of the MLP head architecture in the second stage of TST and V-TST influences the model's calibration, conducted experiments by varying the original thyee-layer MLP head to single-layer and four-layer MLP, while keeping the latent dimensions consistent with the replication experiments in section 5.1 across all setups. For further details, we refer to Appendix C
The results in Tables 4 and show that TST and V-TST with a four-layer MLP improved ECE and MCE with and without distribution shifts. While the performance is close to that of the three-layer MLP, is slightly worse overall when compared to all metrics. For TST, the lower train NLL and higher test NLL show signs of more overfitting in the four-layer configuration. We believe that the additional layer without optimizing the dimension of Z introduces too many trainable parameters. We also noticed that the models converged faster with four layers. Furthermore, the results of TST using a four-layered MLP but with a decreased latent dimension of 2 = 10 show comparable results to the ones with a three layer MLP and latent dimension of Z = 128, highlighting the potential importance of balancing number of layers in the MLP and the latent dimension.
In the single-layer configuration, the results show that re-training the final FC layer in TST improves the MCE by 49.37% and the overall performance across the shift and OOD metrics, still outperforming temperature scaling in most metrics. This indicates that simply re-training the FC layer can effectively enhance model calibration and robustness under distribution shifts. However, the single-layer MLP performs worse compared to its multi-layer counterparts in terms of ECE and MCE on in-distribution data, likely due to the reduced capacity to capture complex calibration patterns. It is interesting to see that even though the results for this model seem consistently promising across NLL, Brier score, OOD data, and SHIFT ECE, this is not reflocted in its ECE, MCE, and SHIFT MCE.
Table 4: Results for size of the MLP in the second stage of trainingon in-distribution data. The value indicates the number of layers in the MLP head of TST and V-TST. We used Z 128 as latent dimension unless stated otherwise.
Model
Accuracy
ECE
MCE
Train NLL
Test NLL
Brier
V-TST, 1= 4, 10
92.50+00.0
1.94+
26.08+5.912
0.0733+0.0006
0.2823+0.0029
0,0118+0.00005
V-TST, 1= m1
91.65++0.05
1.714
27.94+8.612
0.1227=+0.011
0.3971 +000026
0.0141++0.00W
TST. 1
92.5210.00 0.02
3.50-10280
23.36:12.717
0.0547:10.0002
0.2588
0.0115:10.00 00000
TST. 4, 2 2-10
91.6 01.64110.11
1.4
14.7410.997
0.0702-10.000 0.001
0.3052-10.001
0.0125:10.000
TST. 1=
92.52-00.02
2.69
18.3742.763
0.0526110.00001
0.2408:10.0016
0.0113:10.0000
8

---

Published in Transactions on Machine Learning Research (06/2025)
Table 5: Results for size of the MLP in the second stage of training on shifted and OOD data. The value indicates the number of layers in the MLP head of TST and V-TST. We used Z 128 as latent dimension unless stated otherwise.
Model
SHIFT ECE
SHIFT MCE
OOD AUROC OOD FPR95
V-TST, [ = 4, m = 10
11.
28.16:00999
0.802:10.000
0.695=0.000
V-TST, 1= 1,4,11 m1
12.47=0.29
17.07=0.47
0.71=00.004
0.773=0.004
TST, -
13.45=#0.04
30.25:+0.09
0.875:10.002
0.727=0.013
TST, (= 1 10
11.69+0.43
24.54+0.90
0.883=00.001
0.688=0.003
TST, 1-1
11.46=0.27
25.35:10.65
0.884=0.001
0.682=00002
5.4 Ablation Study: V-TST Dependence on Number of Samples Training
To test whether using additional samples n during V-TST improves the model calibration, we re-trained the base WRN-28-10 for CIFAR10 using V-TST with = 5 MC samples and compared it to the original setup with n = 1. As shown in Table G using n = 5 results in slightly better MCE and OOD metrics, while the overall calibration performance remains comparable or slightly worse. Again, the Brier score was consistent with overall performance and aligned with the ECE and MCE calibration metrics. In contrast to our expectations, increasing the number of samples drawn during training did not prove to be beneficial for improving the calibration. However, we suspect that further exploration of a broader range of values would provide deeper insights to draw better conclusions. As expected, also observed reduced variability in most evaluation metrics with n n= 5.
Table 6: Effect of number of samples used in the second stage training with V-TST. Evaluation metrics on in-distribution, shifted data, and OOD data. The value n denotes the number of MC samples used during training, while m represents the number of samples used during testing. We used 128 as latent dimension.
Model
Accuracy
BCE
MCE
Brier
SHIFT
SHIFT MCE
OOD AUROC
OOD FPR95
V-TST = 5, m 10
92.45-10.02
1.58:0.081
1.999
0.011410.00002
9.65::021
20.56=0-18
0.862:0.002
0.697
V-TST,
91.04+00066
1.62+0.103
18. 68.8976.773
0.0143-0.0009
12.57+0.22
20.08-0.39
0.702=00004
0.779++0.004
V-TST, 1, 10
92.44+0.02
1.38=01100
21.26 3..2666.44
0.0114+0.00004
9.41=00.3
20.33-0.66
0.861=0.003
0.702=0.002
V.TST. 1.
90.9410.3
1.7820.14
24.8018.981
0.014500.00007
12.50:00.29
20.10:10.56
0.70620.003
0.778
5.5 Ablation Study:
Focal Loss
We investigated how two-stage training compares to the focal loss, and whether combining the two implicit calibration methods could lead to improved performance. To Lest this, we experimented with using three configurations: using focal loss in both training stages, the first training stage only or the second training stage only. For better ovaluation, we experimented with two different variants of focal loss: using a constant value for 7 (FL) and using an adaptive, sample-dependent schedule (FLA) as described in Mukhoti et al.) (2020). In accordance with Mukhoti et al. (2020), we used 7 for the constant loss and sample-lependent scheme that uses 7 3 for By E [0.2, I. and 7 5 P. [0, where Py denotes the probability of the network to predict the correct label y.
Focal loss base models First, observed that base model trained on constant focal loss (FL-B) adaptive focal loss (FLA-B) is more calibrated on in-distribution and shifted data than a base model trained using cross-entropy loss (CE-B), demonstrating the expected calibration benefits of focal loss. However, V-TST with m = 10 and CE loss (V10-CE2) outperforms base models trained using focal loss (FL-B and FLA-B) in terms of both ECE and MCE scores as well as accuracy. The Brier scores of the focal loss base models were lower than the score of the cross-entropy loss base model, consistent with the overall calibration performance of these three models.
9

---

Published in Transactions on Machine Learning Research (06/2025)
Focal loss in two-stage training When combining the two methods we noticed two key observations. Firstly, using focal loss as training objective for both stages in two-stage training significantly worsens performance on several calibration metrics such as ECE, NLL, Brier score and (mostly) MCE. We noticed a similar trend when comparing two-stage focal loss models to the reproduction results obtained in Table Hence, fully combining two-stage training with focal loss does generally not lead to improved calibration compared to base model trained with focal loss or two-stage training using CE loss. However, secondly, we observed significant improvement of calibration on shift and OOD data. As can be seen in Figure 2 the focal loss base models become overconfident after the second stage using focal loss, however, the plots suggest an overall tendency of models to be less confident on shift data. This is in contrast to other observations, where models become overconfident (Minderer et al. 2021 but explains the improved calibration of the overconfident model on the shift data.
worsenentropp
Focal loss base and cross-cntropy in sccond-stage Using focal loss in the sccond stage on top of a focal loss base model worsens the calibration overall. Therefore, we tested whether combining better- calibrated base model (compared to base model trained with CE loss, see Table ( with second stage that uses CE loss would enhance calibration performance. We found that this combination (VI0-FLA-CE, VI-FLA-CE, FLA-CE) generally improved both ECE and MCE significantly compared to the base model that used adaptive focal loss (FLA--) and the corresponding CE. baseline models (V10-CE2, V1-CE2, CE2) while maintaining a similar accuracy. For instance, including FLA in the first stage for TST (FLA-CE) improved ECE and MCE by 37.1% and 27.4%, respectivcly compared to the CE2 model (sce Table ..hile the best. ECE score was achieved for V-TST (m 10), this was not consistent with the best Lest NLI. Brier score, and OOD performance. Results for using FT--B base model and performing second stage training with CE were similar to those using FLA-B as base model and are thus reported in the Appendix D
While
Table 7: Results for including Focal Loss (FL) and Adaptive Focal Loss (FLA) in the training procedure. Test metrics and evaluation metrics on in-distribution data. The first loss before + sign indicates the loss used for training the base model prior to TST and V-TST. We used a latent dimension of Z = 128. The asterisks (*) indicate the best value across the table.
Model
Loss
Model Name
Accuracy
BCB
MCE
Train NLL
Test NLL
Brier
V-IST,
10
VI0-P12
91.77+0.03
17480.148
20.94++0.934
0.2016+0.00188
0.3694-3 D016
0.01443-0.0000
V-TST, m-
3 FL.
V1-F12
87.71+0.11
11.08+0.482
24.35+1.880
0.31 0,3128+00.002
0.5185+00.005
0.054
0.0213-0000023
TST
3 FL
FL2
91.5510.03
5.9020.384
29.2418.339
1..06150.0043
0.27962:00008
0.0150=0.00011
Temp. WRN
3 FL
FL-TS
91.78
1.59
10.25*
0.0880
0.2524
0.0123
WRN
3 FL
FL-B
91.78
1.57
27.45
0.0688
0.2604
0.0123
V-IST, m 10
5-3 FLA + CB
VI0-FLA-CE
92.04:0.01
0.80=
20.33+1.998 1.81
8.0773+0.0004
0,2861 0008
0.0122200000073
V-TST, 10
1-3 FLA CE
VI-FLA-CE
90.11+00.4
1,52+0.129
20. 20.42+66.533
0,1351
+0.0012
0.4101=00.0035
0.0158800000017
TST
5-3 FLA
FLA-CE
12.10.00:3
2::00.0022
14.92:22.22
0.0095:10.0001
0.2529=000004
0.0119=0.00000
V-TST,
10
5:3 FLA
V10-FLA2 92.02:0.03
11 12:00.01151
L,25
0.1964:0..00 0016
0.3092::000000 00112
0.0141=40.0000
V-TST,
5-3 FLA
VL-FLA2
87.7310.08
12.0810.327
23.7510..339
0.318110.00 0014
0.5259110.005
0.0170 00000
IST Temp. WRN WRN
5-3 FLA
FLAZ
02.6::0001
5. 5.96
27.78:66. 487
0.1190
0 0.2736:2000 030
0.0126::0.00
5-3 FLA
FLA-TS
92.14
1.42
81.25
0.077
0.242*
0.0119
5-3 FLA
FLA-D
92.14
1.97
22.83
0.0602
0.2539
0.0120
V-TST, 10
CD FL
VIGCEEFL
92.40:0.02
9.84:00235
26. 01.922660064
0.06
0.0016:000022
0.0027=0.00008
V-TST, m= 1 IST
CE + FL
VI-CE-FL
89.2610008
10. (0.65:10.400 108
25.35 224
0.26671000008
0.17081:00034
4.0156110.00000
CE + 3LL
CE-FL
92.54:0.02
0.21:0.255
27.5001-4.499
0.0822:0.0021
0.0112:0.002300 0.2023
V-IST, In 10 V-TST,
CB
V10-CE2
02.44:00002
1.88:0.100 1.00
21.26::6.44
8.0655
0,2679 0017
0.0114::00 00000000000044
CE
V1-CE2
90.9+0.03
1,78+001148
24. 14.80+-9.000
0.1116+00.00 16+0000012
0.3732+-0.064
0.01 4.0145HHLD0000
IST
CE
CE2
92.55=00.01
3.34:40016
24 68--6883..8 883
0.0532:0.00000
0.000*
0.2516=000007
0.0113:20.0000
Temp. WRN
CE
CE/TS
92.12
4.7
25.2
0.0616
0.3024
0.0122
WRN
CE
CE.B 92.82
6.06
36.28
0.1001
0.5124
0.0133
Cross-entropy loss base and focal loss in second-stage For completeness, we also tested how using constant focal loss in the second stage top of CE base model would perform (V10-CE-FL, VI-CE-FL, CE-FL) and found a similar behavior to using constant focal loss or adaptive focal loss in both stages of training (V10-FL2, V1-FL2, FL2, V10-FLA2, V1-FLA2, FLA2). Both ECE and MCE scores were higher than those obtained for CE-only two-stage training. An exception to this was the TST model, which reached a similar performance than the CE2 model. Here, Brier Score, NLL, OOD performance and ECE were aligned. Taken the results for these models together with the ones using focal loss in both stages together, we find the overall tendency that using focal loss in the second stage of training is not a good choice.
10

---

Published in Transactions on Machine Learning Research (06/2025)
We hypothesize that the worsened calibration may be due to the way focal loss interacts with the fixed feature representations established in the first stage. Two-stage training relies retraining the classification head to express confidence that aligms with the structure of the latent feature space. However, since focal loss penalizes all output logits to a low level (Tao et al. 2023b). it may be problematic when the feature extractor has produced well-separated clusters, as it prevents the classifier from confidently predicting clear cases. However, this should be explored and tested further.
Table 8: Evaluation metrics shifted and OOD data for the effect of inchuding Focal Loss (FL) and Adaptive Focal Loss (FLA) in the training procedure. The first loss before sign indicates the loss used for training the base model prior to TST and V-TST. We used a latent dimension of Z = 128. The asterisks (*) indicate the best value across the table.
Model
Loss
Model Name
Shift: ECE
Shift MCE
OOD AUROC
OOD FPR95
V-TST, m 10
3 FI.
V10-FL2
3.92:+0.39
6.07
0.799
0.768=0.003
V-TST, m=
3 FL
V1-FL2
2.96=2.15
4.9740.31
0.701 0000
0.850=0.003
TST
3 FL
FL2
2.44
4.75.0..29
0.87610.001
1672_00000
Temp. WRN
3 FL
FL-TS
6,28
11.89
0.875
0.701
WRN
3 FL.
FL-B
11.93
22.21
0.866
0.740
V-TST, m
3-3 FLA + CE
VI0-FLA-CE
10101=0.16
19.66
0.825+0.002
0. 0.727
V-TST, m
5-3 FLA + CE
VI-FLA-CE
12.65+0.32
19.19=:0.56
0.762+0.002
D. 0.790=0.003
TST
5-3 FLA + CE
FLA-CE
11.75+0.07
23.36+0.16
0.877=0.001
0.676=0.002
V-TST,
10
5-3 FLA
V10-FLA2
4.1310.36
6.7940044
0.828110.00 2000
0.714=0.004
V-TST,
5-3 FLA
VIFLAA
2.50+0.15
0.28+0.29
0.735+00.003
0.827-
TST
5-3 FLA
FLA2
2.75++0.06
5.53++0.22
0.887+00.007
0.625+0.000 0.004
Temp. WRN WRN
5-3 FLA
FLA-TS
7.06
13.36
0.88
0.65
3-3 FLA
FLA-B
12.74
24.30
0.88
0.69
V-TST, m. 10
CE + FL
VIO-CE-FL
3.9210.30
6.07.10.42
0.799:10.006 004
U. 0.768=0.003
V-TST, M =1 TST
CE + 3 FL.
VI-CEFI
2.25+0.27*
4.20++0.11
0.745++0.004
0.810-0.005
CE + 3 FL
CE-FI
2.44++0.10
4.75++0.29
0.876+0.001
0.672
V-TST, 10
CE
V10-CE2
9.41=0.3
20.33+0.66
0.851#0.000
0.00220.002
V-TST,
CE
VI-CE2
12.53+0.29
00.19+0.556
0.766+00.03
0.778=0.003
TST
CE2
13.09:10.04
29.32:10100
0,88310.001
0.688=0.003
Temp. WRN
CF
CF-TS
16.44
37.93
0.885
0.698
WRN
CF
CE-B
20.06
45.62
0.874
0.708
Focal loss performance on shifted data Incorporating focal loss in Iwo-stage training any form leads to significantly lower ECE and MCE for shifted data compared focal loss base models and two-stage CE models (see Table D The only exceptions to this are the V-TST models that used adaptive focal loss in the first stage and CE loss in the second stage (VI0-FLA-CE, VI-FLA-CE) but even these perform similarly well compared to V-TST models that only use CF loss (V10-CE2, V1-CE2). Another interesting observation was that models trained in a two-stage fashion with focal loss in both or one of the stages perform better on shifted data than in-distribution data. The exception for this trend were models trained with adaptive focal loss in the first stage and CE loss in the second stage (V10-FLA-CE, V1-FLA-CE, FLA-CE). This is rather surprising finding as the WRN-28-10 trained with CE loss (V10-CE2, VI-CE2, CE2) was better calibrated for in-distribution than shifted data for both CIFARI0 and CIFAR100 (see Table
A possible explanation for the better performance shifted data is that focal loss emphasizes hard exam- ples during training (Lin et al. 2017). leading to better calibration on challenging shifted data but poorer calibration on easier in-distribution data. Additionally, focal loss with a fixed y has different dynamics than the CF loss. Initially, models trained with focal loss show higher weight norms than those trained with CE loss. However, this trend reverses after further training so that the weight norms of the focal loss model become lower, which can be explained by the regularizing effect the focal loss has on the network weights Mukhoti et al. 2020
Focal loss performance on OOD data For the OOD data, we find two interesting results. First, TST generally performs best. both in terms of AUROC and FPR95, regardless of the loss function combination
11

---

Published in Transactions on Machine Learning Research (06/2025)
(even the CE baseline model (CE-B) has the lowest FPR95 score for the TST). This shows the superiority of TST on OOD data. In contrast, V-TST (m 1) performs worst in all configurations for both AUROC and FPR95. Second, among the TST models, using adaptive focal loss in both stages of training (FLA2) reaches the highost AUROC and lowest FPR95 rate overall, outperforming both the CE bascline and the other local loss combinations. Thirdly, the combination of adaptive local loss in the first stage and CE loss the second stage (V10-FLA-CE, V1-FLA-CE,, FLA-CE) mostly worsens the performance (both AUROC and FPR95) compared to using the adaptive focal loss in the first and second stages of training (V10-FLA2, V1-FLA2, FLA2) or the CE baseline (CE-B).
1.8
0.6
Confidence
Confidence
(a) CE Loss In-Distribution
(b) CE Loss SHIFT
Perfect Call bration
Perfect Calibration
2.6
0.6
0.5
0.
(c) Focal Loss In-Distributtin
(d) Focal Loss SHIFT
Figure 2: ECE plots for models with CE and focal loss. (a) In-distribution ECE plot of a V-TST trained
model (Z = 128, m 10, V10-CE2) and the base WRN-28-10 trained CIFAR10 with CF loss (CE-B). (b) Same models with CE loss evaluated shifted data CIFAR10-C. (c) In-distribution ECE plot of V-TST trained model (Z = 32, 1, V1-FLA-B) and the base WRN-28-10 10 trained on CIFAR10 with focal loss (FLA-B). (d) Same models with focal loss evaluated on shifted data CIFARIO-C.
Reliability diagrams Figure 2 shows the ECE plots for the base and V-TST models trained with CE loss and the 5-3 focal loss both stages. Starting from an uncalibrated WRN-28-10 model trained with CE loss in 2a (CE-B), V-TST in the second stage leads to an improved calibration (V10-CE2). 2c shows that the base WRN-28-10 10 that was trained with adaptive focal loss (FLA-B) is already better calibrated than the base WRN-28-10 model trained with CE (CF-B) in 2a The second training stage with adaptive focal loss makes the model underconfident in its predictions but also evens out the overconfidence it showed for images that were around 35% certain. For adaptive focal loss on the shifted dataset (2d], both the base WRN-288100 (FLA-B) and V-TST (VI-FLA-B) models perform worse than on in-distribution data, with a decrease in accuracy. However, for the V-TST model (VI-FLA-B). this results in confidence aligning more closcly with its accuracy, whereas the base model becomes overconfident.
5.6 Ablation Study: Label Smoothing
Following the results from combining focal loss and two-stage training, we investigated whether label smooth- ing, another implicit regularization technique, be similarly combined with two-stage training to improve
12

---

Published in Transactions on Machine Learning Research (06/2025)
model calibration. In order to verify this, experimented with two configurations: using label smoothing in both training stages, and only in the second stage. We chose smoothing factor = 0.1, consistent with prior work on similar network architectures (Müller et al. 2019).
Table 9: Results for including Label Smoothing (LS) in the training procedure. Test metrics and evaluation metrics on in-distribution data. The method column specifies the training strategy. If single entry is applied consistently across both training
present, it indicates that the loss function or regularization stages. If two entries separated by +, the first denotes the configuration for the first stage and the second for the second stage. We used latent dimension of Z = 128. The asterisks (*) indicate the best value across the table.
Model
Method
Model Name
Accuracy
ECE
MCE
Train NLL
Test NLL
Brier
V-TT,
10
LS
V10-L82
92.19::0011
8:0020 100
21.8816.251
0.183820 0108
0,379F:20.00 0009
0.013220.0UXX2
VTST,
LS
VI-LS2
92.37.1.002
8.2610.134
26.2115.897
0,205410.00I
0.113410.0020
0.0.4611000008
TST
IS
182
92.50+00.00
6.01 176
19.57+40.767
0.1358:0.000
0.0010
0.3321=4000013
0.0121=0.00002
Temp, WRN
LS
LS-TS
92.22
9.04
26.93
0.1891
0.3796
0.0133
WRN
1S
LS-B
92.22
6.71
27.09
0.1446
0.3523
0.0129
V-TST, 10
CE
LS
VI0-CE-IS
92.1740.00
6.98#0.155
21.1842.947
0.1505=000018
0:142556.0016
0.812740.00003
V-TST, m
CR
+ 18
V1-CR-IS
91.61+00.05
5,80+001700
24.64444.196
0.1625+00.0019
0.3836+0.0015
0.01 4.011334000000
TST
CB
CE-LS
92.21:0.02
3.42:10.201
14.60=1.890*
0.1046:0.0000
0.2919:0.0044
0.0119=0.00004
V-TST, IN = 10
CE
V16-CE2
92.44:00002
1.38=0.1100*
21. 6.426
0.06655:20 0000
0. 0.2679:20.00) 0017
0.0114::0.000
V-TST,
CE
V1-CE2
90.9:10:03
1.78200118
21,80
0,1116100 0012
0,973210:001
0.0145510:00007
TST
CE
CE2
92.65=:0.01
24.68-t33.88 883
0.0532=0.00022
0.2516=0.0000
0.01 15:00
Temp. WHN
CE
CE-T9
25.2
0.3024
WRN
CE
CE-B
02.52
6.06
36.28
0,1031
0.5124
0.0133
Table 10: Evaluation metrics on shifted and OOD data for the effect of including Label Smoothing (I.S) the training procedure. The method column specifies the training strategy. If single entry is present, it indicates that the loss function entries are separated by a +, the first. denotes the configuration for the first stage and the second for the second stage. We used a latent dimension of Z = 128. The asterisks (*) indicate the best value across the table.
regularization is applied consistently across both training stages. If two
Model
Method
Model Name
Shift: ECE
Shift MCE
OOD AUROC
OOD FPR95
V-TST, m = 10
LS
V10-LS2
5.0420.15
21.5210.2
0.719
0.752#0.002
V-TST, m =1
LS
VI-LS2
5.64=0.24
8.65:+0.45
0.646:00.006
0.811=0.004
TST
LS
LS2
10.67.10.16
26.64110.7
0.831 10.005
0.748 0.001
Temp. WRN
LS
LS-TS
6.91
31.62
0.611
0.763
WRN
T
LS-BB
11.06
34.93
0.612
0.763
V-IST, m = 10
CE + LS
V10-CE-LS VI-CE-LS
3.12#0.07*
11.1040.22
0.821+1.008
0.88800003
V-TST, m =
CE LS
4.19=20.1
5.88+0.28*
0.767+0.003
0.727=0.003
TST
CE LS
CE-LS
5.4910.22
16.1910.98
0.863.10.006
0.657
V-IST, m 10 V-TST, CF
CE
V10-CE2
9.41:0.33
20.3310.66
0.851 003
O. 0.702
VI-CE2
02.53+00.29
20.19+00.56
0.766++0 0.003
0. +778++0000
TST
CE
CE2
13.09
29.32:10.10
0.883:10.001
0.68810.003
Temp. WRN
CE
CE-TS
16.44
37.93
0.885*
0.698
WRN
CE
CE-B 20.06
45.62
0.874
0.708
The results in Table show that using LS alone slightly improves the calibration of the base WRN model on in-distribution data, particularly in terms of MCE, Test NLL, and Brier score. However, when combined with TS (LS-TS) or V-TST (V1-LS2, V10-LS), the calibration performance of the basc model (LS-B) decreases across nearly all metrics and is worse than the ones obtained using CE loss alone (CE-TS, V1-CE2, V10- CE2). While LS does benefit from TST (LS2), the improvements are still smaller compared to using TST with CE loss alone (CE2), except for MCE. Notably, including LS into the second stage of (V)-TST improves the MCE compared to the results obtained using CE loss alone (CE2, V10-CE2, V1-CE2), likely due to its tendency to reduce prediction overconfidence. For TST using T.S only in the second stage (CE-LS), the
13

---

Published in Transactions on Machine Learning Research (06/2025)
performance is comparable to using CE loss alone in both stages (CE2) across all metrics, with a significant improvement in MCF.
Under distribution shifts and OOD data, using LS in the second stage leads to better overall performance as shown in Table I0] When evaluated on shift data, using LS largely improves the ECE and MCE metrics compared to the results obtained with CE loss alone. Similarly to the in-distribution results, using LS only in the second stage leads to better performance than using it in both stages. In both setups, the V-TST models (V10-LS2, V1-LS2, V10-CE-LS, VI-CE-LS) perform best, with m 1 achieving the best MCE and m = 10 the best ECE scores. For OOD data, using LS in the second stage IV10-CE-IS VI-CE-LS, CE-LS) improves the FPR95 metric compared to the CE results (V10-CE2, V1-CE2, CE2) but leads to a slight drop in AUROC, whereas using LS in both stages worsens both metrics.
In summary, we found that integrating LS in two-stage training does not lead to further improvements in calibration on in-distribution data compared to the reproduced results. The exception applying LS only in the second stage of TST. Under distribution shifts and OOD data, using LS in the second stage leads to better overall performance.
5.7 Ablation Study: L2 regularization
Lastly, we invostigated whether a similar pattern can be found for combining two-stage training with an explicit regularization technique such as L2 regularization. Again, we experimented with using L2 regular- ization during both stages of training or only in the second stage of training. We chose À 0.003, based on the results of Joo & Chung 2020) who performed hyperparameter search for ResNet-50 on CIFAR10 to find the best L2 regularization strength.
The results show that the use of T.2 regularization leads to a better calibrated base model (I.2-B) for the in-distribution data in terms of ECE and MCE scores, but increased Brier score (see Table II. For the OOD data, the base model that used L2 regularization performs worse compared to the CE-B model. No clear trend can be seen for the base models in regards to shift data.
Results for including L2 regularization (L2) in the training procedure. Test metrics and evaluation metrics on in-distribution data. The method column specifies the training strategy. If single entry is present, it indicates that the loss function or regularization is applied consistently across both training stages. If two entries are separated by a +, the first denotes the configuration for the first stage and the second for the second stage. We used a latent dimension of Z = 128. The asterisks (*) indicate the best value across the table.
Table 11:
Model
Method
Model Name
Accuracy
ECE
MCE
Train NLL
Test NLL
Brier
V-TST,
10
L2
V1D-12-2
90.67=0.00
2.58=0.171
05.90=0.591
0.1127=0.0010
0.3578=0.0011
0.0146
V-TST, m =1
12
V1-1-222
90:47+0.05
3.53+0.222
18.2 18.29+1..0 878
0.1306:-0.003
0.4232+0.0046
0
L2-2
90.6620.01
1.3220.153
28.79:5.178
0.1052:0.00$$
0.359020.0011
0.00-5
Temp. WRN
L2
12-TS
90.16
2:24
19:27
0.1531
0.3688
0.0152
WRN
L2-B
5.00
0.1160
0.3865
0.0159
V-TTT m= 10 V-ST, m
CB + L2
V10-CE-L2
92.0920.01
1004
8.79
0.0660=0.0005
0,261720.0001
0.0118:00000007
CE
VL-CE-L2
91.2110.11
3.1110.160
14.98111.778
0.082220.0000
0.340910.00 0043
0.013710.00005
TST
CE L.2
CE-L2
92.26:0.02
2. 2.601:02200
19.62:47.328
0.0587=000000
0.2373 0.0024
0.01 100.1940000004
V-TTT in 10 V-TST, in
CE
V10-CE2
92.44:1000
1.3010.100
21.26:
0.0655:0.000:
0.2579
0.0114 0.00004
CR
V1-CF2
90.9+0003
1.78+-01488
24.80-+ 9.981
0.1116-0.0012
0.3732+0.0041
0.014550000007
CE
CE2
92.55-10.01
3.341:00016
24.68:11.888 883
0.0532:0.0002+
0.2516:40.0007*
0.0115
Temp. WRN
CR
CE/TS
02.52
4.7
25.2
0.0616
0.3024
0.0122
WRN
CE
CE-B
92.52
6.06
36.28
0.1031
0.5124
0.0133
Applying L2 regularization in both stages results in worse calibration on in-distribution data for TST (L2- 2), compared to the CE baseline (CE2). Similarly, for V-TST (V10-L2-2, V1-L2-2), using L2 regularization worsens ECF, NL.,, and Brier scores compared to not using L2 regularization (V10-CF-2, V1-CE-2, CE-2, respectively). However, MCE scores improve for these V-TST models.
Using L2 in the second stage of training on top of the CE base model leads to improved calibration in terms of ECE and MCF scores compared to using T.2 in both stages and not using L.2 at all (see Table 11 The
14

---

Published in Transactions on Machine Learning Research (06/2025)
only exception is the V1-CE2 model, which produces lower ECE score compared to the V1-CE-L2 model. Moreover, using T.2 regularization only in the second stage also improves the Brier score compared to using L2 in both stages. Hence, using L2 regularization in the second stage tends to improve calibration, which is in contrast to the results obtained using focal loss.
Under distribution shift, using T.2 regularization in both stages worsens ECE and MCE compared to the results obtained without the explicit regularization, as can be seen in Table 12 However, using L2 regular- ization in the second stage only mostly improves ECE and MCE except for V-TST with 1, where it increases both (V1-CE-L2 compared to V1-CE2). The worst OOD detection performance is associated with using L.2 regularization in both stages, while the best one is achieved by using it only in the second stage.
Taken logether, applying L2 regularization in the second stage can considerably improve calibration compared to using it in both or none of the stages. These results extend from the in-distribution data to shift data and OOD detection. Interestingly, this finding is in contrast the results of the focal loss and, to a lesser extent, LS ablation studies, where using the implicit regularization technique tends to worsen the in-distribution calibration of V-TST compared to the CE baseline if it is applied in the second stage of training.
Table 12: Evaluation metrics on shifted and OOD data for the effect of including L2 regularization (L2) in the training procedure. The method column specifies the training strategy. a single entry is present, it indicates that the loss function or regularization is applied consistently across both training stages. If two entries are separated by a +, the first denotes the configuration for the first stage and the second for the second stage. We used a latent dimension of Z 128. The asterisks (*) indicate the best value across the table.
Model
Method Model Name
Shift ECE
Shift MCE
OOD AUROC
OOD FPR95
V-TST,
10
1.2
V10-L2-2
18.30+0.21
34.56+0.36
0.743=0.002
0.765++0.001
V-TST, m 1
1.2
V1-L22-2
20.50+0.31
03.39++0.6
0.713-0.001
0.780++0.009
TST
L2-2
20.5411.67
38 1.11
0.76-1-0.002
0.762:10.002
Temp. WRN WRN
L2-TS
14.68
26.23
0.825
0.725
L2-B
23.02
40.45
0.824
0.741
V-TST, m 10
CE + L2
V10-CE-L2
8.37=0.25*
17.21=0.42
0.887=0.001
0.637++0.002
V-TST, m=
CE + L2
VI-CE-12
12.92+00.3
25.42+0.67
0.862-0.001
0.687+0.003
TST
CE L2
CE-12
11.80:00.45
25.72=1.19
0.897=0.001
0.632:10.002
V-TST, m 10
CE
V10-CE2
0.4
20.33=0.66
0.851=0.003
0.702=+0.002
V-TST, m 1
CE
V1-CE2
12.53.10.29 .9
20.19
0.006-0.003
0.77811 003
TST
CF
CE2
13.09+00.04
29. 29.32=00.10
0.883=0.001
0.688++0.003
Temp. WRN
CF
CF-TS
16.44
37.93
0.885
0.698
WRN
CE
CE-B
20.06
45.62
0.874
0.708
6 Conclusion
In this paper, aimed to reproduce the main results from Jordahn & Olmos (024)) and extended their findings by performing numerous ablation studies to get new insights into two-stage training. We confirmed their findings that decoupling the training of feature and classification layers improves the calibration of a WRN-28-10 model on CIFAR10 and CIFAR100, with slightly varying effects on CIFAR100. Additionally, we successfully reproduced two of their ablation studies, the CNN ablation study and the ViT fine-tuning ablation study. From the additional ablation studies in this paper, we found that the effects of improved calibration from two-stage training transfer to other similar model architectures, namely ResNet-50 and EfficientNee-B5. However, the effectiveness of two-stage training differs across CNN architectures. For the second stage architecture, we observed the importance of balancing the number of layers in the MLP with the latent dimension Z used. Lastly, we found that using more MC samples during training does not affect the calibration performance of the models positively. Thus, overall, the performance was less affected by our ablation studies than the choice of Z, as could be seen from their additional results.
Finally, we followed the suggestion Jordahn & Olmos (2024) to combine their method with other calibration techniques. Specifically, explored the integration of their approach with two implicit. regularization
15

---

Published in Transactions on Machine Learning Research (06/2025)
methods, focal loss and label smoothing, and one explicit regularization method, L2 regularization. Our investigation led to four key takeaways in regards to in-distribution data. First, all regularization methods improve the calibration of the base WRN-28-10. Second, using additional implicit regularization in both stages of training tends to perform worse than the CE bascline or the corresponding base model before the second stage. Third, combining two-stage training with additional regularization can either enhance or degrade calibration, depending on the regularization and which stages it integrated. For instance, focal loss, when used only in the second stage of V-TST, worsens calibration of the corresponding base models, whereas L2 regularization tends to have the opposite effect. Lastly, applying focal loss only to the base model, followed by TST or V-TST using CE loss, resulted in better-calibrated models compared to the baseline. When looking at the OOD detection performance, we found a tendency for TST to perform better in OOD detection than V-TST and the corresponding base model across most experiments. On the other hand, V-TST generally achieved better performance under distribution shifts. Notably, TST also balances in-distribution and shift performance well when additional implicit regularization used in the socond stage of training.
To conclude, we found that TST and V-TST can further improve calibration when combined with another regularization method. However, both the choice of regularization method and the way it is combined with TST or V-TST significantly influence the method's effectiveness on calibration. We had to conduct several experiments to identify good combination, and our initial attempts even worsened calibration. heleth the approach remains promising, and our ablation studies extend the empirical results of Jordahn & Olmos (2024) A thorough theoretical analysis of these methods remains open for future work.
7 Future Work
Next to the theoretical analysis of (V-)TST, future extensions should also investigate possible explanations as why using focal loss for both stages of training decreased model calibration compared to the focal loss base model. One possible next step could be to apply early stopping based on ECE (as has been tested by Mukhoti et al. (2020 during the second stage of training to avoid over-regularization of the weights. Other idcas would be testing whether the obtained focal loss results are transfcrable to other model architectures beyond WRN-28-10, or further experimenting with the second stage latent dimension size and other datasets.
Additionally, we would be interested to see further work investigating how the TST and V-TST perform when combined with other implicit regularization techniques such as OKO (Muttent et al. 2023 or post-hoc calibration methods such laplace approximation, which had also been suggested in Jordahn & Olmos (2024). They claimed that their method is not alternative to other techniques but rather complementary to them. However, our results show that combining TST and V-TST with other calibration techniques is not always straightforward but requires careful experimental design.
8 Challenges and Limitations
During the study, we encountered some challenges in replicating the results from Jordahn & Olmos (2024) due to missing information and incomplete code, which required us to re-implement parts of it and make assumptions about the training hyperparameters. For instance, details on the MTP head design for the CNN in TST and V-TST, choices for latent dimension, training scods for the base WRNs, and information regard- ing the pre-trained VIT model for fine-tuning were omitted. Additionally, the model structures reported in Tables 1 and 2 were not immediately apparent. The variations in the latent dimension Z for each model could only be seen by inspecting and comparing all results in the Appendix in detail.
The theoretical justification jueotreiial,jusifffaaiiin for the proposed methods by Jordahn & Olmos (2024) could have been discussed in more detail. While the authors provide some intuition for why the decoupled learning of feature and classification layers improves calibration, they offer limited reasoning for selecting the latent dimension. Although they conduct multiple experiments with different Z, the reasoning behind the choice of the latent dimension is unclear, despite different values of Z leading to significant. variation in the results.
16

---

Published in Transactions on Machine Learning Research (06/2025)
References
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby, An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. URL htttp:://ppenreller wtt.a//t....ooomm///d///d/aa..
Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In Proceedings of The 33rd International Conference on Machine Learning. pp. 1050- 1059. PMLR, 2016. URL https::/proceedingg.. mlr. www.ll/.reee////////eee http://////glll.hh.... html ISSN: 1938-7228.
Jakob Gawlikowski, Cedrique Rovile Njieutcheu Tassi, Mohsin Ali, Jongseok Lee, Matthias Humt, Jianxiang Feng, Anna Kruspe, Rudolph Triebel, Peter Jung, Ribana Roscher, et al. A survey of uncertainty in deep neural networks. Artificial Intelligence Review, 56(Suppl 1):1513-1589, 2023.
Arindam Ghosh, Thomas Schaaf, and Matthew Gormley. Adafocal: Calibration-aware adaptive focal loss.
Advances in Neural Information Processing Systems, 35:1583 1595, 2022.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In International conference on machine learning, pp. 1321-1330. PMLR, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015. URL htttps:/arrxvv.or///ee 111.. 03385
Wenchong He, Zhe Jiang, Tingsong Xiao, Zelin Xu, and Yukun Li. survey on uncertainty quantification methods for deep learning, 2025. URL http://arxiv.. /tgg///ioooo/o//eoo abs/abas/////// 13425
Ramya Hebbalaguppe, Rishabh Patra, Tirtharaj Dash, Gautam Shroff, and Lovekesh Vig. Calibrating deep neural networks using explicit regularisation and dynamic data pruning, 2022. URL https://arxiw.orgg
abs/2212.00005
Dan Hendrycks and Thomas Dielterich. Benchmarking neural network robustness to common corruptions and perturbations. Proceedings of the International Conference on Learning Representations, 2019.
Dan Hendryeks and Kevin Gimpel. A baseline for detecting misclassified and out-of distributin examples in neural networks, 2018. URL
021366
Dan Hendryeks, Mantas Mazeika, and Thomas Dietterich. Deep anomaly detection with outlier exposure. Proceedings of the International Conference on Learning Representations, 2019.
Taejong Joo and Uijung Chung. Revisiting explicit regularization in neural networks for well-calibrated predictive uncertainty. ar preprint arXiv:2006.06:3 2020.
Mikkel Jordahn. Calibrated machine learning models: How to get them and why they matter. 2024.
Mikkel Jordahn and Pablo M. Olmos. Decoupling feature extraction and classification layers for calibrated neural networks. In Proceedings of the 41st International Conference on Machine Learning, ICML/24 JMLR.org, 2024.
Masanari Kimura and Hiroki Naganuma. Geometric insights into focal loss: Reducing curvature for enhanced model calibration. Pattern Recognition Letters, 2025.
Diederik P Kingma and Max Welling. Auto-cncoding variational bayes, 2022. URL https://arxiv.org/ abs/1312.
Viacheslav Komisarenko and Meelis Kull. Improving calibration by relating focal loss, temperature scaling, and properness. In F.AT 2024, pp. 1535-1542, TOS Press, 2024.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
17

---

Published in Transactions on Machine Learning Research (06/2025)
Aviral Kumar, Sunita Sarawagi, and Ujwwl Jain. Trainable calibration measures for neural networks from
kernel mean embeddings. In International Conference on Machine Learning, pp. 2014. PMT.R, 2018.
Ya Le and Xuan S. Yang. Tiny imagenet visual recognition challenge, 2015a. semanticscholar.. orgooo///ooo (o:rpusID:16664700
URL https://api.
Yann Le and Xuan Yang. Tiny imagenet visual recognition challenge. CS 231N, 7(7):3, 2015b.
Tsung- Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, and Piotr Dollár. Focal loss for dense object detection. In 2017 IEFE International Conference on Compuler Vision (ICCV), pp. 2999-3007 2017. doi: 10.1109/ICCOC L099IICCV.011..22... 2017.324.
Matthias Minderer, Josip Djolonga, Rob Romijnders, Frances Hubis, Xiaohua Zhai, Neil Houlsby, Dustin Tran, and Mario Lucic. Revisiting the calibration of modern neural networks Aduances in neural infor- malion processing systems, 34:15682 15694, 2021.
Jishnu Mukhoti, Viveka Kulharia, Amartya Sanyal, Stuart Golodetz, Philip Torr, and Puneet Dokania. Calibrating deep neural networks using focal loss. Advances in Neural Information Processing Systems, 33:15288-1133 15299, 2020.
Rafael Müller, Simon Kornblith, and Geoffrey E Hinton. When does label smoothing help? Advances in neural information processing systems, 32, 2019.
Lukas Muttenthaler, Robert Vandermeulen, Qiuyi Zhang, Thomas Unterthiner, and Robert Müller. Set learning for accurate and calibrated models. Xiv preprint arXiv:2307. 02245, 2023.
Mahdi Pakdaman Naeini, Gregory F. Cooper, and Milos Hauskrecht. Obtaining well calibrated probabilities using bayesian binning. In Proceedings of the Twenty-Nin Ninth AAAI Conference on Artificial Intelligence, AAAI115, pp. 2901- -901. AAAI Press, 2015. ISBN 0262511290.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Baolin Wu, Andrew Y Ng, et al. Reading digits in natural images with unsupervised feature learning. In NIPS workshop on deep learning and unsupervised feature learning, volume 2011, pp. 4. Granada, 2011.
Gabriel Pereyra, George Tucker, Jan Chorowski, Lukasz Kaiser, and Geoffrey Hinton. Regularizing neural networks by penalizing confident output distributions. arXiv preprint arxiv 1tt:::70100000 701.06548, 2017.
Afshar Shamsi, Hamzeh Asgharnezhad, AmirReza Tajally, Saeid Nahavandi, and Henry Leung. An uncortainty-aware loss function for training neural networks with calibrated predictions, 2023. URL
2atbb///110.026666 03260
Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In Proceedings of the IEFE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016.
Mingxing Tan and Quoe V. Le. Efficientne:: Rethinking model scaling for convolutional neural networks, 2020. URL
11946
Linwei Tao, Minjing Dong, and Chang Xu. Dual focal loss for calibration. Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett. (eds.), Proceedings of the 40th International Conference Machine Learning, volume 202 of Proceedings of Machine Learning
Research, pp. 33833 33849. PMLR, 23 29 Jul 2023a. URL httpps:/pproceedingg. mir. mlr, press/v202/ tao23a. html
Linwei Tao, Minjing Dong, and Chang Xu. Dual focal loss for calibration, 2023b. URL https://arxiv.
org/ abs/2305. 13665
Sunil Thulasidasan, Gopinath Chennupati, Jeff A Bilmes, Tanmoy Bhattacharya, and Sarah Michalak. On mixup training: Improved calibration and predictive uncertainty for deep neural networks. Advances in neural information processing systems, 32, 2019.
18

---

Published in Transactions on Machine Learning Research (06/2025)
Vishal Thanvantri Vasudevan, Abhinav Sethy, and Alireza Roshan Ghias. Towards better confidence
timation for neural models. In ICASSP 2019 2019 IEFE International Conference Acoustics,
Speech and Signal Processing (ICASSP), pp. 7335 7339, 2019. doi: 10.1109/ICAS0 0019.8683359, URL
hattps://eeexppooss
ISSN: 2379-190X.
Cheng Wang. Calibration in deep learning: A survey of the statc-of-the-art. arv preprint arXiv:2308.01222. 2023.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks, 2017. URL https://arxiv.cr org.ab// 1605.07146
A
Datasets
Table 13: Overview of used datasets and corresponding preprocessing steps.
Datasct
# Images
Classes
Preprocessing
Implementation
CIFARIO
60,000
10
32x32
Normalization, Random Crop, Random Horizontal Flip
PyTorch
CIFARI00
60,000
100
32x32
Normalization, Random Crop, Random Horizontal Plip
PyTorch
SVHN
600,000
10
32x32
Normalization
Py/Torch
CIFARIO-C
60,000
32x32
Corruption Added
Downloaded
CIFAR10-C
60,000
100
32x32
Corruption Added
Downloaded
Tuyy Image:Net
100,000
200
6dx64
Resiæe (224x224), Normolization
Gi Repo
B Reproducing results for CIFAR100
We noticed that our WRNN-28-10 base models for CIFAR100 achieved the lowest validation loss in the early epochs (14, 16, 17, 18, 18, 20, 28 and 61) even though the training accuracy reached around 100% and the validation accuracy kept improving to similar levels reported the original paper. The model reported in Table D is the one that achieved the lowest validation loss in epoch 61 (seed 2). In Table [4 and Table 15 we report the results when using seed 1 (our default seed for all other models). As we tried various seeds, used the code provided by the authors, and followed their instructions, models never reached their reported accuracy. Similar results may have been obtainable if we had known their seed or if we had saved the model with the highest validation accuracy rather than the model with the lowest validation loss. However, we did not explore the latter possibility further, as the original paper also saved the models based on the lowest validation accuracy, and because this could have had negative effects on calibration.
were surprised to see that our base
Table 14: Replication of CIFAR100 results with seed 1. Test metries on in-distribution test datasets. For V-TST, m = and m = 10 indicate using 1 and 10 samples in the MC approximation, respectively. Temp. WRN is the temperature-scaled WRN. 512 was used for V-TST (m = 10) and TST, and = 128 for
V-TST (m = 1).
Dataset
Model
Accuracy
ECE
MCE
Train NLL
Test NLL
Brier
V-TST, m= 10 64.32=0.09
9.92=0.288
21.74::1.349
0.5008=0.000
1.4159=0.0033
0.0049:10.000 0034
V-TST, m=
62.35+:0.07
10.71+0..188
20.48:0.7880
0.6064=0.0049
1.5634=0.004
0.00524:000000
CIFARI00
TST
61.82++0.16
10.32+0..50
20.55+22.400 404
0.6448=0.0089
1.5212=0.00m 0.007
0002=00000000S
Temp. WRN
52.31
2.72
6.42
1.2029
1.8358
0.0062
WRN
52.31
17.38
27.52
1.1451
2.0083
0.0066
Due to the accuracy differences between our CIFAR100 model and the original one, the typical calibration measures ECE and MCE cannot be directly applied compare their calibration. This because ECE is estimator whose bias is dependent the accuracy and is highest for models with accuracy around
19

---

Published in Transactions on Machine Learning Research (06/2025)
Table 15: Replication of CLFAR100 results with seed 1. Evaluation metrics on shifted data and OOD data.
Dataset
Model
SHIFT ECE
SHIFT MCE
OOD AUROC
OOD FPR95
V-TST, m 10
20.30+-0.4
32.39=0.66
0.727+0.002
0.817=0.002
V-TST, m
21.15:10.40
34 34.0110..00
0.71710.000 0.002
0.828-00002
CIFAR100
TST
21.13::0.76
30 82111.1
0.692:10.002
0. 0.846-0.0044
Temp. WRN WRN
11.09
20.46
0.677
0.878
28.15
43.18
0.658
0.885
50% while decreases for models with more extreme accuracies (closer to 0 or 1) (Minderer et al. 2021 Thercfore, one should always consider both accuracy and the calibration measure when comparing models with different classification errors and choose the bin size carefully (Minderer al., 2021] Here, use the same bin size (M 10) in the original paper for comparability. Alternative calibration measures such the NLL the Brier score don't suffer from the same bias and should be taken into account as well (Minderer et al. 2021
C Network architecture details
ViT architecture As the original study does not specify which weights or latent dimension they used, we decided to use the weights of google/vit- vitt-bae-pattch6--25 224-in21k from the transformers package which corresponds to the ViT-B/16 architecture and used the latent dimension Z 128 for both V-TST and TST, As Tiny ImageNet images are too small for the ViT, we had to resize them before using them as an input to and the ViTModel class from the transformer package
the model. We used the CerfoolmmgeCCasssfiaattttt to include the VIT in the code basc of the original paper.
ResNet-50 and EfficientNet-BA. architecture The architectures used for ResNet-50 and EfficientNet- B5 have been adapted to the smaller input image size of the CIFAR10 dataset. For ResNet-50, the parameters of the first convolutional layers have paramcters: stride 1, padding 1, kernel size of x 3. No max-pooling is performed after the first convolution and global average pooling performed over a smaller feature map of 4 x 4.
For EfficientNet-B5, the first convolutional layer has stride and padding of and uses kernel size of x 3. Similarly, the first down-sampling layer following the first convolutional layer is removed and the feature map of the global average pooling is reduced to 2 x 2.
For both ResNet and EfficientNet the number of output classes is changed to 10.
MLP Head Modifications For the single-layer configuration, we simply re-initialized the final FC layer. For the four-layer configuration, we extended the original three-layer MLP to have dimensions (output, 67). [62, 3Z], [3Z, 2] and [Z, #classes].
D Additional Results
Table 16:
Test metrics on in-distribution data for the effect of including Focal Loss (FL) in the training procedure. The first loss before + sign indicates the loss used for training the base model prior to TST and V-TST.
Dataset
Model Name
Model
Loss
Accuracy
ECE
MCE
Train NLL
Test NLL
V-TST, m 10
VI0-FI.CE
3 Fl. + CE
91.88-10.0
0.054
20.60-+2.488
0.0821 +000006
0.2840-0.001
CIFAR10
V-TST, m=
VI-FL-CE
3 F1, 4 CE
80.82=0.07
1.62+40076
35.68-48.88 816
0.1405=0..007
0.4124+4.0000
TST
FL-CE
+ CE
91.86=0011
2.03=0.020
18-48-13.253
0.0652=0.0002
0.2677=0.0005
20

---

Published in Transactions on Machine Learning Research (06/2025)
E Replication of Ablation Studies
Table 17: Replication of test metrics for ablation studies on in-distribution data test-sets using a. simple
CNN on CIFAR10. We used Z 512 for V-TST with m 10, Z 128 for V-TST with m 1, and Z 32 for TST.
Model
Accuracy
ECE
MCE
V-TST, 10
71.25=0.04
1.82+0.07
8.85+1.42
V-TST, m =
67.85+00007
1.25+0008
6.55+1.17
Ours
TST, Z = 128
71.0310.04
1.11110.55
18 18.64110.99
TST
70.55:10.05
0.73
14.8011..88
CNN
70.83
0.80
19.46
V-TST, 10 V-TST, m TST 70.77+0.2
75.9511..
3.27:0.05
10.65-1.33
73.2210..1
1.61:0.3
9.0815.47
Original
0.76+0.17
7.79+4.67
CNN
72.48
1.04
19.89
Table 18: Replication of test metrics for ablation studies on in-distribution data test-sets while fine-tuning VIT on Tiny ImageNet. We used Z 128 for V-TST and TST.
Model
Accuracy
ECE
MCE
V-TST, m 10
84.36:0.04
3.35=10.7
12.46:0.88
V-TST, m
83.61+00.08
1.56++0.88
11.13=1.17
Ours
TST
83.98++0.05
6.11+0..11
22.03-+2.00
FT ViT-B/16
83.38
3.33
13.60
V-TST, 10
86.7-:0.02
1.96:40.07
9.76-10.56
V-TST, m TST
85.6340.08
2.370.27
8.62=0.55
Original
85.99
3.4740.44
14.77=1.69
FT B/16
85.46
2.19
12.54
21


Reviews:
[{"review_id": "ACpWIcJlYA", "reviewer": "Reviewer_fWJz", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper's emphasis on reproducibility is valid and important. \nHowever, the contribution, in its present form, seems better suited for a workshop setting than a full paper at TMLR.\nOverall, the paper is missing an explanation of why additional settings were tested. The conducted ablations beyond the reproduction of prior work are missing a streamlined orientation e.g. in Section 5.2. Why did the authors use these additional models besides having different parameter counts? \n\nFurther, the results on CIFAR-100 are not competitive, with an accuracy of 64% also metrics such as ECE are not reliable anymore. However, a discussion on the reliability of ECE is missing.\nAdditionally, the authors fail to report what is actually used as confidence. The output after Softmax? If so, there is also a lot of literature discussing that the output after Softmax is not a suitable metric for confidence.\n\nTables are not referenced. Thus, it is not clear which Table belongs to which experiment.\nFurther, it is not clear to me what m and n a representing in Table 8.\n\nAnother question: Why is the 2-stage training with the focal loss so bad, can the authors elaborate more on that?\n\nThe authors omit to cite any of their data sources (CIFAR, ImageNet...)\n\nThe conclusion is too long and not precise enough.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper reproduces results from another paper which proposes a two-stage training procedure that decouples feature extraction and classification layers for better network calibration. Further, a few ablation studies have been conducted.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "I would suggest the authors, to shorten the reproduction of prior work or move it to the appendix.\nThe ablation studies require a more focused narrative; please clarify the key takeaways.\nFurthermore, please incorporate the missing information detailed above.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "7q08nF9Kfd", "reviewer": "Reviewer_w9zd", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n\n1. The paper has repeated the most relevant aspects of the original work.\n2. The authors add relevant and interesting ablation studies.\n3. The authors combine the original method with an interesting, relevant methodology in focal loss.\n\n# Weaknesses\n\n1. The authors do not reproduce the original findings to my satisfaction. It is not required for reproducibility studies to replicate the same results, but more steps could be attempted to reproduce exactly before having enough evidence that the results were not reproducible.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors present a reproduction study of the paper “Decoupling Feature Extraction and Classification layers for Calibrated Networks”, referred to as ‘original’ henceforth. The original work proposes a two-stage training (TST) paradigm to improve calibration of (image-domain) neural networks. The original work also proposes a variational (V-TST) version. \n\nThe reproduction effort repeats the main experiments of the original work on 2 out of 3 datasets, and also repeats two ablation studies, namely applying the technique with a smaller, less complex model, and applying it with pre-trained image representations networks.\n\nFurthermore, the authors add two ablation studies, namely on the depth of the classification network, and on the number of monta-carlo samples during training with V-TST.\n\nFinally, the authors extend the work by experimenting with the addition of focal loss next to TST and V-TST.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The main change I would like to see is for Table 1 and Table 2 to match closely with the original work, especially for Cifar100. In paragraph 5.1, the authors note the following:\n\n| We noticed that our base models for CIFAR100 achieved the lowest validation loss and were obtained in the\nearly epochs, 18 and 16, even though the training accuracy reached around 100% and the validation accuracy\nkept improving to similar levels reported in the original paper. Similar results may have been obtainable by\nexploring more seeds or saving the model with the highest validation accuracy rather than the model with\nthe lowest validation loss.\n\nCan the authors clarify whether their experiments use the base model with the lowest validation LOSS, and whether the original code uses base models with the lowest validation ACCURACY? I think this is a very relevant difference. It is usually the case that the validation accuracy can increase at the same as the validation loss, as the model increases in confidence during training, which leads to higher loss values on samples where the model was confidently wrong. At the same time, the validation accuracy still increases as there are things to learn. This has, of course, a large impact on calibration, as the longer you train, the less calibrated the network will be. Therefore, I can imagine that this small change in experimental methodology can lead to large effects in the calibration results.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "IJuc7oOZlm", "reviewer": "Reviewer_UyE4", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n- S1: Experimental settings are well described.\n\n## Weaknesses\n- W1: Replication of TST itself is not a contribution since Jordahn & Olmos released the code. Is there any problem when running the code or suspicion about the results?\n- W2: The motivation for the replication is unclear. Clarifying the importance of (V-)TST and why we should replicate the results would strengthen the motivation.\n- W3: The justification of the proposed method's design, combining the focal loss with TST, is unclear.\n- W4: The related work section should be more extensive. Reviewing existing calibration studies and the importance of the TST and focal loss would strengthen the paper.\n- W5: For examining TST, experimenting with datasets not used in Jordahn & Olmos's paper or evaluating with other metrics such as Brier score [a] would bring new insights.\n\n\n[a] https://en.wikipedia.org/wiki/Brier_score", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper replicates the existing work (Jordahn & Olmos, 2024) results and proposes combining the focal loss with (V-)TST, which is the method by Jordahn & Olmos, to further improve overconfidence and calibration in image classification.\nThis paper also extends the experiments by Jordahn & Olmos with more ablation and model variants.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "## Major changes\n- For W1 and W2: Clarifying the motivation of replicating Jordahn & Olmos's results is necessary. It would also be compelling to explain why TST is essential and worth replicating since many calibration studies exist. If there was any problem when replicating TST, describing the problem and how it was addressed would support the contribution of the replication.\n- For W3: Unveiling the problem of TST and describing how the focal loss solves the TST's problem would be interesting and strengthen the reasonability of the method. The current paper just combines the two methods and not motivated.\n- For W4: Reviewing other existing calibration studies, explaining their problems, and explaining why TST and focal loss are focused on in this paper would strengthen the motivation.\n- For W5: Experimenting with datasets not used in Jordahn & Olmos's paper or evaluating with other metrics such as Brier score would bring new insights. For instance, Cars, Pets, Food, Birds, Flowers, domain adaptation datasets, etc., would be beneficial for practitioners.\n\n\n## Minor changes\n- The citation of the most important work in this paper, Jordahn & Olmos (2024), is inaccurate. It has been [accepted to ICML2024](https://proceedings.mlr.press/v235/jordahn24a.html) but cited as an arxiv paper.\n- Adding the scores from the original paper to Tabs. 1-4 would be helpful for checking the replication.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors have addressed some of the concerns raised by the reviewers. However, the two reviewers still have concerns regarding the justification for the method design and necessary explanations. The reviewers recommend the authors provide more explanations for performance drivers and ablation studies, and conduct more experiments to show the combination of TST with various calibration methods beyond focal loss in an online setting.
Decision: Accept (poster)