PaperlayerView exact revision on arXiv

Source-linked AI summary

Decoupling Feature Extraction and Classification Layers for Calibrated Neural Networks

Mikkel Jordahn, Pablo M. Olmos

arXiv:2405.01196v3cs.LGstat.ML

TL;DR

Over-parameterized DNNs can be poorly calibrated, a concern for safety-related classification. The paper decouples feature extraction and classification through two-stage training, with a variational variant that adds a prior over hidden features. It reports improved calibration and attributes further improvement to variational training.

  • Problem

    Over-parameterized DNNs trained with cross-entropy tend to be overconfident, while reliable probabilities matter in safety-related applications.

    Source: §1 · p. 1

  • Method

    The paper trains feature extraction and classification in separate stages, then proposes a variational version with a probabilistic prior on hidden features.

    Source: §1 · pp. 1–2, §3 · p. 3

  • Results

    The authors report that two-stage training improves calibration and that variational training with a prior on hidden features improves it further.

    Source: §1 · pp. 1–2, §7 · p. 8

  • Takeaways & Limitations

    In the paper's account, freezing feature extraction reduces model flexibility, while variational training augments features and supports mapping multiple latent points to the same label.

    Source: §7 · p. 8

    V-TST prediction requires more Monte Carlo samples to significantly improve performance, and its training is slightly more expensive than TST.

    Source: §6 · p. 8

Abstract

from arXiv · show

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.

1. Introduction

Calibration matters in safety-related classification, but over-parameterized DNNs trained with cross-entropy tend to be overconfident. The paper presents two-stage and variational two-stage training to improve calibration by decoupling feature extraction from classification.

1. Introduction

  • Over-parameterized DNNs trained with cross-entropy tend to be overconfident, while calibration concerns the reliability of assigned class probabilities.The paper frames calibration as important in applications such as medical diagnostics and self-driving systems.Source: §1 · p. 1
  • The paper situates its approach alongside Bayesian neural networks, post-hoc recalibration, and data-augmentation methods for calibration.It distinguishes post-hoc recalibration as occurring after regular model training.Source: §1 · p. 1
  • The paper reports that decoupling feature extraction and classification can significantly improve DNN calibration, and presents TST and V-TST as two training procedures.TST retrains the classification layers while feature extraction is fixed; V-TST adds a probabilistic prior and variational training.Source: §1 · pp. 1–2
  • The authors report that a Gaussian prior on final hidden-layer outputs with variational classifier training further improves calibration.This is the additional calibration strategy in V-TST.Source: §1 · p. 2

2. A Motivational Example

A WRN example compares one-stage end-to-end training with retraining the classifier while freezing convolution layers. The two-stage model retains similar accuracy, while its high-entropy predictions are more associated with cluster boundaries.

2. A Motivational Example

  • In the end-to-end model, class structure is visible in the hidden space, but it is largely lost among the points with greatest prediction uncertainty.The end-to-end model's overall classifier accuracy is 91.41%.Source: §2 · p. 2
  • 91.61% versus 91.41% accuracy; after two-stage training, high-entropy points appear more clearly separated and generally lie near class-cluster boundaries.The calibration plots corroborate improved calibration, and the entropy histogram shows fewer near-zero-entropy predictions.Source: §2 · pp. 2–3
  • The example compares a WRN trained end-to-end with one whose FC classifier is retrained from scratch while convolution layers are frozen.The reported accuracy remains roughly the same after the second-stage retraining.Source: §2 · p. 2

3. Two-Stage Training

Two-Stage Training first fits the network end-to-end, then freezes feature extraction and reinitializes and retrains the FC classifier. V-TST keeps the feature extractor fixed while regularizing the latent feature space with a Gaussian prior and training classifier layers using an ELBO-like objective.

3. Two-Stage Training

  • TST trains end-to-end with cross-entropy, freezes feature-extraction parameters, reinitializes the FC layers, and retrains those layers on the same dataset.Stage 2 is not pre-training and fine-tuning: both stages use identical training data.Source: §3 · p. 3
  • The authors hypothesize that choosing the Stage 2 classifier structure, particularly the dimensionality of latent output z, can further affect calibration.This motivates V-TST's additional structure in the feature space.Source: §3 · p. 3
  • V-TST places a Gaussian prior on latent z and trains the FC layers with an ELBO-like objective while keeping feature-extraction parameters fixed.The prior is kept simple and independent of the image x to impose stronger regularization in z space.Source: §3 · p. 3
  • V-TST uses a conditional variational posterior for z that depends on classifier input x; its mean and variance networks include frozen feature extraction and trainable FC layers.The FC parameters are optimized to maximize the lower bound.Source: §3 · p. 3
  • At prediction time, V-TST estimates p(y|x) by Monte Carlo sampling from the variational posterior.The ELBO is optimized over training pairs using mini-batch optimization and Monte Carlo sampling.Source: §3 · p. 3

