Beyond the Mean: Distribution-Aware Loss Functions for Bimodal Regression

TL;DR

This paper introduces distribution-aware loss functions for bimodal regression that effectively model complex predictive uncertainties, outperforming traditional methods and mixture models in stability and fidelity.

Contribution

The authors propose a novel family of loss functions integrating Wasserstein and Cramér distances, enabling deep models to recover bimodal distributions reliably without instability.

Findings

Wasserstein loss achieves a 45% reduction in Jensen-Shannon Divergence on bimodal datasets.

The approach matches MSE in unimodal tasks while improving bimodal distribution modeling.

The method outperforms Mixture Density Networks in fidelity and robustness.

Abstract

Despite the strong predictive performance achieved by machine learning models across many application domains, assessing their trustworthiness through reliable estimates of predictive confidence remains a critical challenge. This issue arises in scenarios where the likelihood of error inferred from learned representations follows a bimodal distribution, resulting from the coexistence of confident and ambiguous predictions. Standard regression approaches often struggle to adequately express this predictive uncertainty, as they implicitly assume unimodal Gaussian noise, leading to mean-collapse behavior in such settings. Although Mixture Density Networks (MDNs) can represent different distributions, they suffer from severe optimization instability. We propose a family of distribution-aware loss functions integrating normalized RMSE with Wasserstein and Cram\'er distances. When applied to standard deep regression models, our approach recovers bimodal distributions without the volatility of mixture models. Validated across four experimental stages, our results show that the proposed Wasserstein loss establishes a new Pareto efficiency frontier: matching the stability of standard regression losses like MSE in unimodal tasks while reducing Jensen-Shannon Divergence by 45% on complex bimodal datasets. Our framework strictly dominates MDNs in both fidelity and robustness, offering a reliable tool for aleatoric uncertainty estimation in trustworthy AI systems.