Optimal Transport Dropout for Structured Predictive Uncertainty

Giacomo Lorenzon; Francesco Regazzoni

cs.LGSep 27, 2026

Optimal Transport Dropout for Structured Predictive Uncertainty

Authors: Giacomo Lorenzon, Francesco Regazzoni

Organizations: MOX–Department of Mathematics, Politecnico di Milano, Italy

Abstract

Deterministic neural networks and neural operators provide point predictions with no intrinsic measure of reliability. Yet, predictive uncertainty may stem from irreducible outcome variability, finite data, or limitations of the chosen model class. Monte Carlo dropout offers a computationally convenient way to construct a predictive distribution through stochastic feature masking, without training multiple independent networks or explicitly inferring a posterior over model parameters. However, its perturbation law is largely prescribed a priori and typically factorised across latent coordinates. We introduce Optimal Transport Dropout (OTD), which instead learns the predictive mapping and the law of its latent perturbations jointly. Starting from a simple independent reference distribution, OTD transports latent perturbations through a learnable flow and propagates them through the predictive neural network, thereby inducing a structured predictive law. Training uses the strictly proper Energy Score, while a kinetic-action term geometrically regularises the transport. Synthetic benchmarks show that OTD captures multimodal predictive distributions, generates meaningful dispersion when the model is misspecified, and exhibits contracting dispersion as more training data or greater model capacity are provided. For a field-valued partial differential equation surrogate, predictive dispersion strongly aligns with the spatial pattern of prediction errors. On this task, compared with Monte Carlo dropout, OTD yields more accurate predictions and better-calibrated, substantially narrower intervals. On real-world regression benchmarks, it further shows competitive accuracy and better probabilistic predictions compared to established baselines. OTD therefore offers a way to learn structured predictive uncertainty without explicit posterior inference or ensembles of independently trained predictors.

Figures & tables

Figure 1 : Bimodal aleatoric experiment. Left: kernel density estimation of the predictive conditional distribution. Right: principal component analysis of the first two components of the sampled and the transported masks.

Figure 3 : From under- to over-parametrisation experiment. The four images show the predictions of the trained models in blue with increasing hidden widths

M

. Above them, in grey, the histograms represent the distribution of the transported mask values across hidden units.

	Predictive performance				PICP % (Width)
Method	RMSE	MAE	$S_{\text{ES}}$	MACE	50%	75%	80%	90%	95%
Deterministic	0.163	0.070	0.070	–	–	–	–	–	–
Dropout	0.172	0.102	0.071	0.016	50.5 (0.068)	74.0 (0.115)	78.4 (0.128)	87.5 (0.164)	92.7 (0.193)
OTD	0.128	0.044	0.031	0.009	50.9 (0.025)	76.7 (0.044)	81.5 (0.050)	90.5 (0.066)	94.9 (0.081)

Table 2 : Ginzburg–Landau equilibrium fields experiment. Predictive performance and marginal calibration. Interval entries report PICP in per cent, with relative width in parentheses.

Panel of Figure 4 — Figure 4 : Ginzburg–Landau equilibrium fields experiment. OTD ensemble mean, local predictive discrepancy, and pairwise dispersion for two selected test fields. These maps illustrate the spatial alignment between predictive discrepancy $\epsilon_{F}$ and predictive dispersion $\epsilon_{D}$ .

Model	RMSE	MAE	$S_{\mathrm{ES}}$	MACE
MC Dropout	$0.1992\pm 0.0158$	$0.1228\pm 0.0088$	$0.0885\pm 0.0075$	$0.0381\pm 0.0045$
MC Dropout + ES	$0.1902\pm 0.0133$	$0.1083\pm 0.0067$	$0.0812\pm 0.0056$	$0.1657\pm 0.0106$
OTD-0	$0.2070\pm 0.0077$	$0.1215\pm 0.0019$	$0.0896\pm 0.0012$	$0.1665\pm 0.0018$
OTD, $\lambda_{\mathrm{kin}}=0$	$\mathbf{0.1442\pm 0.0141}$	$0.0668\pm 0.0037$	$0.0460\pm 0.0028$	$0.0536\pm 0.0273$
OTD	$0.1497\pm 0.0152$	$\mathbf{0.0661\pm 0.0084}$	$\mathbf{0.0455\pm 0.0053}$	$\mathbf{0.0328\pm 0.0225}$

