Towards Intrinsically Calibrated Uncertainty Quantification in Industrial Data-Driven Models via Diffusion Sampler

Abstract: In modern process industries, data-driven models are important tools for real-time monitoring when key performance indicators are difficult to measure directly. While accurate predictions are essential, reliable uncertainty quantification (UQ) is equally critical for safety, reliability, and decision-making, but remains a major challenge in current data-driven approaches. In this work, we introduce a diffusion-based posterior sampling framework that inherently produces well-calibrated predictive uncertainty via faithful posterior sampling, eliminating the need for post-hoc calibration. In extensive evaluations on synthetic distributions, the Raman-based phenylacetic acid soft sensor benchmark, and a real ammonia synthesis case study, our method achieves practical improvements over existing UQ techniques in both uncertainty calibration and predictive accuracy. These results highlight diffusion samplers as a principled and scalable paradigm for advancing uncertainty-aware modeling in industrial applications.

• The paper introduces a diffusion sampling framework using the Schrödinger bridge paradigm and stochastic control to achieve intrinsic uncertainty calibration.
• It leverages neural diffusion samplers to accurately represent complex, multimodal, and non-Gaussian posterior distributions, outperforming traditional UQ methods.
• Empirical studies in penicillin fermentation and ammonia synthesis demonstrate enhanced predictive performance and robust uncertainty calibration.

DiffUQ: Diffusion-Based Uncertainty Quantification for Industrial Data-Driven Modeling

Motivation and Theoretical Framework

Industrial process models increasingly rely on data-driven methodologies to infer key quality indicators that are difficult to measure in real time. However, a fundamental challenge persists: point predictions without reliable uncertainty quantification (UQ) cannot support risk-sensitive operations or guarantee the robustness required in safety-critical environments. This work addresses intrinsically calibrated UQ, emphasizing posterior sampling without post-hoc calibration, which is largely absent in traditional industrial model deployments.

The core innovation is a posterior sampling framework grounded in the Schrödinger bridge (SB) paradigm, connecting Bayesian inference to finite-horizon stochastic optimal control (SOC). This enables principled sampling from high-dimensional, possibly multimodal and non-Gaussian posteriors ubiquitous in industrial applications. Posterior inference is reformulated as a stochastic transport problem and solved through neural diffusion samplers, with the SOC cost function relaxed for tractable end-to-end optimization.

Figure 1: Overview of the DiffUQ framework integrating probabilistic regression, posterior sampling via diffusion SDEs, and uncertainty-aware prediction aggregation.

Diffusion Sampler Construction and Algorithmic Details

The methodology employs a probabilistic regression model, $p(y|x,\theta)$, that incorporates both aleatoric and epistemic uncertainty. For flexible noise modeling, two neural networks provide the mean and input-dependent precision. The parameter posterior $p(\theta|\mathcal{D})$ is sampled by simulating controlled SDEs, where the optimal drift $u_\phi(t,\theta)$ is learned via minimization of a composite objective balancing control effort (running cost) and posterior-fit (terminal penalty). Training leverages the Euler–Maruyama scheme for SDE discretization; inference uses parallel trajectory simulation and aggregation of model outputs.

Sensitivity analyses demonstrate robustness to step size ($\Delta t$) and diffusion coefficient ($\gamma$) choices, reinforcing suitability for deployment in dynamic industrial contexts.

Figure 2: Effect of different $\Delta t$ during the training phase shows negligible impact on calibration and accuracy metrics.

Figure 3: Variation of $\Delta t$ and $\gamma$ demonstrate invariance in loss dynamics and predictive performance.

Comparison with Conventional UQ Methods

The DiffUQ framework is evaluated against MC dropout, deep ensembles (DE), mean-field variational inference (MFVI), stochastic gradient Langevin dynamics (SGLD), Stein variational gradient descent (SVGD), and MAP. Unlike MC dropout and MFVI, which rely on restrictive unimodal or factorized assumptions leading to systematic under-coverage, DiffUQ exploits diffusion samplers for high-fidelity coverage of complex posterior landscapes. The method avoids finite-particle degeneracies common in SVGD and slow mixing in SGLD.

In synthetic benchmarks, DiffUQ consistently achieves superior mode coverage, as illustrated for non-Gaussian and multimodal distributions.

Figure 4: MFVI, SGLD, and SVGD inadequately sample complicated distributions, whereas diffusion samplers accurately represent multimodal and ill-posed posteriors.

Industrial Case Study: Penicillin Fermentation

Penicillin fermentation data (IndPenSim) is used to assess Raman-based PAA soft sensor modeling. Despite the linearity of the task, DiffUQ yields optimal uncertainty calibration (lowest NLL, ECE, and MCE) and superior predictive accuracy on all metrics compared to alternatives. The performance of MFVI is competitive due to linear model simplicity; however, DiffUQ demonstrates greater resilience to finite-sample effects and mixing inefficiency.

Figure 5: IndPenSim provides the industrial-scale simulation platform for benchmarking uncertainty-aware modeling.

Industrial Case Study: High-Low Transformer in Ammonia Synthesis

For real-world nonlinear modeling, predicting residual CO concentrations in the ammonia plant HLT unit is investigated. DiffUQ employs deep neural architectures and delivers systematically improved calibration and accuracy versus baselines. SVGD occasionally matches calibration but at the cost of worse predictive precision. DiffUQ maintains robustness across sample sizes, highlighting sample efficiency deriving from exploration amortized pre-inference and learned transport dynamics.

Figure 6: Diagram of the ammonia synthesis HLT unit contextualizing the challenging CO concentration prediction task.

Practical and Theoretical Implications

DiffUQ’s posterior calibration obviates the need for post-hoc validation, critical in industrial settings with scarce or expensive ground-truth data. The strong empirical results on both linear and nonlinear modeling tasks confirm that diffusion samplers, guided by stochastic control, are a scalable foundation for uncertainty-aware industrial modeling. Algorithmic regularization induced by diffusion processes ensures well-behaved training landscapes and monotonic convergence.

Model capacity studies indicate non-monotonic dependence on drift network size: under a fixed numerical discretization budget, excessive capacity may amplify discretization error and degrade calibration, evidencing a practical tradeoff between expressivity and simulation accuracy.

Figure 7: NLL vs. drift network width and depth validates existence of optimal network capacity for calibration under fixed discretization.

Conclusion

This study advances industrial uncertainty quantification by introducing DiffUQ, a diffusion-based posterior sampling framework for calibrated probabilistic modeling. The method surpasses conventional UQ approaches in calibration and predictive accuracy, is resilient to hyperparameter variations, and exhibits stable training dynamics. These features make DiffUQ particularly suitable for industrial deployment where operational reliability and sample efficiency are paramount. Future research should explore more computationally efficient diffusion samplers and alternative control-based sampling formulations to further optimize uncertainty quantification in large-scale industrial settings.