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

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

cs.LGMar 20, 2026

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

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

Organizations: Department of Engineering, Universitat Pompeu Fabra, Roc Boronat St., Barcelona, 08018, Barcelona, Spain. · Department of Engineering, University of Porto, Praça de Gomes Teixeira, Porto, 4099-002, Portugal. · Faculty of Sciences, University of Porto, 4169-007 Porto, Portugal. · INESC TEC, 4200-465 Porto, Portugal.

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.

Figures & tables

Panel of Figure 1 — Figure 1: Bimodal distribution of targets (a). Standard regression models trained with MSE fail to capture the underlying bimodal distribution (b).

ID	Dataset	Type	Target Modality
Stage I: Synthetic Proof of Concept
1	Inverse Square	Synthetic Tabular	Conditional Bimodal
2	Two Path	Synthetic Tabular	Conditional Bimodal
Stage II: Controlled Separation (Unimodal $\to$ Bimodal)
3	Airfoil	Real-World Tabular	Unimodal*
4	Bike Sharing (Hour)	Real-World Tabular	Unimodal*

Table 1: Summary of the 17 datasets used across the four-stage experimental protocol.

Panel of Figure 2 — Figure 2: Qualitative comparison of predictive distributions (Orange) versus Ground Truth (Blue) on the Two Path dataset.

Configuration	Test Loss	RMSE $\downarrow$	Wasserstein $\downarrow$	JS Div $\downarrow$	$\Delta_{BC}$ $\downarrow$
Wasser-Simple	0.535	2.62	1.80	0.63	0.29
Wasser-Range	0.581	2.90	1.71	0.59	0.34
Cramer-Simple	0.542	2.77	1.73	0.62	0.14
Cramer-Range	0.576	2.94	1.38	0.54	0.39
Wasser-Simple (Def)	0.777	3.65	0.30	0.37	0.03
MDN (NLL)	0.071	3.82	1.61	0.54	0.15

Table 2: Performance comparison on synthetic datasets.

	Unimodal ( $S=0$ )		Transition ( $S=0.5$ )		Bimodal ( $S=1$ )
Configuration	RMSE	JS Div	RMSE	JS Div	RMSE	JS Div	$\Delta_{BC}$
Proposed Methods
Wasser-Simple	0.939	0.427	0.467	0.387	0.509	0.441	0.105
Wasser-Range	0.977	0.430	0.491	0.366	0.532	0.436	0.171
Cramer-Simple	0.951	0.457	0.481	0.392	0.501	0.472	0.129
Cramer-Range	0.988	0.440	0.480	0.378	0.518	0.441	0.156

Table 3: Performance snapshots at critical separation intervals (

S

) in synthetic datasets.

Panel of Figure 3 — Figure 3: Evolution of predictive densities on Bike Sharing (Hour) . Blue Area: Ground Truth. Orange Area: Predicted Distribution.

Configuration	Test Loss	RMSE $\downarrow$	Wasserstein $\downarrow$	JS Div $\downarrow$	$\Delta_{BC}$ $\downarrow$
Wasser-Simple	0.232	0.330	0.083	0.306	0.045
Wasser-Range	0.251	0.342	0.079	0.300	0.112
Cramer-Simple	0.222	0.324	0.085	0.306	0.046
Cramer-Range	0.249	0.344	0.101	0.329	0.071
Wasser-Simple (Def)	0.342	0.367	0.081	0.291	0.018
MDN (NLL)	-0.246	0.525	0.084	0.300	0.033

Table 4: Aggregated performance on Real-World Bimodal Datasets.

Panel of Figure 4 — 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.

Configuration	Test Loss	RMSE $\downarrow$	Wasserstein $\downarrow$	JS Div $\downarrow$	$\Delta_{BC}$ $\downarrow$
Wasser-Simple	0.092	0.088	0.052	0.402	0.060
Wasser-Range	0.109	0.088	0.058	0.434	0.032
Cramer-Simple	0.099	0.108	0.071	0.437	0.061
Cramer-Range	0.131	0.118	0.086	0.413	0.040
Wasser-Simple (Def)	0.237	0.153	0.112	0.452	0.044
MDN (NLL)	-2.105	0.069	0.024	0.338	0.046

