Uncertainty Estimation via Bayesian Decision Layer with Wasserstein Geometric Matching

Reliable uncertainty estimation should distinguish predictive confidence from the degree of compatibility between an input and the data distributions associated with different classes. Existing approaches typically characterize uncertainty through predictive probabilities, posterior variability, or distances to class structures, but do not jointly model these aspects. Probabilistic methods do not necessarily quantify compatibility between input- and class-level distributions, whereas distance-aware methods commonly operate on point representations. To address these challenges, we introduce the Wasserstein Generalized Bayesian Decision Layer (W-GBDL), which derives uncertainty from the compatibility between an input-conditioned distribution and learnable class reference distributions. W-GBDL represents each input and each class with variational distributions and measures their compatibility using the Wasserstein distance. This geometry provides a unified basis for training and minimum-risk prediction, while the entropy of the resulting class probabilities captures predictive ambiguity. Experiments across image and language tasks show improved uncertainty estimation and calibration while maintaining competitive classification performance.