Generative Neural Operators through Diffusion Last Layer

Neural operators provide a powerful framework for learning discretization invariant mappings between function spaces, but standard deterministic models do not capture predictive uncertainty. We introduce diffusion last layer (DLL), a modular probabilistic output head for neural operator backbones. DLL represents target fields through an input dependent low rank expansion inspired by the Karhunen--Lo`eve expansion and learns a conditional diffusion model over the corresponding coefficient space. This design enables efficient distributional modeling while preserving the structural advantages of operator learning. On stochastic PDE benchmarks with random forcing, DLL achieves strong distributional fidelity and performs competitively with pixel space and conventional latent diffusion baselines. In deterministic long horizon rollout tasks, DLL improves rollout stability over the underlying backbone and provides useful estimates of predictive uncertainty under compounding autoregressive errors. These results suggest that diffusion modeling in learned coefficient spaces offers a practical route to uncertainty aware neural operators. Code is available at https://github.com/sungwpark/dll-no

Scientific simulations often require predicting complex physical quantities, such as pressure, temperature, or fluid motion. Modern machine learning models can learn these predictions from data, but many standard models produce only a single answer and do not indicate how uncertain that answer is. This is limiting when uncertainty arises from random effects, limited data, or chaotic dynamics. We propose Diffusion Last Layer, a lightweight component that can be added to an existing scientific prediction model to generate multiple plausible outcomes. Instead of generating the full physical field directly, it represents the field using a compact learned description and generates predictions in that smaller space. This makes uncertainty estimation more efficient while preserving physical structure. Experiments on stochastic and chaotic systems show that the proposed method improves distributional prediction and makes long term forecasts more stable, suggesting a practical path toward more reliable machine learning tools for scientific simulation.