Optimal Transport Dropout for Structured Predictive Uncertainty

• Predictive uncertainty may reflect irreducible outcome variability, finite observations, or a model family that cannot represent the target relation exactly.These sources can coexist and are not always separately identifiable; this work models their combined predictive effect.Source: §1 · pp. 1–2
• Scientific surrogates replace expensive simulations with cheaper predictions, but inaccurate predictions may occur without an explicit indication of failure.Uncertainty estimates are needed to identify predictions that should not be trusted while preserving the surrogate’s computational advantage.Source: §1 · p. 2
• Standard Monte Carlo dropout is lightweight but uses a perturbation law prescribed largely a priori, typically through factorised masks.Flexible learned transformations of reference distributions motivate OTD’s alternative.Source: §1 · p. 2
• OTD transports latent perturbations through a learnable flow and propagates them through the predictive network, jointly learning a structured predictive law.The transported law can represent dependencies across mask coordinates, unlike standard factorised dropout.Source: §1 · p. 2

• OTD transports a reference law into a continuous latent perturbation law and trains the induced predictive distribution with the Energy Score and kinetic-action regularisation.Its objective does not require a likelihood, variational posterior over masks, or trajectory-balance loss.Source: §2 · p. 3
• OTD applies flow-based generation in perturbation space, unlike post-hoc Monge–Kantorovich ranking, which operates on uncertainty scores rather than the law inducing predictions.Source: §2 · p. 3
• Unlike residual fitting or operator-parameter generation, OTD learns correlated feature perturbations and is not tied to a particular neural-operator architecture.The cited alternatives include REEF-GP [28] and GenUQ.Source: §2 · p. 3
• The Energy Score supports sample-based predictive-distribution training without a tractable density, while OTD adds transport discretisation and kinetic-regularisation contributions to error analysis.The passage also motivates complementary spatial diagnostics for field-valued predictions.Source: §2 · p. 3
• Predictive dispersion can mix intrinsic and epistemic components and need not capture model bias or dataset and procedural variability.Accordingly, OTD dispersion is interpreted through reducibility, error-alignment, and misspecification diagnostics rather than identified a priori with epistemic uncertainty.Source: §2 · pp. 3–4

3.1 Optimal Transport Dropout: from independent to transported masks

• A random mask modulates a neural network’s hidden representation, and the resulting prediction is written as Gθ(x, Z).The described single-hidden-layer setup extends to multiple modulated layers.Source: §3.1 · p. 4
• OTD transforms a factorised continuous relaxation of Bernoulli masks through a learnable flow before applying the transported mask to the hidden representation.The transported law is shared across inputs, while the predictive distribution depends on each input and can include dependence across mask coordinates.Source: §3.1 · p. 4

3.2 Learning objective

• The predictive and transport parameters are learned jointly using the Energy Score for the induced predictive law and kinetic action to regularise the latent flow.The data-fitting term constrains perturbations through their effect on predictions; the kinetic action regularises the transport directly.Source: §3.2 · pp. 4–5
• The Energy Score enables training at the predictive-distribution level from samples, and is strictly proper when the output metric is of strong negative type.For deterministic targets, the passage relates the score to the Energy Distance.Source: §3.2 · pp. 4–5
• OTD places stochasticity inside the predictive model and optimises it for the quality of the induced conditional distribution, rather than attaching an uncertainty estimator afterward.Source: §3.2 · p. 5

3.3 Predictive dispersion under model misspecification

• Under misspecification, Energy-Score minimisation targets the best-in-class predictive law rather than necessarily recovering the data-generating law.For deterministic targets, dispersion is interpreted as a candidate signal of finite-data and model-class uncertainty; with complete information and correct specification, the optimal law collapses to the target.Source: §3.3 · p. 6

4.1 Controlled experiments