Table 3 : Ginzburg–Landau equilibrium fields: controlled comparison of the Energy Score, transport flow, and kinetic regularisation. Results are mean

\pm

standard deviation over three seeds.

Figure 5 : Experimental dataset. Negative log-likelihood comparison of OTD (blue) against other models (grey).

Appendix figures & tables22 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6 : Bimodal aleatoric experiment. Empirical correlation of the latent variables transported by the flow in the bimodal dataset experiment.

Latent mask
Hyperparameter	Value
Concrete parameter $p$	$0.5$
Concrete temperature $\tau$	$1.0$
Latent flow
Velocity hidden widths	$(4,4)$
Activation	GeLU

Appendix

Table 4 : Bimodal aleatoric experiment. List of hyperparameters.

Latent mask
Hyperparameter	Value
Concrete parameter $p$	$0.5$
Concrete temperature $\tau$	$1.0$
Latent flow
Velocity hidden widths	$(2,2)$
Activation	GeLU

Appendix

Table 5 : Reduced-Order Model misspecification experiment. List of hyperparameters.

Latent mask
Hyperparameter	Value
Mask dimension	$4,8,16,32$
Concrete parameter $p$	$0.5$
Concrete temperature $\tau$	$1.0$
Latent flow
Velocity hidden widths	$(4)$

Appendix

Table 6 : From under- to over-parametrisation experiment. List of hyperparameters.

Hyperparameter	Value
Hidden-layer depth	$5$
Reference hidden width	$128$
Deterministic hidden widths	$[128]\times 5$
Vanilla-dropout hidden widths	$[142]\times 5$
OT-dropout hidden widths	$[256]\times 5$
Activation	GeLU

Appendix

Table 7 : Ginzburg–Landau equilibrium fields experiment. Experimental settings.

Figure 7 : Ginzburg–Landau equilibrium fields experiment. Equilibrium field examples for 16 values of the input parameter $\mu$ . As can be seen, the map $\mu\mapsto u_{\mu}$ is highly non-linear, albeit continuous.

Sample	$\epsilon_{F}$	$\epsilon_{D}$	$S_{\text{ES}}$	Mean error	Std. dev.
00000	0.0242	0.0253	0.0116	0.0167	0.0268
00011	0.0196	0.0232	0.0080	0.0114	0.0243
00036	0.0192	0.0229	0.0077	0.0106	0.0228
00050	0.0257	0.0277	0.0118	0.0167	0.0274
00058	0.0208	0.0253	0.0082	0.0109	0.0265
00077	0.0257	0.0288	0.0113	0.0171	0.0286

Appendix

Table 8 : Per-sample results for the Ginzburg–Landau experiment.

Figure 8 : Ginzburg–Landau equilibrium fields experiment, sample 3. From top to bottom: ground truth, deterministic predictor, MC Dropout, OTD.

Figure 9 : Ginzburg–Landau equilibrium fields experiment, sample 6. From top to bottom: ground truth, deterministic predictor, MC Dropout, OTD.

Figure 10 : Ginzburg–Landau equilibrium fields experiment, sample 9. From top to bottom: ground truth, deterministic predictor, MC Dropout, OTD.

Figure 11 : Ginzburg–Landau equilibrium fields experiment, sample 12. From top to bottom: ground truth, deterministic predictor, MC Dropout, OTD.

Panel of Figure 12 — Figure 12 : Ginzburg–Landau equilibrium fields experiment. Spatial alignment between the first ( $\epsilon_{F}$ ) and the second term ( $\epsilon_{D}$ ) of the Energy score.

	$\varepsilon_{F}$	$\varepsilon_{D}$	Mean error	Std. dev.
$\varepsilon_{F}$	1	0.9765	–	0.9735
$\varepsilon_{D}$	0.9765	1	0.9441	–
Mean error	–	0.9441	1	0.9294
Std. dev.	0.9735	–	0.9294	1

