Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Abolfazl Mohammadi-Seif, Carlos Soares, Rita P. Ribeiro, Ricardo Baeza-Yates

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ér 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.

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

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ér 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.
Paper Structure (31 sections, 6 equations, 7 figures, 6 tables)

This paper contains 31 sections, 6 equations, 7 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Aleatoric Uncertainty and Mixture Models
  4. Deep Imbalanced and Non-Parametric Regression
  5. Application: Complexity and Difficulty Assessment
  6. Methodology
  7. Component 1: Pointwise Fidelity
  8. Component 2: Distributional Fidelity
  9. Component 3: Range Alignment
  10. Scale Invariance and Final Loss
  11. Experimental Setup
  12. Datasets
  13. Synthetic and Real-World Data (Stages I--III)
  14. Image Data Generation (Stage IV)
  15. Data Preprocessing
  16. ...and 16 more sections

Figures (7)

  • Figure 1: Bimodal distribution of targets (a). Standard regression models trained with MSE fail to capture the underlying bimodal distribution (b).
  • Figure 2: Qualitative comparison of predictive distributions (Orange) versus Ground Truth (Blue) on the Two Path dataset.
  • Figure 3: Evolution of predictive densities on Bike Sharing (Hour). Blue Area: Ground Truth. Orange Area: Predicted Distribution.
  • Figure 4: Qualitative comparison on the Protein Structure dataset. Top Row: MSE and HMLP collapse to unimodal averages in sparse regions, while MDN captures bimodality but exhibits disjoint peaks. Bottom Row: MLPQ and Proposed method (Center) provide a strong approximation. Notably, increasing to $\alpha=7$ (Right) results in high coverage of the bimodal support.
  • Figure 5: Distributional Fidelity (JS Divergence) across all Image Datasets.
  • ...and 2 more figures