• OTD reproduces an input-dependent transition from a Dirac measure where target branches coincide to bimodality where they separate.The bimodality is also reflected in latent space.Source: §4.1 · p. 7, Figure 1 · p. 7
• In the misspecified reduced-order setting, OTD achieves lower L1, L2, Energy Score, and MACE than MC Dropout.The comparison is reported for the aggregate predictive metrics in Table 1.Source: Table 1 · p. 7, §4.1 · p. 7, Figure 2 · p. 7
• At fixed sample size, predictive dispersion contracts as hidden width increases once the model becomes sufficiently expressive.Stochasticity persists when representational capacity is insufficient.Source: §4.1 · p. 8, Figure 3 · p. 8

4.2 Functional uncertainty quantification

• As training data increase, ensemble dispersion decreases while accuracy improves; spatial dispersion also rank-aligns with discrepancy, assigning more uncertainty to larger-error regions.Figure 4 illustrates the spatial alignment for two selected test fields.Source: §4.2 · pp. 8–9, Figure 4 · p. 9
• OTD achieves better accuracy and probabilistic performance than Dropout, with nearly nominal 90% coverage and substantially sharper intervals.Table 2 reports predictive performance and marginal calibration; interval entries give PICP and relative width.Source: §4.2 · p. 8, Table 2 · p. 9

4.3 Regression benchmarks on real-world data

• Across the considered datasets, OTD remains competitive while improving the quality of the returned predictive information.The aggregate results indicate that this does not require an ensemble of independently trained networks.Source: §4.3 · p. 10

Discussion and limitations

• Predictive dispersion may reflect systematic surrogate error as well as uncertainty from limited data, so it should not be read as a separate estimate of either source.Even for a deterministic target, Energy Score minimisation can favor non-zero dispersion when the predictive mean is inaccurate.Source: §5 · pp. 10–11
• OTD dispersion does not explicitly separate aleatoric, model-class, and finite-data contributions, and it does not define a Bayesian posterior.The uncertainty arises from internal perturbations learned jointly with the predictor.Source: §5 · p. 11
• OTD adds computational cost through predictive sampling and latent-flow integration, and its direct Monte Carlo Energy Score scales quadratically with sample count.The authors report stable performance with few samples and integration steps, but validation on larger architectures and output spaces remains necessary.Source: §5 · p. 11
• OTD offers a way to enrich stochastic predictors beyond independent perturbations, combining latent transport, proper scoring rules, and kinetic regularisation.The authors describe a variational connection with stochastic control and optimal transport as ongoing work.Source: §5 · pp. 10–11

AI use statement

• LLMs assisted with retrieval and discovery of potentially relevant related work, which the authors subsequently read and verified.They also used LLMs for minor language editing to improve grammar and readability.Source: AI use statement · p. 11
• The authors state that they independently conceived, carried out, and verified the research ideas, theoretical results and proofs, and experiments.They report that LLMs were not used to generate ideas, derive results, conduct experiments, or analyse data.Source: AI use statement · pp. 11–12
• The authors take full responsibility for the paper’s content.This statement follows their account of the limited LLM uses.Source: AI use statement · p. 12

A. Diagnostic metrics

• RMSE and MAE assess the accuracy of predictive means, but neither establishes that predictive dispersion or distributional shape is calibrated.RMSE is particularly sensitive to large discrepancies, whereas MAE penalises errors linearly and is less sensitive to large errors and outliers.Source: A. Diagnostic metrics · p. 15
• The Energy Score evaluates predictive samples against observations while accounting for the predictive distribution’s internal dispersion; lower values indicate better predictive performance.Aggregated scores are averaged over the test set.Source: A. Diagnostic metrics · p. 15
• PICP measures empirical coverage of central predictive intervals, with calibration indicated when PICP(α) is approximately α.Coverage alone does not assess interval sharpness, so it should be paired with interval width.Source: A. Diagnostic metrics · p. 16
• Sharpness is quantified by mean central-interval width, where smaller widths are desirable only alongside adequate coverage.Narrow intervals alone do not establish calibrated uncertainty.Source: A. Diagnostic metrics · p. 16
• MACE measures average absolute discrepancies between empirical and nominal coverage across the selected levels, with zero denoting perfect calibration.The metrics capture distinct aspects of predictive performance and should be interpreted jointly.Source: A. Diagnostic metrics · p. 16