4. Related Work

Related work spans Bayesian parameter-prior methods, implicit regularization, data augmentation, and post-hoc calibration. The paper distinguishes V-TST from other variational and Gaussian-feature-space approaches.

4. Related Work

  • Bayesian neural networks place priors on some or all parameters and approximate a posterior, with examples including Deep Ensembles and MC-Dropout (Lakshminarayanan et al., 2017; Gal & Ghahramani, 2016).Other cited examples include Laplace Approximation, Mean-Field BNNs, and Node-Based BNNs (Dusenberry et al., 2020; Trinh et al., 2022).Source: §4 · p. 4
  • TST is an implicit regularization method for uncertainty quantification, complementary to Focal Loss; related approaches also use Mix-Up or Label-Smoothing, while Temperature Scaling is post-hoc (Lin et al., 2020; Mukhoti et al., 2020; Thulasidasan et al., 2019; M¨uller et al., 2019; Guo et al., 2017).Source: §4 · p. 4
  • Unlike the Variational Classifier of Dhuliawala, Sachan, and Allen (2024), V-TST trains feature extractors supervisedly and uses a stochastic final layer with an ELBO for the marginal label distribution p(y).The cited classifier uses an auxiliary discriminator and mini-max ELBO optimization, and trains end-to-end; the paper describes V-TST's ELBO objective as simpler and fast-converging.Source: §4 · p. 4
  • Wan, Zhong, Li, and Chen (2018a) impose Gaussian structure on classifier features with a likelihood regularizer, but do not use an ELBO or two-stage training and report no calibration improvement.Source: §4 · p. 4

5. Experiments

Across WRN and ViT image-classification experiments, the authors evaluate two-stage training against standard and temperature-scaled baselines. They report improved calibration with mostly comparable accuracy, alongside sampling and architecture ablations that qualify where these results hold.

5. Experiments

  • Experiments compare WRN-28-10 on CIFAR10, CIFAR100, and SVHN and ViT on CIFAR10, using frozen feature-extraction layers for Stage 2.Evaluation covers accuracy, ECE, MCE, and NLL, plus shifted and OOD data; Stage 2 uses 10 reinitialization seeds for TST and V-TST.Source: §5 · p. 4

5.1. Benchmark Results

  • V-TST lowers base-WRN ECE by 78.40% on CIFAR10, 73.18% on CIFAR100, and 87.64% on SVHN, with comparable accuracy except a slight degradation at m = 1.Figures 4 and 5 illustrate ECE improvements across bins rather than only in particular bins.Source: Table 1 · p. 5, §5.1 · p. 5, Figure 4 · p. 5
  • On shifted data, TST and V-TST improve calibration and outperform temperature scaling, but they generally perform worse than baselines on OOD detection.Both methods also decrease training- and test-set NLL; Figure 6 shows increased prediction entropy and Gaussian distributions within each V-TST class.Source: §5.1 · pp. 5–6, Table 2 · p. 6

5.2. MC Samples Effect on V-TST Performance

  • Increasing V-TST test-time Monte Carlo samples improves accuracy across models and ECE for all tested latent dimensions except 32.The authors note that tuning sample count m could improve reported results.Source: §5.2 · p. 6, Figure 7 · p. 7

5.3. Ablation Study: Under-parametrized Architecture

  • On the simpler CNN, V-TST worsens ECE but improves MCE, while TST improves both calibration metrics at the cost of accuracy.The authors characterize this architecture as not over-parameterized because its training and test accuracy are almost identical.Source: §5.3 · p. 6, Table 3 · p. 7

5.4. Ablation Study: Training Models End-to-End

  • One-step training of adapted architectures does not match the two-stage calibration metrics, although variational end-to-end training improves over CE-only end-to-end training.This comparison suggests the VAE-like stochastic layer itself has regularization properties that improve calibration and uncertainty quantification.Source: §5.4 · p. 7, Table 4 · p. 7

