PaperlayerView exact revision on arXiv

Source-linked AI summary

Factorized Fourier Neural Operators

Alasdair Tran, Alexander Mathews, Lexing Xie, Cheng Soon Ong

arXiv:2111.13802v4cs.LGcs.CE

TL;DR

Numerical PDE simulation can be prohibitively expensive, while existing neural operators face stability and geometry-related limitations. F-FNO addresses these issues with factorized Fourier layers, improved residual connections, and targeted training strategies, outperforming FNO and geo-FNO across PDEs and geometries while achieving larger time steps and substantial speedups.

  • Problem

    Existing neural operators face stability issues and deteriorating performance on complex geometries and noisy data, despite the computational expense of conventional PDE solvers.

  • Method

    F-FNO combines dimension-wise Fourier factorization, post-activation residual connections, and training strategies including Markov constraints, Gaussian noise, and cosine learning-rate decay.

  • Results

    F-FNO outperforms FNO and geo-FNO across PDE systems and geometries, reducing Navier–Stokes error by 83% versus FNO and achieving an order-of-magnitude speedup over a pseudo-spectral solver.

  • Takeaways & Limitations

    F-FNO provides a flexible learned PDE solver that scales to deeper networks and higher-dimensional problems while supporting larger time steps than numerical solvers.

Abstract

from arXiv · show

We propose the Factorized Fourier Neural Operator (F-FNO), a learning-based approach for simulating partial differential equations (PDEs). Starting from a recently proposed Fourier representation of flow fields, the F-FNO bridges the performance gap between pure machine learning approaches to that of the best numerical or hybrid solvers. This is achieved with new representations - separable spectral layers and improved residual connections - and a combination of training strategies such as the Markov assumption, Gaussian noise, and cosine learning rate decay. On several challenging benchmark PDEs on regular grids, structured meshes, and point clouds, the F-FNO can scale to deeper networks and outperform both the FNO and the geo-FNO, reducing the error by 83% on the Navier-Stokes problem, 31% on the elasticity problem, 57% on the airfoil flow problem, and 60% on the plastic forging problem. Compared to the state-of-the-art pseudo-spectral method, the F-FNO can take a step size that is an order of magnitude larger in time and achieve an order of magnitude speedup to produce the same solution quality.

1 INTRODUCTION

PDE solvers are computationally expensive, motivating learned operators; F-FNO addresses stability, scalability, and efficiency limitations through factorized representations, improved residual connections, and training strategies.

  • Motivation: PDE simulations often require computationally expensive numerical solvers because most real-world problems lack closed-form solutions.Such simulations can consume millions of core hours and terabytes of storage.
  • Existing neural operators: FNOs learn mappings between infinite-dimensional function spaces and offer larger time steps, super-resolution, and a shared architecture across PDEs.geo-FNO extends this approach to structured meshes and point clouds.
  • Existing neural operators: First-generation neural operators suffer stability issues, with performance deteriorating on complex geometries and noisy data.The authors also observe worsening performance for FNO and geo-FNO in their experiments.
  • F-FNO approach: F-FNO factorizes Fourier features by dimension, reducing model complexity by an order of magnitude and enabling higher-dimensional problems such as 3D plastic forging.The architecture also places residual connections after activation to support deeper networks.
  • F-FNO approach: Training combines teacher forcing, Markov constraints, Gaussian input noise, and cosine learning-rate scheduling to improve operator performance.The paper identifies these techniques as important components of its training setup.
  • Evaluation: F-FNO outperforms prior operators across multiple PDE systems and geometries, including an 83% error reduction versus FNO on Navier–Stokes.The paper evaluates regular-grid, structured-mesh, and point-cloud settings.

2 RELATED WORK

PDE solution methods range from classical discretization and hybrid correction schemes to pure learned operators, including Fourier-based approaches for efficient function-space mappings.

  • Classical methods: Classical PDE solvers discretize space using finite element, finite difference, finite volume, or pseudo-spectral methods, with finer resolution increasing computational cost.Simplified physical models have traditionally been used to reduce this cost.
  • Hybrid methods: Hybrid methods accelerate numerical solvers through lower-resolution grids or learned replacements for computationally expensive solver components.Examples include data-driven discretization and neural-network interpolation.
  • Hybrid methods: Learned correction adds a neural residual or super-resolution term to the output of a coarse numerical step.This approach is less specialized than many hybrid methods but less general than pure machine learning methods.
  • Pure machine learning: Pure machine-learning approaches learn the field directly with neural operators based on graph networks, low-rank decompositions, or Fourier transforms.They can also incorporate physical constraints through loss functions based on conservation laws.
  • Fourier representations: Fourier-based operators efficiently model long-range dependencies using frequency-space convolutions and FFT complexity O(n log n).FNO is presented as a Fourier-based approach for learning mappings between function spaces.