Appendix

Table 9 : Ginzburg–Landau equilibrium fields experiment. Pairwise Spearman rank correlations.

Sample	$\rho(\varepsilon_{F},\varepsilon_{D})$	$\rho(e,\sigma)$	$\rho(\varepsilon_{F},\sigma)$	$\rho(e,\varepsilon_{D})$
00000	0.6290	0.1138	0.5881	0.1264
00011	0.6909	-0.0576	0.6559	-0.0404
00036	0.7722	0.1702	0.7455	0.1886
00050	0.8393	0.3215	0.8049	0.3426
00058	0.8659	0.3631	0.8326	0.3810
00077	0.9168	0.4516	0.8773	0.4826

Appendix

Table 10 : Ginzburg–Landau equilibrium fields experiment. Per-case pixel-wise Spearman rank correlations. The largest correlation in each row is shown in bold.

	$\varepsilon_{F}$	$\varepsilon_{D}$	Mean error	Std. dev.
$\varepsilon_{F}$	1	0.9263	–	0.9095
$\varepsilon_{D}$	0.9263	1	0.6085	–
Mean error	–	0.6085	1	0.5926
Std. dev.	0.9095	–	0.5926	1

Appendix

Table 11 : Pairwise Spearman rank correlations across samples.

$L$	RMSE	MAE	$S_{\mathrm{ES}}$	MACE
$0$	$0.1704\pm 0.0138$	$0.0949\pm 0.0074$	$0.0698\pm 0.0027$	$0.1736\pm 0.0132$
$1$	$0.1417\pm 0.0188$	$0.0555\pm 0.0086$	$0.0399\pm 0.0071$	$0.0367\pm 0.0309$
$2$	$0.1392\pm 0.0199$	$0.0542\pm 0.0065$	$0.0386\pm 0.0053$	$0.0255\pm 0.0095$
$4$	$\mathbf{0.1368\pm 0.0098}$	$\mathbf{0.0533\pm 0.0043}$	$\mathbf{0.0377\pm 0.0028}$	$\mathbf{0.0105\pm 0.0072}$

Appendix

Table 12 : Sensitivity to the number of integration steps at

\tau=1

. Results are mean

\pm

standard deviation over three seeds.

$\tau$	RMSE	MAE	$S_{\mathrm{ES}}$	MACE
$0.05$	$0.1412\pm 0.0133$	$0.0588\pm 0.0044$	$0.0422\pm 0.0029$	$0.0458\pm 0.0196$
$0.10$	$0.1400\pm 0.0048$	$0.0601\pm 0.0013$	$0.0422\pm 0.0014$	$0.0453\pm 0.0447$
$0.50$	$\mathbf{0.1372\pm 0.0245}$	$0.0546\pm 0.0084$	$0.0387\pm 0.0068$	$0.0344\pm 0.0146$
$1.00$	$0.1392\pm 0.0199$	$\mathbf{0.0542\pm 0.0065}$	$\mathbf{0.0386\pm 0.0053}$	$\mathbf{0.0255\pm 0.0095}$

Appendix

Table 13 : Sensitivity to the Concrete temperature using two integration steps. Results are mean

\pm

standard deviation over three seeds.

Dataset	CADET	CART-h	CART-k	CKDE	LSCDE	NKDE	CDE-HT	LinCDE	MDN	NF	CPFN	OTD
energy	3.55	3.09	3.06	2.47	3.38	3.00	2.93	2.93	2.78	2.86	2.73	1.08
synchronous	-2.93	-1.63	-1.86	-3.59	-1.25	-1.57	-2.11	-1.85	-2.94	-2.64	-4.48	-5.02
localization	-0.23	-0.55	-0.01	-0.26	-0.61	-0.28	-0.66	-0.95	-0.68	-0.43	4.19	1.03
toxicity	1.80	1.50	1.38	1.32	1.34	1.55	1.53	1.29	1.24	1.23	2.13	1.40
concrete	4.17	3.75	3.93	3.32	3.66	3.91	3.72	3.47	2.97	3.18	3.40	3.07
slump	3.42	3.55	3.43	2.35	2.91	3.08	3.34	2.98	2.23	2.39	4.48	2.40

