Research map / Architectures and efficiency / RNNs and normalization
Decoupling Feature Extraction and Classification Layers for Calibrated Neural Networks
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Two-Stage Training (TST) that freezes feature extractor and retrains a re-initialized classifier head, plus a variational variant (V-TST) with a Gaussian latent on the last hidden layer, to improve calibration of WRN and ViT without hurting accuracy.
Problem
Over-parameterized DNNs trained end-to-end with cross-entropy are overconfident and poorly calibrated, a problem in safety-critical settings such as healthcare.
Approach
Decouple feature extraction from classification: train end-to-end, then freeze the convolution/attention layers and retrain re-initialized FC layers in a second stage on the same data; V-TST further places a Gaussian latent prior on the last hidden layer and trains the head with an ELBO.
Key results
- V-TST reduces WRN ECE by 78.40% on CIFAR10, 73.18% on CIFAR100 and 87.64% on SVHN while keeping accuracy (92.58% vs 92.53% on CIFAR10).
- TST reaches 92.59% CIFAR10 accuracy with ECE 2.18 vs 5.88 for the WRN baseline.
- On CIFAR10 ViT, V-TST m=10 gets ECE 1.23 vs 13.71 for base ViT at 82.84% accuracy.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Freeze the feature extractor and retrain the classifier | WRN-28-10 with TST · CIFAR10, CIFAR100, and SVHN test sets | ECE: 5.88 -> 2.18; 21.74 -> 7.07; 3.48 -> 0.43 | Helped |
| Add a Gaussian latent prior and train the classifier variationally | WRN-28-10 with V-TST, m=10 · CIFAR10, CIFAR100, and SVHN test sets | ECE: 5.88 -> 1.27; 21.74 -> 5.83; 3.48 -> 0.47 | Helped |
| Use V-TST with 10 MC samples instead of TST | WRN-28-10 · CIFAR10 test | ECE: 2.18 -> 1.27 | Helped |
| Use V-TST with 10 MC samples instead of TST | WRN-28-10 · CIFAR100 test | ECE: 7.07 -> 5.83 | Helped |
| Use V-TST with 10 MC samples instead of TST | WRN-28-10 · SVHN test | ECE: 0.43 -> 0.47 | Hurt |
| Freeze the feature extractor and retrain the classifier | WRN-28-10 with TST · CIFAR10-C and CIFAR100-C | shift ECE: 20.27 -> 11.62; 40.47 -> 17.44 | Helped |
| Freeze the feature extractor and retrain the classifier | WRN-28-10 with TST · SVHN rotated test | shift ECE: 47.06 -> 30.97 | Helped |
| Apply two-stage training to an under-parametrized CNN (ablation) | Small CNN · CIFAR10 test | accuracy: 72.48 -> 70.77 | Hurt |
| Apply TST to an under-parametrized CNN (ablation) | Small CNN · CIFAR10 test | ECE: 1.04 -> 0.76 | Helped |
| Apply V-TST to an under-parametrized CNN (ablation) | Small CNN · CIFAR10 test | ECE: 1.04 -> 3.27 | Hurt |
| Train the adapted architecture end-to-end with CE instead of two-stage TST (ablation) | WRN-28-10 · CIFAR10 test | ECE: 2.18 -> 6.74 | Hurt |
| Train the variational adapted architecture end-to-end instead of two-stage V-TST (ablation) | WRN-28-10 · CIFAR10 test | ECE: 1.27 -> 4.72 | Hurt |
Surprising findings
- TST and V-TST sometimes reduce OOD detection performance relative to the end-to-end WRN baseline.
- V-TST worsens ECE on an under-parameterized CNN, indicating the gains depend on over-parameterization.
Limitations
- Two-stage training adds cost, and V-TST needs more MC samples at inference for best results.
- TST and V-TST can reduce OOD detection performance relative to baselines.
- Gains are not shown for under-parameterized models, where V-TST worsens ECE.
Experimental setup
Datasets: CIFAR10 CIFAR100 SVHN CIFAR10-C CIFAR100-C
Benchmarks: CIFAR10 CIFAR100 SVHN CIFAR10-C CIFAR100-C
Metrics: accuracy ECE MCE NLL AUROC FPR95
Baselines: WRN 28-10 end-to-end Temperature scaling WRN ViT end-to-end CNN end-to-end
Builds on: Wide Residual Networks Vision Transformer temperature scaling variational autoencoder reparameterization
Closest papers by meaning
- Calibration Bottleneck: Over-compressed Representations are Less Calibratable (ICML 2024 · similarity 0.86)
- Balancing Two Classifiers via A Simplex ETF Structure for Model Calibration (CVPR 2025 · similarity 0.85)
- A Benchmark Study on Calibration (ICLR 2024 · similarity 0.85)
- Feature Clipping for Uncertainty Calibration (AAAI 2025 · similarity 0.84)
- Uncertainty Weighted Gradients for Model Calibration (CVPR 2025 · similarity 0.84)
- Towards Calibrated Multi-label Deep Neural Networks (CVPR 2024 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.