Generative Neural Operators through Diffusion Last Layer
This work addresses uncertainty quantification for neural operators in scientific computing, offering a practical add-on for improved reliability, though it is incremental as it builds on existing neural operator backbones.
The paper tackles the challenge of uncertainty quantification in neural operators for stochastic systems by introducing the diffusion last layer (DLL), a lightweight probabilistic head that improves generalization and uncertainty-aware prediction in stochastic PDE benchmarks, while also enhancing rollout stability and providing epistemic uncertainty estimates in deterministic settings.
Neural operators have emerged as a powerful paradigm for learning discretization-invariant function-to-function mappings in scientific computing. However, many practical systems are inherently stochastic, making principled uncertainty quantification essential for reliable deployment. To address this, we introduce a simple add-on, the diffusion last layer (DLL), a lightweight probabilistic head that can be attached to arbitrary neural operator backbones to model predictive uncertainty. Motivated by the relative smoothness and low-dimensional structure often exhibited by PDE solution distributions, DLL parameterizes the conditional output distribution directly in function space through a low-rank Karhunen-Loève expansion, enabling efficient and expressive uncertainty modeling. Across stochastic PDE operator learning benchmarks, DLL improves generalization and uncertainty-aware prediction. Moreover, even in deterministic long-horizon rollout settings, DLL enhances rollout stability and provides meaningful estimates of epistemic uncertainty for backbone neural operators.