# Density-Softmax: Efficient Test-time Model for Uncertainty Estimation and Robustness under Distribution Shifts

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

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 60 of its 83 references that are in the map.
- First posted on arXiv 13 Feb 2023.

## Builds on (most cited first, 30 of 60)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [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)
- [Improved Training of Wasserstein GANs](https://mlsanity.com/paper/1704.00028/) (2017)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Spectral Normalization for Generative Adversarial Networks](https://mlsanity.com/paper/1802.05957/) (2018)
- [A Simple Unified Framework for Detecting Out-of-Distribution Samples and Adversarial Attacks](https://mlsanity.com/paper/1807.03888/) (2018)
- [Normalizing Flows for Probabilistic Modeling and Inference](https://mlsanity.com/paper/1912.02762/) (2019)
- [Energy-based Out-of-distribution Detection](https://mlsanity.com/paper/2010.03759/) (2020)
- [Can You Trust Your Model's Uncertainty? Evaluating Predictive Uncertainty Under Dataset Shift](https://mlsanity.com/paper/1906.02530/) (2019)
- [Evidential Deep Learning to Quantify Classification Uncertainty](https://mlsanity.com/paper/1806.01768/) (2018)
- [Deep Neural Networks as Gaussian Processes](https://mlsanity.com/paper/1711.00165/) (2017)
- [GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration](https://mlsanity.com/paper/1809.11165/) (2018)
- [Out-of-Distribution Detection with Deep Nearest Neighbors](https://mlsanity.com/paper/2204.06507/) (2022)
- [A Simple Baseline for Bayesian Uncertainty in Deep Learning](https://mlsanity.com/paper/1902.02476/) (2019)
- [Predictive Uncertainty Estimation via Prior Networks](https://mlsanity.com/paper/1802.10501/) (2018)
- [Weight Uncertainty in Neural Networks](https://mlsanity.com/paper/1505.05424/) (2015)
- [Do Deep Generative Models Know What They Don't Know?](https://mlsanity.com/paper/1810.09136/) (2018)
- [Accurate Uncertainties for Deep Learning Using Calibrated Regression](https://mlsanity.com/paper/1807.00263/) (2018)
- [Lipschitz regularity of deep neural networks: analysis and efficient estimation](https://mlsanity.com/paper/1805.10965/) (2018)
- [ViM: Out-Of-Distribution with Virtual-logit Matching](https://mlsanity.com/paper/2203.10807/) (2022)
- [Revisiting the Calibration of Modern Neural Networks](https://mlsanity.com/paper/2106.07998/) (2021)
- [Simple and Principled Uncertainty Estimation with Deterministic Deep Learning via Distance Awareness](https://mlsanity.com/paper/2006.10108/) (2020)
- [BatchEnsemble: An Alternative Approach to Efficient Ensemble and Lifelong Learning](https://mlsanity.com/paper/2002.06715/) (2020)
- [Regularisation of Neural Networks by Enforcing Lipschitz Continuity](https://mlsanity.com/paper/1804.04368/) (2018)

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

- [Density-Regression: Efficient and Distance-Aware Deep Regressor for Uncertainty Estimation under Distribution Shifts](https://mlsanity.com/paper/2403.05600/) (2024)
- [Variance-Aware Linear UCB with Deep Representation for Neural Contextual Bandits](https://mlsanity.com/paper/2411.05979/) (2024)
- [Disentangling Uncertainties by Learning Compressed Data Representation](https://mlsanity.com/paper/2503.15801/) (2025)
- [Efficient Nearest Neighbor based Uncertainty Estimation for Natural Language Processing Tasks](https://mlsanity.com/paper/2407.02138/) (2024)
- [Source-Free Domain-Invariant Performance Prediction](https://mlsanity.com/paper/2408.02209/) (2024)
- [Q-Learning with Shift-Aware Upper Confidence Bound in Non-Stationary Reinforcement Learning](https://mlsanity.com/paper/2510.03181/) (2025)
- [On the challenges of detecting MCI using EEG in the wild](https://mlsanity.com/paper/2501.17871/) (2025)
- [A Simple and Effective Method for Uncertainty Quantification and OOD Detection](https://mlsanity.com/paper/2508.00754/) (2025)
- [Calibrated Uncertainty Sampling for Active Learning](https://mlsanity.com/paper/2510.03162/) (2025)
- [VOLTA: The Surprising Ineffectiveness of Auxiliary Losses for Calibrated Deep Learning](https://mlsanity.com/paper/2604.08639/) (2026)
- [U-SEG: Uncertainty in SEGmentation -- A systematic multi-variable exploration](https://mlsanity.com/paper/2605.15421/) (2026)
- [Courtroom Analogy: New Perspective on Uncertainty-Aware Classification](https://mlsanity.com/paper/2605.25616/) (2026)

## Abstract (arXiv)

Sampling-based methods, e.g., Deep Ensembles and Bayesian Neural Nets have become promising approaches to improve the quality of uncertainty estimation and robust generalization. However, they suffer from a large model size and high latency at test-time, which limits the scalability needed for low-resource devices and real-time applications. To resolve these computational issues, we propose Density-Softmax, a sampling-free deterministic framework via combining a density function built on a Lipschitz-constrained feature extractor with the softmax layer. Theoretically, we show that our model is the solution of minimax uncertainty risk and is distance-aware on feature space, thus reducing the over-confidence of the standard softmax under distribution shifts. Empirically, our method enjoys competitive results with state-of-the-art techniques in terms of uncertainty and robustness, while having a lower number of model parameters and a lower latency at test-time.