B. Energy-Score Optimisation

• Over unrestricted predictive distributions, the unique Energy Score minimizer for a deterministic observation is the Dirac measure at that observation.Positive dispersion can be optimal only when the admissible predictive family excludes this optimum.Source: B. Energy-Score Optimisation · p. 17
• For a Gaussian predictive family with an inaccurate mean, Energy Score minimization has a non-zero optimal scale, so deterministic targets can still yield predictive dispersion.When the predictive mean matches the observation, the optimum instead occurs at zero scale.Source: Gaussian Scale Family · p. 17
• For non-degenerate centred location-scale families, Energy Score is convex in scale; its minimizer is zero when the predictive mean matches the observation and strictly positive otherwise.When the mean is inaccurate, a global minimizer exists.Source: General Centred Location–Scale Families · p. 18
• For these families, the optimal scale is proportional to the absolute residual, with the proportionality coefficient determined by the distribution of the centred random variable.For the Gaussian case, the coefficient is 1/√log 2.Source: Other preliminary results · p. 19
• Within a constrained symmetric two-point family, the Energy Score prefers a bimodal law over the misspecified deterministic prediction, with the optimal separation equal to the absolute residual.The optimum is a* = |e|.Source: Other preliminary results · p. 19
• Constraining only the predictive mean can leave Energy Score optimization without an attained infimum, so further restrictions on the predictive family are necessary.The construction can place almost all probability mass at the observation while sending vanishing mass arbitrarily far away to preserve the mean.Source: Other preliminary results · pp. 19–20
• In OTD, the transported mask family determines admissible predictive laws, while kinetic regularisation discourages excessive latent deformation but does not directly control output-space dispersion.Within a restricted family, dispersion can reduce Energy Distance and absorb part of residual error, so predictive spread may reflect model-class constraints as well as aleatoric uncertainty.Source: Other preliminary results · p. 20

C. OTD Implementation Details

• A relaxed independent Bernoulli mask is sampled from uniform noise, with temperature controlling how closely it approaches binary masks.Small temperatures approach binary masks; at p = 1/2 and τ = 1, the relaxed mask is uniform on (0, 1).Source: C. OTD Implementation Details · p. 20
• The learned velocity field transports relaxed masks in unconstrained logit space, after which the masks are bounded for use by the predictive network.The predictive network receives masks in (0, 1)^dz.Source: C. OTD Implementation Details · pp. 20–21
• Transported masks are partitioned across hidden layers and multiplicatively modulate their corresponding representations.Source: C. OTD Implementation Details · p. 21
• The Monte Carlo Energy Score balances agreement with the observation against internal predictive dispersion, and its samples support inference summaries such as variances and intervals.Source: C. OTD Implementation Details · p. 21
• The training objective adds a kinetic-action penalty to the Energy Score, with the two terms estimated from independent Monte Carlo samples.Source: C. OTD Implementation Details · p. 22

D.1 Bimodal Dataset

• OTD learns a non-trivial correlation pattern in transported masks in the bimodal experiment.The experiment tests whether a shared transported mask law can induce locally bimodal output distributions through the input-dependent predictive network.Source: §D.1 · pp. 22–23

D.2 Reduced-Order Model misspecification Dataset

• The reduced-order experiment uses a deliberately small predictor whose mean can capture the low-frequency target component but not the higher-frequency component.This controlled misspecification setting compares predictive dispersion with the unresolved target component rather than observational noise.Source: §D.2 · p. 23

D.3 Latent collapse experiment