Appendix

Table 14 : Experimental dataset. Results reported as mean NLL. The better, the darker.

Dataset	RMSE	$S_{\text{ES}}$	Sharpness	MACE
energy	0.857	0.464	1.231	0.323
synchronous	0.001	0.001	0.006	0.064
localization	0.829	0.395	1.866	0.142
toxicity	1.028	0.569	1.605	0.278
concrete	5.172	2.965	6.830	0.355
slump	2.345	1.613	0.721	0.690

Appendix

Table 15 : Experimental dataset. Summary of the OTD results across datasets. MACE is computed against nominal coverage levels

\alpha\in\{0.50,0.80,0.90,0.95\}

Figure 13 : Role of the kinetic regularisation. Grey dots denote the reference masks, while blue dots denote the transported masks. Larger values of $\lambda_{\mathrm{kin}}$ produce smoother transport trajectories but increasingly restrict the deformation of the reference measure.

Figure 14 : OTD timing-sweep configurations for the Ginzburg–Landau experiment. All variants were trained using

50\%

of the training cases and batch size

8192

. Timings are reported as mean

\pm

standard deviation over three seeds. Each run used

5

warmup epochs, excluded from timing, followed by

1000

measured training epochs.

Method	$d_{z}$	RMSE	MAE	$S_{\text{ES}}$	MACE	Sharpness	$\mathrm{PICP}_{0.90}$
OTD	1536	$\mathbf{0.134\pm 0.003}$	$\mathbf{0.043\pm 0.000}$	$\mathbf{0.031\pm 0.000}$	$\mathbf{0.038\pm 0.020}$	$0.046\pm 0.005$	$\mathbf{0.911\pm 0.024}$
Low-latent OTD	512	$0.140\pm 0.003$	$0.045\pm 0.001$	$0.033\pm 0.001$	$0.086\pm 0.038$	$0.039\pm 0.004$	$0.811\pm 0.031$
Low-latent OTD	128	$0.141\pm 0.021$	$0.045\pm 0.006$	$0.033\pm 0.005$	$0.067\pm 0.050$	$0.040\pm 0.008$	$0.829\pm 0.047$
Low-latent OTD	32	$0.143\pm 0.008$	$0.045\pm 0.002$	$0.034\pm 0.001$	$0.087\pm 0.066$	$\mathbf{0.037\pm 0.000}$	$0.809\pm 0.067$
Conditional OT	1536	$0.169\pm 0.007$	$0.078\pm 0.005$	$0.050\pm 0.002$	$0.130\pm 0.031$	$0.142\pm 0.003$	$0.968\pm 0.014$
Conditional OT	512	$0.156\pm 0.013$	$0.059\pm 0.011$	$0.041\pm 0.006$	$0.076\pm 0.013$	$0.089\pm 0.021$	$0.937\pm 0.002$

Appendix

Table 17 : Ginzburg–Landau results comparing the original OTD experiment with the low-dimensional OTD variants and Conditional OT variants. RMSE, MAE,

S_{\text{ES}}

, MACE, sharpness, and

\mathrm{PICP}_{0.90}

are evaluated at the pixel level on a

64\times 64

grid. MACE is averaged over

\alpha\in\{0.50,0.75,0.80,0.90,0.95\}

Explore similar work

May 1, 2026cs.LG

Possibilistic Predictive Uncertainty for Deep Learning

Deep neural networks achieve impressive results across diverse applications, yet their overconfidence on unseen inputs necessitates reliable epistemic uncertainty modeling. Existing methods for uncertainty modeling face a fundamental dilemma: Bayesian approaches provide principled estimates but remain computationally prohibitive, while efficient second-order predictors lack rigorous connections between their specific objectives and epistemic uncertainty quantification. To resolve this dilemma, we introduce Dirichlet-approximated possibilistic posterior predictions (DAPPr), a principled framework grounded in possibility theory. We define a possibilistic posterior over parameters, project it to the prediction space via supremum operators, and approximate the projected posterior using learnable Dirichlet possibility functions. This projection-and-approximation strategy yields a simple training objective with closed-form solutions. Despite its simplicity, extensive experiments across diverse benchmarks show that DAPPr achieves competitive or superior uncertainty quantification performance over state-of-the-art second-order predictors while maintaining both principled derivation and computational efficiency. Code is available at https://github.com/MaxwellYaoNi/DAPPr.