5.5. Ablation Study: Visual Transformer

  • On CIFAR10, V-TST maintains ViT accuracy while significantly improving ECE and MCE.This extends the reported calibration improvements beyond convolutional architectures.Source: Table 5 · p. 8, §5.5 · p. 8

6. Limitations of TST and V-TST

TST and V-TST add a second training stage, but that stage converges quickly; V-TST also increases prediction cost when using more Monte Carlo samples. The methods can sometimes reduce OOD performance, and likelihood-based OOD detection remains an open investigation.

6. Limitations of TST and V-TST

  • TST's second stage converges quickly, but V-TST requires more training time and Monte Carlo samples at prediction time.For WRN, best validation loss came after 1.167 epochs on average for TST and 10.33 for V-TST; CIFAR10 prediction took 0.932 seconds for TST and 1.406 seconds for V-TST with m = 100.Source: §6 · p. 8
  • TST and V-TST sometimes reduce performance on OOD data, despite fitting training data well, including in terms of calibration.The authors suggest that training stochastically with respect to inputs does not teach the model about unseen modalities.Source: §6 · p. 8
  • The authors propose BNN methods and prior likelihoods for new test points as future OOD directions, noting varied success for likelihood-based detection in generative models.The cited example is work on VAEs.Source: §6 · p. 8

7. Conclusion

The authors conclude that jointly training feature extraction and classification layers is a main driver of poor calibration in over-parameterized DNNs. They argue that freezing features constrains the classifier, while a prior on the feature space and variational training further improve calibration.

7. Conclusion

  • Jointly training feature extraction and classification layers is identified as a main driver of poor calibration in WRN and ViT models.Freezing feature extraction layers reduces model flexibility; the authors hypothesize this limits how far Stage 2 can move points from decision boundaries to increase label likelihood artificially.Source: §7 · p. 8
  • A probabilistic prior on the feature space further improves calibration by sampling features during variational training, which maps multiple z points to the same label.The end-to-end ablation also found better calibration for variational training than for regular end-to-end training.Source: §7 · p. 8

A. Training Details for WRN and ViT

The appendix specifies the WRN and ViT training setups, Stage 2 model configurations, SVHN rotation severities, a small-CNN setting, and conventions for reporting additional WRN results.

A. Training Details for WRN and ViT

  • WRNs were trained for up to 600 epochs and ViTs for up to 1000, using Adam at 10^-4 with validation-loss early stopping.WRN validation used 15% of CIFAR10 and SVHN training data and 5% of CIFAR100; the ViT setup used a patch size of 4 and about 9.5 million parameters.Source: A. Training Details for WRN and ViT · p. 11

B. Training details for TST and V-TST

  • Stage 2 reinitializes and trains the classifier layers on the same data, testing latent dimensions Z = 2, 8, 32, 128, 256, and 512.V-TST uses two MLPs to parameterize the variational distribution's mean and scale; its scale output uses a log-sigmoid activation.Source: B. Training details for TST and V-TST · p. 11

C. SVHN Rotations

  • SVHN shifted-data metrics are evaluated on the full test set at rotation severities of 10, 45, 90, 135, and 180 degrees.Source: C. SVHN Rotations · p. 11

D. Small CNN Architecture

E. Additional WRN Experiment Results

  • Additional WRN TST and V-TST results report means and standard errors over 10 seeds; m denotes samples used to approximate the predictive expectation.E-to-E and Var. E-to-E label the corresponding TST and V-TST architectures trained end-to-end in one stage.Source: E. Additional WRN Experiment Results · pp. 11–12

E.1. CIFAR10 WRN Two-Stage Results

The supplied appendix table captions identify WRN results by dataset and evaluation setting, covering in-distribution, shifted-data, and OOD tests; they do not state the outcomes.

E.1. CIFAR10 WRN Two-Stage Results

E.2. CIFAR100 WRN Two-Stage Results

E.3. SVHN WRN Two-Stage Results

F. CIFAR10 CNN Two-Stage Results

The CIFAR10 CNN results are reported separately for in-distribution, shifted, and out-of-distribution settings.

G. CIFAR10 ViT Two-Stage Results

The CIFAR10 ViT results are reported separately for in-distribution, shifted, and out-of-distribution settings.

Loading 2405.01196v3…