3 THE FACTORIZED FOURIER NEURAL OPERATOR

F-FNO learns PDE solution operators by combining Fourier-based operator layers with coordinate deformation for irregular geometries. Its factorized spectral representation, post-activation residual connections, and flexible inputs reduce complexity and support deeper, higher-dimensional models.

  • Neural-operator formulation: F-FNO maps an input function a to an output function u through lifting, nonlinear operator layers, projection, and optional coordinate deformation.The coordinate map deforms irregular physical geometries into a regular computational space for geo-FNO-style processing.
  • Factorized representation: Fourier factorization processes each problem dimension independently, reducing parameters from O(LH^2M^D) to O(LH^2MD).Sharing weight matrices across layers can reduce the parameter count further to O(H^2MD).
  • Improved operator layer: The improved operator layer places residual connections after the activation and adds a two-layer feedforward component.These changes are designed to preserve more of the layer input and support deeper networks.
  • Scaling and flexibility: F-FNO supports higher-dimensional problems, including 3D plastic forging, while continuing to improve performance in deep networks.The paper attributes this capability to the combination of factorized transforms and residual connections.
  • Scaling and flexibility: Relevant field-evolution information, such as viscosity or external forcing, can be included in the input representation across different PDEs.This input flexibility is presented as enabling straightforward generalization to different PDEs.
  • Training strategy: The training strategy enforces a first-order Markov property by feeding only the current step rather than the previous 10 time steps.This mirrors the information used by a numerical solver.

4 DATASETS AND EVALUATION SETTINGS

The evaluation covers regular-grid Navier–Stokes tasks and irregular-geometry elasticity, airfoil, and plasticity problems. Models are trained with task-specific schedules and evaluated using normalized mean squared error, with vorticity correlation added for TorusKochkov.

  • Regular-grid PDEs: Four Torus datasets evaluate turbulent flows on a 2D periodic grid representing the surface of a 3D torus.TorusLi benchmarks against the original FNO, while TorusVis and TorusVisForce vary viscosity and forcing to test generalization.
  • Regular-grid PDEs: TorusKochkov uses 2048x2048 Carpenter-Kennedy simulations downsampled to smaller grids for speed–accuracy analysis.Only 32 training trajectories are generated to test learning in a low-data regime.
  • Irregular geometries: Elasticity, Airfoil, and Plasticity evaluate point-cloud, mesh, and structured-mesh tasks on irregular geometries.Their targets are stress, Mach number, and a three-dimensional displacement-related output over space and time, respectively.
  • Training and metrics: Training uses 100,000 steps on regular-grid datasets and 200 epochs on irregular-geometry datasets, with warmup followed by cosine learning-rate decay.The learning rate is warmed to 2.5 × 10^-3 during the first 500 steps.
  • Training and metrics: Normalized mean squared error is the primary comparison metric, while TorusKochkov additionally uses vorticity correlation.The correlation evaluation measures the time until correlation drops below 95%.

5 RESULTS FOR NAIVER-STOKES ON A TORUS

