# Decoupling Feature Extraction and Classification Layers for Calibrated Neural Networks

arXiv: https://arxiv.org/abs/2405.01196 · HTML: https://mlsanity.com/paper/2405.01196/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 37 references that are in the map.
- First posted on arXiv 2 May 2024.

## Builds on (most cited first, 20 of 20)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Wide Residual Networks](https://mlsanity.com/paper/1605.07146/) (2016)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution](https://mlsanity.com/paper/2202.10054/) (2022)
- [Revisiting the Calibration of Modern Neural Networks](https://mlsanity.com/paper/2106.07998/) (2021)
- [Calibrating Deep Neural Networks using Focal Loss](https://mlsanity.com/paper/2002.09437/) (2020)
- [On Mixup Training: Improved Calibration and Predictive Uncertainty for Deep Neural Networks](https://mlsanity.com/paper/1905.11001/) (2019)
- [Being Bayesian, Even Just a Bit, Fixes Overconfidence in ReLU Networks](https://mlsanity.com/paper/2002.10118/) (2020)
- [Efficient and Scalable Bayesian Neural Nets with Rank-1 Factors](https://mlsanity.com/paper/2005.07186/) (2020)
- [End-to-end Training for Whole Image Breast Cancer Diagnosis using An All Convolutional Design](https://mlsanity.com/paper/1711.05775/) (2017)
- [Rethinking Feature Distribution for Loss Functions in Image Classification](https://mlsanity.com/paper/1803.02988/) (2018)
- [Machine learning based disease diagnosis: A comprehensive review](https://mlsanity.com/paper/2112.15538/) (2021)
- [Bin-wise Temperature Scaling (BTS): Improvement in Confidence Calibration Performance through Simple Scaling Techniques](https://mlsanity.com/paper/1908.11528/) (2019)
- [Robust Variational Autoencoders for Outlier Detection and Repair of Mixed-Type Data](https://mlsanity.com/paper/1907.06671/) (2019)
- [Tackling covariate shift with node-based Bayesian neural networks](https://mlsanity.com/paper/2206.02435/) (2022)

## Built on by (most cited first, 5 of 5)

- [Balancing Two Classifiers via A Simplex ETF Structure for Model Calibration](https://mlsanity.com/paper/2504.10007/) (2025)
- [SphOR: A Representation Learning Perspective on Open-set Recognition for Identifying Unknown Classes in Deep Learning Models](https://mlsanity.com/paper/2503.08049/) (2025)
- [Deep Neural Network Calibration by Reducing Classifier Shift with Stochastic Masking](https://mlsanity.com/paper/2508.09116/) (2025)
- [FedLAS: Feature-Modulated Bidirectional Label Smoothing for Neural Network Calibration](https://mlsanity.com/paper/2606.28654/) (2026)
- [Calibration in One Go: Jointly Learning Accurate and Calibrated Neural Networks](https://mlsanity.com/paper/2610.13693/) (2026)

## Abstract (arXiv)

Deep Neural Networks (DNN) have shown great promise in many classification applications, yet are widely known to have poorly calibrated predictions when they are over-parametrized. Improving DNN calibration without comprising on model accuracy is of extreme importance and interest in safety critical applications such as in the health-care sector. In this work, we show that decoupling the training of feature extraction layers and classification layers in over-parametrized DNN architectures such as Wide Residual Networks (WRN) and Visual Transformers (ViT) significantly improves model calibration whilst retaining accuracy, and at a low training cost. In addition, we show that placing a Gaussian prior on the last hidden layer outputs of a DNN, and training the model variationally in the classification training stage, even further improves calibration. We illustrate these methods improve calibration across ViT and WRN architectures for several image classification benchmark datasets.