• The latent-collapse experiment varies hidden width around the teacher's 16-segment capacity to examine how output dispersion and transported masks change with model capacity.Below the threshold, variability may compensate for unresolved structure; with greater capacity, the predictor can reduce sensitivity to the mask.Source: §D.3 · p. 24

D.4 Ginzburg–Landau experiment

• The Ginzburg–Landau task maps a scalar parameter and spatial coordinates to equilibrium fields, using 204 training and 52 test fields.The dataset contains 256 parameter values, and each grid point is treated as a neural-field observation.Source: §D.4 · pp. 25–26

Results.

• Across Ginzburg–Landau cases, larger field errors generally accompany greater Energy Score dispersion and predictive standard deviation, but sample 00255 is under-dispersed.The error–dispersion associations are also reported at sample and pointwise levels; rank correlations capture ordering, not quantitative agreement.Source: Results · pp. 26–28
• For Ginzburg–Landau cases, predictive dispersion generally aligns with error, although the particularly difficult sample 00255 remains under-dispersed.Source: Results · pp. 27–28, Figure 12 · p. 33

D.5 UCI/OpenML Protocol

• On real-world regression datasets, OTD provides competitive point accuracy and meaningful predictive distributions, with several cases showing good calibration.Evaluation uses five-fold cross-validation and sample-based predictive metrics.Source: §D.5 · p. 35

E. The role of the kinetic regularisation

• Increasing kinetic regularisation smooths latent trajectories but can restrict reference-measure deformation and reduce the expressiveness of the learned distribution.The regularisation weight is chosen to balance trajectory regularity against transport flexibility.Source: E. The role of the kinetic… · p. 36, Figure 13 · p. 37

F. Sensitivity analysis: timing sweep

• At full grid resolution, reducing KES from 4 to 1 did not reduce runtime, while increasing it to 8 produced a comparable time.Removing kinetic regularisation yielded a 1.03× speedup; zero or one flow step yielded 1.06× and 1.05× speedups, respectively.Source: F. Sensitivity analysis: timing sweep · p. 37, Table 16 · p. 38
• Four flow steps took a time comparable to the reference, indicating a modest measured runtime contribution from flow integration in this configuration.The sweep measured 1000 epochs after five warmup epochs, averaged over three seeds.Source: F. Sensitivity analysis: timing sweep · p. 37, Figure 14 · p. 38
• The 32×32 and 16×16 grids took 148.82 and 144.81 seconds, compared with 145.76 seconds for the 64×64 reference.The timings do not establish which operation dominates computational cost, and the adequacy of flow time discretisation should be checked against a finer integration grid.Source: F. Sensitivity analysis: timing sweep · p. 38, Table 16 · p. 38

G. A generalisation of OTD

• The conditional formulation makes latent dimension dz a free choice by transporting a latent code and concatenating it to the neural-field input.This decouples the stochastic dimension from the total number of hidden features.Source: G. A generalisation of OTD · p. 39
• For mask-based OTD, dz = 1536 has the lowest RMSE, MAE, and Energy Score, while dz = 512 remains close with RMSE 0.1397±0.0025 and SES 0.0330±0.0013.The matched dz = 1536 model reports RMSE 0.1341±0.0026 and SES 0.0310±0.0003.Source: G. A generalisation of OTD · p. 39, Table 17 · p. 39
• The smallest mask-based model, dz = 32, yields the sharpest intervals but worse calibration, while the low-latent OTD variants take about 0.185–0.189 seconds per epoch.Reducing dz mainly reduces transport dimension rather than total wall-clock cost in those OTD runs.Source: G. A generalisation of OTD · pp. 39–40
• For conditional models, reducing dz improves predictive metrics relative to dz = 1536, with dz = 32 best among those models for RMSE, MAE, and SES.The matched conditional model takes approximately 0.952 seconds per epoch, compared with 0.153 seconds for dz = 32.Source: G. A generalisation of OTD · pp. 39–40, Table 17 · p. 39