On Torus benchmarks, F-FNO improves accuracy and deep-network scaling over FNO while reducing parameters, enabling larger time steps and faster simulations than numerical solvers. Its performance also depends on training strategy and input representation.

  • Comparison against FNO: F-FNO is substantially more accurate than FNO across network depths, uses fewer parameters, and remains up to two orders of magnitude faster at inference than Crank-Nicolson.Training time is similar, while inference is generally longer than FNO inference.
  • Training strategies: Teacher forcing improves FNO accuracy below 24 layers, while the first-order Markov assumption improves it further; additional history does not help.The Markov variant uses only one step of history.
  • Training strategies: Residual connections after the nonlinearity allow convergence at 24 layers, while normalization, Gaussian noise, and cosine decay further improve performance.Without Gaussian noise, early validation loss can explode.
  • Fourier factorization: 35%: Fourier factorization reduces error from 3.73% to 2.41% at 24 layers while reducing parameter count by an order of magnitude.Weight sharing reduces the parameter count further to 1M with little additional performance change at deep layers.
  • Trade-off between speed and accuracy: F-FNO reaches similar performance using 64x64 inputs compared with DNS on 128x128 grids, while achieving an order-of-magnitude speedup.At equal spatial resolution, F-FNO can also produce a vorticity field closer to ground truth than DNS.
  • Optimal step size: F-FNO can use a time step at least an order of magnitude larger than a stable numerical-solver step; on TorusKochkov, its sweet spot is around 0.2.Unlike numerical solvers, it does not need smaller steps at higher spatial resolution.
  • Flexible input representations: With vorticity, viscosity, and forcing as inputs, F-FNO achieves 2% error; removing viscosity doubles error, while removing forcing increases it by an order of magnitude.Redundant features do not significantly hurt performance, and different datasets benefit from different input feature sets.
  • Flexible input representations: Coordinate encoding improves performance even when grid indices already contain positional information, suggesting different model components use the two encodings.The Fourier layer uses absolute grid positions, while the pointwise physical-space layer relies on raw coordinates.

6 RESULTS FOR PDES ON POINT CLOUDS AND MESHES

F-FNO scales better with network depth and hidden size than geo-FNO across irregular-geometry PDEs, delivering lower prediction errors on elasticity, airfoil flow, and plastic forging.

  • Results: geo-FNO scaling deteriorates with depth, including non-convergence beyond 12 layers for Airfoil and Plasticity.It also appears stuck in a local minimum beyond 8 layers on Elasticity and can worsen as hidden size increases.
  • Results: 31% lower prediction error on Elasticity point clouds, from 2.51% to 1.74%.F-FNO continues improving with deeper networks and larger hidden size.
  • Results: 57% lower prediction error on the 2D transonic airfoil-flow problem, from 1.35% to 0.58%.The improvement accompanies continued gains with deeper networks and larger hidden size.
  • Results: 60% lower prediction error on plastic forging, from 0.45% to 0.18%.A 24-layer F-FNO with 11M parameters outperforms a 12-layer geo-FNO with 57M parameters.

7 CONCLUSION

The paper concludes that Fourier factorization, improved residual connections, and better training enable F-FNO to outperform prior neural operators across PDE geometries and domains.

  • Conclusion: F-FNO outperforms the state of the art on PDEs across a variety of geometries and domains.The approach combines factorized Fourier transforms, improved residual connections, and an improved training setup.
  • Conclusion: Fourier transforms provide an efficient way to learn neural operators with long-range spatial dependencies.The paper identifies this capability as a central strength of Fourier-based operator learning.

A APPENDIX

The appendix documents datasets, resource scaling, super-resolution, energy-spectrum behavior, transform ablations, and detailed benchmark results across regular and irregular geometries.

  • Datasets: TorusVis and TorusVisForce broaden fluid-dynamics data with more varied viscosities and forcing functions than earlier datasets.The appendix notes that the earlier TorusLi dataset lacked a validation set.
  • Resource scaling: F-FNO variants show better predictions as inference time and model resources increase.Figure A.1 tracks parameter count and inference time, with error bars showing min-max values over three trials when applicable.
  • Zero-shot super-resolution: F-FNO supports zero-shot super-resolution from training on 32x32 and 64x64 grids to inference on 128x128 and 256x256 grids.Performance degrades on grid sizes unseen during training.
  • Energy spectrum: F-FNO trained on 64x64 grids captures the long-term energy-spectrum trend better than DNS on a grid four times larger.Its spectrum still has substantially lower high-wavenumber energy than 2048x2048 DNS because only the top 16 modes are retained.
  • Transform ablation: F-CNO outperforms F-FNO on Airfoil at deeper layers, while the two methods perform mostly similarly on Plasticity at equal depth.The cosine transform is real-valued and therefore halves the number of parameters.
  • Detailed results: Detailed appendix tables report three-trial mean N-MSE values with min-max ranges for TorusLi, Airfoil, Elasticity, and Plasticity.The tables provide expanded benchmark results and input-representation configurations.
  • Transform ablation: Increasing depth raises parameter count in the Airfoil and Plasticity transform comparison.Figure A.3 plots test loss against parameter count and reports min-max error bars over three trials.
  • Visual comparisons: The appendix compares visual vorticity correlations between F-FNO predictions and DNS across resolutions on TorusKochkov.The figure includes zero-shot super-resolution and energy-spectrum analyses.
Loading 2111.13802v4…