Yao Ni, Jeremie Houssineau, Yew-Soon Ong +1

Nanyang Technological University · A*STAR · The University of New South Wales +1

Predictive Uncertainty Posterior Predictive Distribution Bayesian Neural Networks Dynamic Epistemic Logic Deep Learning Posterior

Ask PDF

May 22, 2026stat.ML

Dirichlet-Based Monte Carlo Dropout for Uncertainty Estimation in Neural Networks

Traditional neural networks provide deterministic predictions without inherent uncertainty estimates. While Bayesian Neural Networks (BNNs) offer a principled approach to uncertainty quantification, their computational complexity limits scalability. Monte Carlo (MC) Dropout, initially introduced as a regularization technique, has been shown to approximate Bayesian inference by enabling probabilistic modeling through multiple stochastic forward passes. In this work, we enhance uncertainty estimation in deep learning by integrating a Dirichlet-based framework within MC Dropout. Specifically, we leverage the formulation proposed by Sensoy et al. (2018), where class probabilities are modeled using a Dirichlet distribution, allowing for a more informative uncertainty representation. The proposed approach maintains the computational efficiency of MC Dropout while improving the quality of uncertainty estimates. We discuss the theoretical foundations of our method and compare it with existing uncertainty quantification techniques. The results highlight the effectiveness of the proposed method in producing well-calibrated uncertainty estimates, offering a practical solution for uncertainty-aware deep learning models.

Rouaa Hoblos, Noura Dridi, Noureddine Zerhouni +1

FEMTO-ST · Universit´e Marie et Louis Pasteur, SUPMICROTECH, CNRS, Institut FEMTO-ST, F-25000 Besan¸con, France

Uncertainty-Aware Classification Uncertainty Quantification Bayesian Neural Networks Uncertainty Dropout Neural Network

Ask PDF

Nov 19, 2025stat.ML

Beyond Uncertainty Sets: Leveraging Optimal Transport to Extend Conformal Predictive Distributions to Multivariate Settings

Conformal prediction (CP) constructs uncertainty sets for model outputs with finite-sample coverage guarantees. Yet ranking scores is straightforward only when they are scalar-valued, limiting CP to real-valued scores or ad-hoc one-dimensional reductions. Vector-valued scores arise naturally in multi-output regression and model aggregation, where each predictor in an ensemble provides its own score. Optimal transport (OT) defines vector ranks and center-outward multivariate quantile regions, though generally with asymptotic coverage guarantees. Applying a fixed transport map learned from calibration data to a new point introduces an uncontrolled approximation error. We restore finite-sample, distribution-free coverage by conformalizing vector-valued OT quantile regions. Each candidate's rank is defined by transporting the calibration scores augmented with that candidate's score, preserving the symmetry needed for validity. This appears to require a continuum of OT problems. However, we prove that the optimal assignment is piecewise constant across a fixed polyhedral partition of score space. This lets us characterize the entire prediction set in

O(n^3)

time, matching the cost of a single assignment solve. It also addresses a limitation of prediction sets: they indicate which outcomes are plausible, but not their relative likelihood. In one dimension, conformal predictive distributions (CPDs) fill this gap by producing a predictive distribution with finite-sample calibration. Extending CPDs beyond one dimension remained an open problem. We construct, to our knowledge, the first multivariate CPDs with finite-sample calibration: a center-outward predictive distribution whose derived uncertainty regions have conformal coverage. We present both conservative and exact randomized versions; the latter generalizes the classical Dempster-Hill procedure.

Eugene Ndiaye

Apple

Online Conformal Prediction Distributionally-Robust Optimization

Ask PDF