Table 5: Performance on the Trimodal Bike Sharing dataset.

Complexity	Configuration	RMSE $\downarrow$	Wasserstein $\downarrow$	JS Div. $\downarrow$	$\Delta_{BC}$ $\downarrow$
Unimodal	Wasser-Simple	1.050	0.254	0.441	0.151
	Wasser-Range	1.060	0.253	0.471	0.134
	Cramer-Simple	1.036	0.280	0.536	0.146
	Cramer-Range	1.068	0.251	0.506	0.149
	Wasser-Simple (Def)	1.089	0.172	0.157	0.147
	MDN	1.066	0.261	0.218	0.040

Table 6: Comprehensive performance on Image Complexity Assessment.

Figure 5: Distributional Fidelity (JS Divergence) across all Image Datasets.

Figure 6: Pareto Efficiency Frontier (stage IV Bimodal Datasets).

Panel of Figure 7 — Figure 7: Qualitative comparison on CIFAR-100 .

Aleatoric Uncertainty

Aleatoric uncertainty is irreducible variability in outcomes, arising from noisy measurements, ambiguous inputs, or stochastic processes.

See all 44 papers

Efficient Analytic Uncertainty Quantification for Multi-Modal Regression
Jun 23, 2026Kun Jin, James Harrison, Jiawei Li +8Quantile Regression Uncertainty Quantification
1Google · 2Google DeepMind
Reinforcement Learning over Predictive Distributions for LLM Regression
May 20, 2026Jungsoo Park, Hyungjoo Chae, Ethan Mendes +4Distributional RL RL for Language Models
Georgia Institute of Technology · Allen Institute for AI
Error Aware Distribution Prediction for Lightweight Implicit Neural Representations
Jul 11, 2026Zhimin Li, Jake D. Balla, Joshua A. LevineImplicit Neural Representations Uncertainty Quantification
Vanderbilt University · University of Arizona

Jun 23, 2026J/K move · Enter open · S save

Efficient Analytic Uncertainty Quantification for Multi-Modal Regression

Kun Jin, James Harrison, Jiawei Li, Sihan Liu+7

1Google · 2Google DeepMind

Quantile Regression Uncertainty Quantification Multimodal Regression Variational Inference

Open paper AskPDF arXiv

Efficient uncertainty quantification (UQ) is essential for trustworthy large-scale learning. Existing UQ methods for regression tasks mainly operate under the assumption that the conditional label marginal satisfies single-peak parametric models, e.g., Gaussians, where the negative log-likelihood function simplifies to the mean square error. However, such single-peak assumptions fail in regression tasks featuring multi-modal distributions. On the other hand, semi-parametric methods which achieve strong regression performance for multi-modal distributions often lack efficient quantification on their prediction variances. In this work, we extend UQ techniques based on Variational Bayesian Inference (VBI) to two widely used semi-parametric regression models that yield histogram-like reconstructions of the conditional label densities: Quantile Regression (QR) and Classification Restoration (CR). Our approach introduces a unified, distribution-agnostic framework that simultaneously achieves accurate estimation of complex conditional distributions and highly efficient UQ. Theoretically, our method is grounded in novel formulations of QR and CR within the VBI framework, yielding analytic Evidence Lower Bounds (ELBO) to streamline training and a closed-form or analytically approximated predictive density for efficient inference. Empirically, we evaluate our methods on three large-scale regression benchmarks with multi-modal label distributions. Our framework outperforms state-of-the-art multi-modal regression baselines, and even matches predictive performance of computationally expensive ensemble models. Furthermore, by leveraging epistemic uncertainty estimation, our approach enables highly data-efficient active learning strategies.

Abstract

Aleatoric Uncertainty

Explore similar work

Efficient Analytic Uncertainty Quantification for Multi-Modal Regression

Reinforcement Learning over Predictive Distributions for LLM Regression

Error Aware Distribution Prediction for Lightweight Implicit Neural Representations