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

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 15 of its 43 references that are in the map.
- First posted on arXiv 20 Mar 2026.

## Builds on (most cited first, 15 of 15)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://mlsanity.com/paper/1905.11946/) (2019)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [An Empirical Study of Example Forgetting during Deep Neural Network Learning](https://mlsanity.com/paper/1812.05159/) (2018)
- [On The Power of Curriculum Learning in Training Deep Networks](https://mlsanity.com/paper/1904.03626/) (2019)
- [Learning with a Wasserstein Loss](https://mlsanity.com/paper/1506.05439/) (2015)
- [Delving into Deep Imbalanced Regression](https://mlsanity.com/paper/2102.09554/) (2021)
- [Deep Learning Through the Lens of Example Difficulty](https://mlsanity.com/paper/2106.09647/) (2021)
- [The Cramer Distance as a Solution to Biased Wasserstein Gradients](https://mlsanity.com/paper/1705.10743/) (2017)
- [Overcoming Limitations of Mixture Density Networks: A Sampling and Fitting Framework for Multimodal Future Prediction](https://mlsanity.com/paper/1906.03631/) (2019)
- [Balanced MSE for Imbalanced Visual Regression](https://mlsanity.com/paper/2203.16427/) (2022)

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

- [Risk-Calibrated Learning: Minimizing Fatal Errors in Medical AI](https://mlsanity.com/paper/2604.12693/) (2026)
- [Improving Model Safety by Targeted Error Correction](https://mlsanity.com/paper/2605.02544/) (2026)

## Abstract (arXiv)

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.