Research map / Probabilistic and scientific modeling / Uncertainty and Gaussian processes
Density-Softmax: Efficient Test-time Model for Uncertainty Estimation and Robustness under Distribution Shifts
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 Density-Softmax, a sampling-free single-model method that scales softmax logits by a normalizing-flow density of Lipschitz-constrained features, improving uncertainty and robustness at near-deterministic test-time cost.
Problem
Sampling-based uncertainty methods (Deep Ensembles, BNNs) give good uncertainty and robustness but are large and slow at test time, while deterministic models are efficient but overconfident under distribution shift.
Approach
Combine a gradient-penalty 1-Lipschitz feature extractor with a normalizing-flow density on the feature space and multiply the classifier logits by the scaled likelihood before softmax, so that predictions flatten to uniform for far-from-training features in a single forward pass.
Key results
- On CIFAR-10-C with Wide ResNet-28-10, Density-Softmax reaches cNLL 0.68 and cAcc 79.2%, the best among methods, with 36.58M parameters versus 145.99M for Deep Ensembles.
- Latency on RTX A5000 is about 520 ms/sample for WRN-28-10 and 300 ms for ResNet-50, nearly identical to deterministic ERM and far below Deep Ensembles (~1520 ms).
- On CIFAR-10 IID test, NLL is 0.137 versus 0.114 for Deep Ensembles, so it does not beat ensembles on in-distribution NLL.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| remove 1-Lipschitz regularization (ablation) | Density-Softmax · CIFAR-10-C defocus_blur_5 | accuracy: 91.2% -> 78.6% | Hurt |
| add the density function p(Z;α) (ablation) | Density-Softmax · CIFAR-10 | ECE: 0.03 -> 0.01 | Helped |
| add 1-Lipschitz regularization (λ=1e-4) (ablation) | Density-Softmax · CIFAR-10 | ECE | Helped |
Surprising findings
- Removing the Lipschitz constraint (lambda=0) worsened calibration ECE, suggesting the 1-Lipschitz regularizer helps uncertainty as well as robustness.
- On CIFAR-10 IID the method's NLL (0.137) is worse than Rank-1 BNN (0.128) and Deep Ensembles (0.114).
Limitations
- Uncertainty quality depends on the quality of the density model; density estimation can be non-trivial in practice.
- Training is longer than deterministic ERM because the Jacobian penalty and multi-stage training are needed.
- A trade-off exists between IID accuracy and OOD robustness under the Lipschitz constraint.
Experimental setup
Datasets: CIFAR-10 CIFAR-100 ImageNet CIFAR-10-C CIFAR-100-C ImageNet-C CIFAR-10.1 SVHN two moons toy dataset
Benchmarks: CIFAR-10 CIFAR-10-C CIFAR-10.1 CIFAR-100 CIFAR-100-C ImageNet ImageNet-C SVHN two moons
Metrics: NLL accuracy ECE AUPR latency parameter count
Baselines: Deterministic ERM Rank-1 BNN Heteroscedastic SNGP DDU (w/o TS) NatPN BatchEnsemble Deep Ensembles MIMO
Builds on: Deep Ensembles DDU Spectral normalization / gradient penalty Lipschitz networks Brier score minimax risk
Closest papers by meaning
- Deep Deterministic Uncertainty: A New Simple Baseline (CVPR 2023 · similarity 0.88)
- Uncertainty Estimation Using a Single Deep Deterministic Neural Network (ICML 2020 · similarity 0.88)
- Training-Free Uncertainty Estimation for Dense Regression: Sensitivity as a Surrogate (AAAI 2022 · similarity 0.88)
- Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles (NeurIPS 2017 · similarity 0.87)
- Adaptive Uncertainty Estimation via High-Dimensional Testing on Latent Representations (NeurIPS 2023 · similarity 0.87)
- SURE: SUrvey REcipes for building reliable and robust deep networks (CVPR 2024 · similarity 0.85)
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.