Generative Neural Operators through Diffusion Last Layer

TL;DR

This paper introduces the diffusion last layer (DLL), a lightweight probabilistic extension for neural operators that models uncertainty in function-to-function mappings, improving generalization and stability in scientific computing tasks.

Contribution

The paper proposes DLL, a novel probabilistic head for neural operators that efficiently models uncertainty in stochastic and deterministic PDE solutions.

Findings

DLL improves uncertainty quantification in stochastic PDEs.

DLL enhances stability in deterministic long-horizon rollouts.

DLL outperforms existing methods in generalization on benchmarks.

Abstract

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\`eve 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.