Why decoupled feature-extraction and classifier training improves neural network calibration: a theoretical characterization of TST/V-TST, its variational design choices, and its out-of-distribution limit

Jordahn and Olmos (ICML 2024) showed empirically that decoupling the training of a deep classifier into two stages — train the full network, then freeze the feature extractor and retrain only the classification head, optionally with a variational bottleneck (Two-Stage Training, TST, and its variational variant V-TST) — substantially improves calibration of over-parametrized networks at essentially no accuracy cost, but at times degrades out-of-distribution (OOD) detection. The mechanism is only hypothesized (freezing features prevents the new classifier from pushing data points arbitrarily far from decision boundaries), and the variational architecture's design choices and Monte Carlo sample budget are set heuristically. This problem asks for a rigorous, predictive theoretical account of the TST/V-TST calibration phenomenon: why and when decoupling improves calibration, why the benefit coexists with degraded OOD detection, and what this theory implies for the variational design choices and the prediction-time Monte Carlo sample count. The 2024–2026 follow-up literature (a TMLR reproducibility study with empirical extensions, classifier-side calibration methods, and applied adopters) is exclusively empirical, leaving the theoretical core open.

Background

A probabilistic classifier is calibrated if among test points to which it assigns confidence $c$, roughly a fraction $c$ are actually classified correctly; modern deep neural networks (DNNs), especially over-parametrized ones trained with cross-entropy (CE) loss, are systematically overconfident, i.e. predicted confidence exceeds empirical accuracy (Guo et al., 2017). Calibration is commonly measured by the Expected Calibration Error (ECE), a binned average of |confidence − accuracy|, its worst-bin variant MCE, and related scores; post-hoc fixes such as temperature scaling (rescaling logits by a single factor fitted on held-out data) correct confidence but leave the learned representation untouched. Theory on the sources of miscalibration exists mainly for simplified settings: Bai et al. (2021) proved that margin-based empirical risk minimization is overconfident even in a well-specified, non-over-parametrized logistic model, so over-parametrization alone is not the whole story. Wang and Zhang (2024) showed that over-compressed representations are harder to calibrate post hoc (the calibration bottleneck) and proposed, based on a weak-classifier hypothesis (a head that is not over-trained lets the representation learn more calibratable features), a progressive layer-peeled training method (PLP) that peels off layers from the top during training. On the Bayesian side, Kristiadi et al. (2020) showed that a last-layer Laplace approximation — a Gaussian posterior over only the final layer's weights — removes much of the overconfidence of ReLU networks, establishing that confining uncertainty to the last layer is a powerful lever.

Jordahn and Olmos (ICML 2024; arXiv:2405.01196) contribute an empirical training-based lever in the same spirit. In Two-Stage Training (TST), Stage 1 trains a network (Wide Residual Network, WRN, or Vision Transformer, ViT) end-to-end with CE; Stage 2 discards and re-initializes the classification head (replacing it with a small MLP) and retrains only this head on frozen Stage-1 features. In the variational variant V-TST, the Stage-2 head is preceded by a stochastic latent layer $\mathbf{z}$ of dimension $Z$: a Gaussian prior $p(\mathbf{z})$ is placed on the last hidden representation, a factorized Gaussian variational posterior $q(\mathbf{z}\mid \mathbf{x})$ with diagonal covariance is trained together with the classifier by maximizing the evidence lower bound (ELBO), i.e. CE minus the Kullback–Leibler (KL) divergence between $q(\mathbf{z}\mid\mathbf{x})$ and the prior, and prediction averages the class probabilities over Monte Carlo (MC) samples of $\mathbf{z}$. On CIFAR-10/100 and SVHN, TST and especially V-TST markedly reduce ECE/MCE while retaining accuracy; an ablation shows the benefit weakens or disappears on an under-parametrized CNN that does not overfit, and end-to-end training of the same adapted architecture does not reproduce the gains. For stability the authors restrict $q$ to a diagonal covariance, and they downscale the KL term by $Z$ because it otherwise dominates the CE loss. However, TST and V-TST at times degrade OOD detection — flagging inputs from classes or datasets never seen in training, measured by the area under the ROC curve (AUROC) and the false-positive rate at 95% true-positive rate (FPR95) — and the authors observe that fitting training data well, even in calibration terms, teaches nothing about unseen modalities. The authors leave open, as future work, whether combining two-stage decoupling with Bayesian or other uncertainty methods — for instance a last-layer Bayesianization, or an OOD score built from the prior likelihood $p(\mathbf{z})$ — can yield calibrated in-distribution probabilities together with reliable OOD detection.

The paper is explicit that the phenomenon is understood only at the level of a hypothesis: freezing the feature extractor may improve calibration because the Stage-2 classifier can no longer move data points arbitrarily far from the decision boundaries to artificially increase likelihood, whereas joint training can co-adapt features and classifier to do so. No formal model, theorem, or quantitative prediction supports or refutes this picture, and the relationship between the in-distribution calibration gain and the OOD degradation is unexplained.

Problem Statement

Provide a rigorous, predictive theoretical account of the calibration benefit of decoupled training in over-parametrized neural network classifiers, in the setting of Jordahn and Olmos (ICML 2024), resolving the following linked questions, which a single coherent theory should address in full, with question (1) as its primary target:

(1) Mechanism (primary). In a formal model of Two-Stage Training — Stage 1 trains feature extractor and classifier jointly by cross-entropy minimization; Stage 2 freezes the features and retrains a re-initialized classification head (optionally through a variational latent bottleneck, V-TST) — prove, refute, or replace the paper's hypothesis that freezing the feature extractor improves calibration because the Stage-2 classifier can no longer move data points arbitrarily far from the decision boundaries to artificially inflate likelihood. The characterization should be quantitative enough to predict the observed boundary conditions: substantial calibration improvement (e.g. in ECE/MCE) at essentially unchanged accuracy for over-parametrized WRN/ViT on benchmarks like CIFAR-10/100 and SVHN, weakening or vanishing for an under-parametrized CNN that does not overfit, and non-reproduction under end-to-end training of the same adapted architecture.

(2) Out-of-distribution limit. Explain, within the same theory, why TST/V-TST can degrade OOD detection (AUROC/FPR95), and characterize when calibrated in-distribution probabilities and reliable OOD detection are compatible.

(3) Variational design choices. Account, within the theory or an equivalent analysis, for the V-TST design decisions that the paper makes heuristically: latent dimension $Z$, the strength and form of the Gaussian prior on $\mathbf{z}$, the diagonal-covariance approximation to the variational posterior, and the downscaling of the KL term by $Z$.

(4) Monte Carlo sample count. Give a principled rule for selecting the number of prediction-time MC samples $m$ in V-TST under a latency or compute constraint, replacing the current empirical tuning.

The common scientific objective is to convert the empirical TST/V-TST phenomenon into predictive understanding: a theory that answers (1) and thereby explains and governs (2)–(4).

The parts are independently refutable: for (1), a proof, a refutation, or a replacement of the paper's hypothesis counts as resolving it only if it makes the three boundary-condition predictions stated under (1) above; (2)–(4) are each resolved by an account grounded in that same theory.

References

1. Mikkel Jordahn and Pablo M. Olmos, Decoupling Feature Extraction and Classification Layers for Calibrated Neural Networks, Proceedings of the 41st International Conference on Machine Learning (ICML 2024), PMLR 235:22530–22550, 2024. arXiv:2405.01196, https://proceedings.mlr.press/v235/jordahn24a.html
2. Chuan Guo, Geoff Pleiss, Yu Sun and Kilian Q. Weinberger, On Calibration of Modern Neural Networks, Proceedings of the 34th International Conference on Machine Learning (ICML 2017), PMLR 70:1321–1330, 2017. arXiv:1706.04599, https://proceedings.mlr.press/v70/guo17a.html
3. Yu Bai, Song Mei, Huan Wang and Caiming Xiong, Don't Just Blame Over-parametrization for Over-confidence: Theoretical Analysis of Calibration in Binary Classification, Proceedings of the 38th International Conference on Machine Learning (ICML 2021), PMLR 139:566–576, 2021. arXiv:2102.07856, https://proceedings.mlr.press/v139/bai21c.html
4. Agustinus Kristiadi, Matthias Hein and Philipp Hennig, Being Bayesian, Even Just a Bit, Fixes Overconfidence in ReLU Networks, Advances in Neural Information Processing Systems 33 (NeurIPS 2020), 2020. arXiv:2002.10118, https://arxiv.org/abs/2002.10118
5. Deng-Bao Wang and Min-Ling Zhang, Calibration Bottleneck: Over-compressed Representations are Less Calibratable, Proceedings of the 41st International Conference on Machine Learning (ICML 2024), PMLR 235:52156–52170, 2024. https://proceedings.mlr.press/v235/wang24cm.html
6. Johanna D'Ciofalo Khodaverdian, Eric Banzuzi and Katharina Deckenbach, A Reproducibility Study of Decoupling Feature Extraction and Classification Layers for Calibrated Neural Networks, Transactions on Machine Learning Research (TMLR), 2025. OpenReview forum 5Hwzd48ILf, https://openreview.net/forum?id=5Hwzd48ILf; dblp record, https://dblp.org/rec/journals/tmlr/KhodaverdianBD25
7. Jiani Ni, He Zhao, Jintong Gao, Dandan Guo and Hongyuan Zha, Balancing Two Classifiers via A Simplex ETF Structure for Model Calibration, Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR 2025), 30712–30721, 2025. DOI 10.1109/CVPR52734.2025.02860, arXiv:2504.10007, https://arxiv.org/abs/2504.10007
8. Jiani Ni, He Zhao, Yibo Yang and Dandan Guo, Deep neural network calibration by reducing classifier shift with stochastic masking, Pattern Recognition 177:113217, 2026. DOI 10.1016/j.patcog.2026.113217, arXiv:2508.09116, https://arxiv.org/abs/2508.09116
9. Mikkel Jordahn, Jonas Vestergaard Jensen, Mikkel N. Schmidt and Michael Riis Andersen, On Local Posterior Structure in Deep Ensembles, 2025. arXiv:2503.13296, https://arxiv.org/abs/2503.13296
10. Nadarasar Bahavan, Sachith Seneviratne and Saman Halgamuge, SphOR: A Representation Learning Perspective on Open-set Recognition for Identifying Unknown Classes in Deep Learning Models, 2025. arXiv:2503.08049, https://arxiv.org/abs/2503.08049
11. Thiru Thillai Nadarasar Bahavan, Sachith Seneviratne and Saman K. Halgamuge, FedLAS: Feature-Modulated Bidirectional Label Smoothing for Neural Network Calibration, 2026. arXiv:2606.28654, https://arxiv.org/abs/2606.28654
12. Robert Platt, Rossella Arcucci and Cédric M. John, Mineralogical Classification of CRISM Hyperspectral Data Under Uncertainty With Hybrid Neural Networks, Journal of Geophysical Research: Planets 131:e2025JE009473, 2026. DOI 10.1029/2025JE009473, https://doi.org/10.1029/2025JE009473