Optimal Transport Dropout for Structured Predictive Uncertainty
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.
Related
- Beyond Empirical Support: Structured Outlier Generation via Sinkhorn Optimal Transport
- Dropout Universality: Scaling Laws and Optimal Scheduling at the Edge-of-Chaos
- Bagged Martingale Posteriors: Calibrated Uncertainty Quantification for Predictive Resampling
- Predictive Uncertainty for Neural CAE Surrogates
- Nonnegativity Shapes Predictive Uncertainty: Complementary Model Fits and Quantile Construction
BibTeX
@misc{lorenzon2026optimal,
title = {{Optimal Transport Dropout for Structured Predictive Uncertainty}},
author = {Giacomo Lorenzon and Francesco Regazzoni},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.33377},
eprint = {2609.33377},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.33377},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Optimal Transport Dropout for Structured Predictive Uncertainty
AU - Giacomo Lorenzon
AU - Francesco Regazzoni
PY - 2026
DA - 2026/09/27
PB - arXiv
DO - 10.48550/arxiv.2609.33377
UR - https://arxiv.org/abs/2609.33377
AB - 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.
ER -