Research map / Probabilistic and scientific modeling / Differential equations and PDEs
Light-Weight Diffusion Multiplier and Uncertainty Quantification for Fourier Neural Operators
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Replaces the dense spectral multiplier in Fourier Neural Operators with a heat-kernel diffusion multiplier having one learnable time per channel, and places priors over those times for variational Bayesian uncertainty quantification.
Problem
Fourier Neural Operators are overparameterized (parameters grow exponentially with domain dimension) and provide no native uncertainty quantification, with existing post hoc UQ methods ignoring neural operator structure.
Approach
Restrict the integral transform kernel to the heat kernel so each channel has a single learnable diffusion time, add spectral gradient features for anisotropy, and treat diffusion times as latent variables with a Gaussian prior and full-rank variational posterior so output samples are spatially correlated.
Key results
- On Darcy flow, DINOZAUR with gradient features reaches RL2 0.75 (x10^-2) with 28.0k parameters, versus 0.90 for FNO with 361.4k parameters.
- On ShapeNet car, the Bayesian DINOZAUR_B achieves the lowest NLL (2.767) and miscalibration area (0.168) among compared UQ methods, versus 3.090 NLL for GINO with dropout.
- On Navier-Stokes, DINOZAUR reaches RL2 15.68 (x10^-2) with about 109.7k parameters, versus 15.90 for FNO with about 12.65M parameters.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace dense FNO spectral multipliers with diffusion multipliers and include gradient features | FNO · Darcy flow | RL2 (×10^-2): 0.90 ± 0.03 -> 0.75 ± 0.10 | Helped |
| Replace dense FNO spectral multipliers with diffusion multipliers and include gradient features | FNO · Navier-Stokes | RL2 (×10^-2): 15.90 ± 0.20 -> 15.68 ± 0.21 | Helped |
| Replace tensorized GINO multipliers with diffusion multipliers and include gradient features | TGINO · ShapeNet car | RL2 (×10^-2): 7.27 ± 0.16 -> 7.04 ± 0.06 | Helped |
| Replace tensorized GINO multipliers with diffusion multipliers and include gradient features | TGINO · Ahmed bodies | RL2 (×10^-2): 8.19 ± 0.49 -> 7.99 ± 0.07 | Helped |
| Add Bayesian diffusion-time uncertainty instead of MC Dropout | DINOZAUR · ShapeNet car | NLL: 2.996 ± 0.010 -> 2.767 ± 0.011 | Helped |
| Add Bayesian diffusion-time uncertainty instead of Laplace approximation | DINOZAUR · ShapeNet car | MA: 0.266 ± 0.018 -> 0.168 ± 0.001 | Helped |
| Add Bayesian diffusion-time uncertainty instead of MC Dropout | DINOZAUR · ShapeNet car | IS: 13.704 ± 0.189 -> 11.667 ± 0.136 | Helped |
| Add gradient features (ablation) | DINOZAUR diffusion blocks · Darcy flow | RL2 (×10^-2): 1.08 ± 0.06 -> 0.75 ± 0.10 | Helped |
| Add gradient features (ablation) | DINOZAUR diffusion blocks · Navier-Stokes | RL2 (×10^-2): 20.10 ± 0.31 -> 15.68 ± 0.21 | Helped |
| Add gradient features (ablation) | DINOZAUR diffusion blocks · ShapeNet car | RL2 (×10^-2): 8.34 ± 0.80 -> 7.04 ± 0.06 | Helped |
| Add gradient features (ablation) | DINOZAUR diffusion blocks · Ahmed bodies | RL2 (×10^-2): 9.59 ± 0.62 -> 7.99 ± 0.07 | Helped |
| Add explicit gradient features (ablation) | FNO and tensorized FNO baselines · PDE benchmarks | RL2 | Mixed |
Surprising findings
- Without gradient features, the pure diffusion multiplier (DINOZAUR_nograd) performs poorly, attributed to the radial symmetry of diffusion.
- Reducing parameter count in GINO (via TGINO) improved performance on non-uniform datasets, indicating overparameterization.
Limitations
- Pure diffusion multiplier is radially symmetric and struggles to train without gradient features, which add compute time.
- Bayesian formulation is sensitive to prior hyperparameter choices.
- Parameter count still scales with width and depth, and results on Navier-Stokes are not better than dense baselines without gradient features.
Experimental setup
Datasets: Darcy flow Navier-Stokes 2D ShapeNet car Ahmed bodies
Benchmarks: Darcy flow Navier-Stokes 2D ShapeNet car Ahmed bodies
Metrics: RL2 NLL miscalibration area interval score peak GPU memory training time parameter count
Baselines: FNO TFNO GINO TGINO MC Dropout Laplace approximation
Builds on: Fourier Neural Operator DiffusionNet GINO Tensorized FNO Laplace approximation
Closest papers by meaning
- Using Uncertainty Quantification to Characterize and Improve Out-of-Domain Learning for PDEs (ICML 2024 · similarity 0.89)
- Maximum Update Parametrization and Zero-Shot Hyperparameter Transfer for Fourier Neural Operators (ICML 2025 · similarity 0.88)
- Linearization Turns Neural Operators into Function-Valued Gaussian Processes (ICML 2025 · similarity 0.87)
- Amortized Fourier Neural Operators (NeurIPS 2024 · similarity 0.87)
- Guaranteed Approximation Bounds for Mixed-Precision Neural Operators (ICLR 2024 · similarity 0.87)
- Beyond Regular Grids: Fourier-Based Neural Operators on Arbitrary Domains (ICML 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.