Decoupling Feature Extraction and Classification Layers for Calibrated Neural Networks

TL;DR

Calibrated predictions are crucial in safety-critical settings, but over-parameterized DNNs tend to be poorly calibrated under cross-entropy training. The authors propose decoupling feature extraction from classification via Two-Stage Training (TST) and a variational variant (V-TST); freezing the feature extractor and retraining the classifier improves calibration with little accuracy loss, and adding a Gaussian prior with ELBO further enhances calibration. They validate on WRN and ViT across CIFAR-10, CIFAR-100, and SVHN, including distribution-shift and OOD scenarios, showing substantial ECE/MCE reductions and overall robust calibration improvements, though OOD performance can be mixed. The approach offers a practical, low-cost path to better-calibrated neural networks for image classification.

Abstract

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.

Figures (7)

• Figure 1: Last hidden layer outputs 2D t-SNE plots of test points from a WRN trained in one stage (left column), and two stages (center column) on CIFAR10 with accuracy $91.41\%$ and $91.61\%$ respectively. In the top row of t-SNE plots, (a) and (c), we plot all test points with equal color intensity, whilst in the second row, (b) and (d), higher color intensity indicates higher entropy of the prediction for a given point. Colors indicate which class the point truly belongs to. In (e) we show a calibration plot for the two models, and in (f) we show a histogram over the entropies of the predictions made by each model on the test set.
• Figure 2: Block Diagram of Model $M$ and belonging parameters. Red background denotes frozen parameters.
• Figure 3: Latent Variable Model
• Figure 4: ECE plot of V-TST trained model and base WRN trained on SVHN.
• Figure 5: ECE plot of V-TST trained model and base WRN trained on CIFAR100.