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

TL;DR

Density-Softmax is a deterministic, sampling-free model that improves uncertainty estimation and robustness under distribution shifts, with lower computational cost compared to traditional sampling-based methods.

Contribution

It introduces a novel density-softmax framework combining a Lipschitz-constrained feature extractor with softmax, reducing over-confidence and computational costs.

Findings

Achieves competitive uncertainty and robustness results.

Has fewer parameters and lower test-time latency.

Theoretically minimizes uncertainty risk and is distance-aware.

Abstract

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.