PaperlayerView exact revision on arXiv

Source-linked AI summary

Towards Multi-spatiotemporal-scale Generalized PDE Modeling

Jayesh K. Gupta, Johannes Brandstetter

arXiv:2209.15616v2cs.LGcs.CV

TL;DR

Deep neural PDE surrogates need to model multi-scale spatio-temporal phenomena while generalizing across equation parameters and time-scales. The paper compares FNO, ResNet, and U-Net approaches for fluid mechanics, transfers modern U-Net improvements, and evaluates FNO integrations. U-Net-based methods perform best on the presented tasks, while single surrogate models show promising parameter and time-scale generalization.

  • Problem

    Neural PDE surrogates must handle complex multi-scale spatio-temporal phenomena, but generalization across different PDE parameters and time-scales remains challenging.

  • Method

    The paper comprehensively compares FNO, ResNet, and U-Net-based fluid-mechanics surrogates, transfers computer-vision U-Net improvements, and studies FNO layers within U-Nets.

  • Results

    U-Net-based methods perform best on the presented fluid-mechanics tasks, with the best architectures generalizing across different PDE parameters and time-scales within a single surrogate model.

  • Takeaways & Limitations

    The comparison provides strong baselines for developing neural PDE surrogates and indicates that modernized U-Nets can support parameter and time-scale generalization.

  • Takeaways & Limitations

    The work does not address long-rollout stability, invariant preservation, sampling regularities, domain topologies and geometries, boundary conditions, or RANS modeling.

Abstract

from arXiv · show

Partial differential equations (PDEs) are central to describing complex physical system simulations. Their expensive solution techniques have led to an increased interest in deep neural network based surrogates. However, the practical utility of training such surrogates is contingent on their ability to model complex multi-scale spatio-temporal phenomena. Various neural network architectures have been proposed to target such phenomena, most notably Fourier Neural Operators (FNOs), which give a natural handle over local & global spatial information via parameterization of different Fourier modes, and U-Nets which treat local and global information via downsampling and upsampling paths. However, generalizing across different equation parameters or time-scales still remains a challenge. In this work, we make a comprehensive comparison between various FNO, ResNet, and U-Net like approaches to fluid mechanics problems in both vorticity-stream and velocity function form. For U-Nets, we transfer recent architectural improvements from computer vision, most notably from object segmentation and generative modeling. We further analyze the design considerations for using FNO layers to improve performance of U-Net architectures without major degradation of computational cost. Finally, we show promising results on generalization to different PDE parameters and time-scales with a single surrogate model. Source code for our PyTorch benchmark framework is available at https://github.com/microsoft/pdearena.

1. Introduction

Neural PDE surrogates address the high cost of solving differential-equation models, but must handle multi-scale spatio-temporal phenomena and generalize across PDE parameters and time-scales. The paper compares FNO, ResNet, and U-Net-like approaches for these challenges in fluid mechanics.

  • Expensive PDE solution techniques have motivated deep neural network surrogates, particularly for fluid-dynamics applications.
  • Generalizing across PDE parameters and time-scales is notoriously difficult because small parameter changes can shift flows between laminar and turbulent regimes.
  • Physical phenomena span different spatial and temporal scales, from local blizzards to global heat waves governed by the same underlying principles.
  • FNOs use Fourier modes to process global and local spatial information, while U-Nets use downsampling, upsampling, and skip connections.
  • The paper presents a side-by-side comparison of FNO, ResNet, and U-Net-like architectures for complex multi-scale spatio-temporal phenomena and studies generalization across PDE parameters and time-scales.

2. Preliminaries

The preliminaries explain how Fourier transforms characterize convolutional filters and how FNOs use mode-wise spectral processing to represent spatial scales. They also introduce the incompressible Navier–Stokes formulations studied in the paper.

  • Spatial scales: Low Fourier modes encode global information, whereas high modes encode local information; FNO layers process both through mode-wise multiplication.This spectral parameterization provides a direct mechanism for handling multiple spatial scales.
  • Fourier Neural Operators: FNO layers interleave discrete Fourier transforms, mode-wise complex weight multiplication, and inverse transforms, with cutoff frequencies controlling retained modes.A 1 × 1 convolution provides a residual path, while Fourier modes above the cutoff are set to zero.
  • Convolutional layers: CNN layers apply learned filters to feature maps, with filter size determining the spatial region from which each output gathers information.Small filters are common, while dilation and flexible filter sizes offer ways to adapt receptive fields across layers.
  • U-Net filter analysis: Figure 3 averages absolute Fourier-mode values across filters in the first convolution layer of U-Net downsampling blocks to compare their spectral properties.The figure examines how filter spectra change across successive downsampling blocks.
  • Fourier transforms and convolution: Fourier analysis diagonalizes circular convolution through the discrete Fourier transform, allowing filters to be represented by their Fourier modes.The circulant convolution matrix shares Fourier eigenvectors, while the diagonal matrix contains the filter’s Fourier modes.
  • Fluid mechanics equations: The paper studies incompressible two-dimensional Navier–Stokes dynamics in velocity and vorticity-stream formulations, including convection, viscosity, pressure, forcing, and incompressibility.The vorticity formulation separates the dynamics into a parabolic vorticity-transport equation and an elliptic Poisson equation for the streamfunction.

3. PDE Surrogates

This section defines the benchmarked PDE-surrogate architectures and motivates operator learning for generalization across equation parameters and time discretizations. It compares ResNets, FNOs, U-Nets, Fourier U-Nets, and UNO-style designs.

  • ResNets: ResNet baselines use eight residual blocks with 3 × 3 convolutions, shortcut connections, group normalization, and GeLU activations without spatial down-projection.Without downsampling or pooling, these ResNets lack a built-in mechanism for jointly handling local and global information.
  • FNOs: FNO baselines vary layer count, channel count, and the number of non-zero Fourier modes, combining a 1 × 1 convolution path with a Fourier path.The architectures use two embedding layers, two output layers, GeLU activations, and no normalization.
  • U-Nets: U-Nets combine spatial downsampling and upsampling blocks with skip connections between corresponding stages, and the benchmark includes original and modernized variants.U-Netbase replaces batch normalization with group normalization and bias weights, while modern designs incorporate Wide ResNet-style blocks and attention mechanisms.
  • Fourier U-Nets: Fourier U-Nets replace lower downsampling and upsampling blocks with Fourier blocks containing two FNO layers and residual connections.The study tests substitutions at different depths and also implements UNO, which uses dimension-wise linear downsampling and mode scheduling.
  • Operator learning: Neural PDE surrogates are useful when their costly training can be amortized through generalization across PDE parameters and time discretizations.The paper frames this capability within operator learning, which relates solution functions across domains through learned operators.
  • Parameter conditioning: The evaluation tests FNO- and U-Net-like architectures on mappings that vary force terms across solution spaces and use different time windows.These experiments target parameter conditioning and temporal-scale generalization with surrogate models.

4. Experiments

The experiments benchmark PDE surrogate architectures across diverse fluid-mechanics formulations, parameter settings, and time windows. U-Net-based methods generally perform strongly, while parameter conditioning favors U-Nets over FNO variants.

  • Benchmark design: The benchmark evaluates Fourier, ResNet, and U-Net-like architectures across velocity and vorticity formulations, parameter conditioning, and multiple time horizons.Tasks use PDE-solver-backed datasets and probe Fourier versus U-Net approaches, formulation differences, and generalization across parameters and time-scales.
  • Shallow water equations: The shallow-water experiments compare models on 192 × 96 grids with 48-hour temporal resolution, predicting pressure and velocity or vorticity fields.The velocity formulation predicts a scalar pressure field and vector wind velocity; the vorticity formulation predicts scalar pressure and scalar wind vorticity.
  • Shallow water equations: U-Net architectures achieve the lowest losses in the shallow-water vorticity-stream formulation, while architectures with dedicated local and global information flow perform well in the velocity formulation.The compared local-global architectures include Dilated ResNet, FNO, and U-Net methods.
  • Parameter conditioning: Parameter-conditioning experiments train models across unseen buoyancy-force values and time windows ranging from 0.375 s to 20 s.Conditioning information uses sinusoidal Fourier embeddings for the continuous time-window and force parameters, injected into residual blocks or feature maps.
  • Parameter conditioning: U-Net-based methods perform best across unseen buoyancy-force evaluations, whereas replacing lower U-Net blocks with FNO blocks does not improve generalization.The evaluation averages results over 208 unseen force values between 0.2 and 0.5.
  • Parameter conditioning: Conditioning is more difficult for FNO layers, although alternative parameter-embedding methods may improve their conditioning behavior.The authors relate this observation to prior findings that FNO-like architectures can be sensitive to noise.

5. Conclusion

The paper establishes strong baselines for neural PDE surrogates through comprehensive comparisons of ResNet, FNO, and U-Net architectures, while demonstrating promising parameter and time-scale generalization. It also identifies important limitations and directions for future work.

  • Conclusion: The study provides a comprehensive comparison of ResNet, FNO, and U-Net approaches for fluid mechanics PDE surrogate modeling.It aims to establish strong baselines for future neural PDE surrogate research.
  • Conclusion: The original U-Net already functions as a powerful neural PDE surrogate and outperforms FNOs on the presented tasks.Recent architectural improvements to U-Nets produced further significant performance improvements.
  • Conclusion: Best-performing U-Net architectures generalize across different PDE parameters and time-scales within a single surrogate model.FNO layers during early U-Net downsampling improved performance in some circumstances but negatively affected parameter and time-scale generalization.
  • Conclusion: The benchmark codebase is intended as a starting point for further investigations of neural PDE surrogates.The framework supports side-by-side comparisons of approaches examined in the study.
  • Limitations & Future Work: The work does not address long-rollout stability, invariant preservation, sampling regularity, domain topology and geometry, or boundary-condition generalization.It also models Navier–Stokes directly rather than in Reynolds-averaged form and leaves Vision Transformer comparisons for future work.

B. Experiments

This appendix supports Section 4 of the main paper.

  • B. Experiments: The appendix provides supporting material for Section 4 of the main paper.No additional experimental content is specified in the supplied passage.
  • B. Experiments: Its stated role is to supplement the main paper's experimental section.The passage identifies the relationship to Section 4 without describing methods or results.
  • B. Experiments: The supplied text frames the appendix as supplemental rather than as an independent analysis.The passage contains no further details about the experiments themselves.

B.1. Experimental details

The experiments use SMSE and RMSE-based objectives and evaluate one-step and five-step rollout losses across tuned neural PDE surrogate architectures. Training uses AdamW with scheduled learning rates, while runtime and resource costs are also reported.

  • Loss functions and metrics: The study evaluates one-step loss with N_t=1 and rollout loss with N_t=5 across scalar fields and vector-field components.Both metrics use the summed mean squared error framework.
  • Loss functions and metrics: Models can alternatively be trained with the summed relative MSE loss introduced by Li et al. (2020a).The appendix presents RMSE as an alternative to the summed MSE objective.
  • Training and model selection: Training uses AdamW for 50 epochs with learning rates of [10^-4, 2 · 10^-4], weight decay 10^-5, cosine annealing, and linear warmup.Reported models minimize SMSE during optimization.
  • Training and model selection: ResNet and FNO configurations are selected by varying depth, channels, normalization, Fourier modes, and activation functions.Reported ResNets use group normalization with one group, while FNOs use GeLU and no normalization.
  • Computational resources: Experiments use 4×16 GB NVIDIA V100 machines, with average training times ranging from 2 h to 140 h.Parameter-conditioning runs are identified as the most expensive.
  • Runtime comparison: Runtime is measured after 10 warmup iterations over 100 runs on one 16 GB NVIDIA V100 with batch size 8.UNO is reported as much slower during multi-GPU cloud training.

B.2.4. U-Net

The U-Net experiments compare embedding and output kernel sizes and implement progressively modified architectures with normalization, residual connections, and restricted attention.

  • U-Net variants: U-Net embeddings and outputs use 3 × 3 kernels, with 1 × 1-kernel variants ablated for fair comparison against FNO and ResNet architectures.The ablation isolates the effect of embedding and output kernel size.
  • U-Net2015: U-Net2015 uses four downsampling layers, a bottleneck, four upsampling layers, batch normalization, and GeLU activations.Its channel multipliers are (2, 2, 2, 2), and it uses no bias weights.
  • U-Netbase: U-Netbase replaces batch normalization with one-group group normalization, adds bias weights, and removes the bottleneck layer.It retains channel multipliers of (2, 2, 2, 2).
  • U-Netmod: U-Netmod uses channel multipliers (1, 2, 2, 4), residual connections, pre-normalization, pre-activations, and zero-initialized second convolution layers.These changes are applied within downsampling and upsampling blocks.
  • U-Netmod,attn: Adding attention to every downsampling and upsampling block made training unstable, so attention was restricted to middle blocks after downsampling.The implementation uses one attention layer with a residual bypass connection.

B.2.5. Parameter Conditioning

The paper conditions PDE surrogates on scalar values such as prediction time window and force strength using projected sinusoidal embeddings. It compares additive, normalization-based, and spatial-spectral mechanisms across U-Net and FNO-like architectures.

  • Embedding: Sinusoidal embeddings encode scalar conditions including prediction time window and force strength.These embeddings provide positional encodings for scalar values.
  • Projection: Two-layer feed-forward projections expand each embedding to 4× hidden channels before combining conditions and injecting them into network blocks.A subsequent linear layer passes the combined embedding to each block.
  • Conditioning: Addition and AdaGN are evaluated as conditioning mechanisms, with AdaGN restricted to U-Net architectures because it requires normalization layers.Addition can also be extended to FNO layers, whereas the implemented FNO architectures lack normalization.
  • Conditioning: Addition injects projected conditions after the first convolution, while AdaGN splits projections into scale and shift terms for normalized activations.Both mechanisms apply conditioning to all blocks.
  • Conditioning: Spatial-Spectral embedding adds conditioning in both Fourier and spatial domains for FNO-like architectures.The paper ablates this mechanism against Addition, including U-FNets with FNO blocks in the downsampling path.

B.3. Shallow water equations.

The shallow-water experiments evaluate surrogate architectures in velocity-function and vorticity-stream-function formulations across one- and two-day prediction windows. They report rollout and one-step errors, architectural ablations, and resource comparisons.

  • Experimental setup: Shallow-water inputs consist of the respective fields at the previous two timesteps, with pressure and vorticity fields normalized for training.The equations use a regular grid with periodic boundary conditions.
  • Two-day predictions: Example 48 h rollouts compare predicted pressure and wind or vorticity fields with ground truth in the two formulations.Velocity-function results use U-F1Netmodes16, while vorticity-stream-function results use U-F2Netmodes16,16.
  • Error evaluation: Figures report rollout and one-step errors for 2-day velocity-function and vorticity-stream-function predictions, plus 1-day velocity-function predictions.Results are averaged over three random seeds and plotted with logarithmic y-axes.
  • Architecture ablations: U-Net ablations compare 1 × 1 and 3 × 3 kernels in encoding and decoding for 2-day shallow-water predictions.The comparisons cover both velocity-function and vorticity-stream-function formulations.
  • Tabulated comparisons: Tables report summed mean-squared errors for shallow-water rollout and one-step predictions across formulations, horizons, and training objectives.A separate table compares parameter count, runtime, and memory for FNO, UNO, and U-FNet architectures.

B.4. Navier-Stokes equations.

The Navier–Stokes experiments evaluate surrogate architectures in velocity-function form using scalar and vector fields on a 128×128 grid. They include rollout, one-step, architectural-ablation, and resource-oriented comparisons.

  • Experimental setup: Navier–Stokes data use 128×128 spatial resolution, ∆t = 1.5 s, viscosity ν = 0.01, buoyancy factor (0, 0.5)^T, and closed-domain boundary conditions.The simulation runs for 21 s and samples 14 time points.
  • Example rollouts: Example ∆t = 1.5 s rollouts compare predicted scalar and vector velocity fields with ground truth for FNO and U-Net surrogates.The FNO example uses FNO96-4modes32,32, while the U-Net example uses U-F1Netmodes16.
  • Error evaluation: Navier–Stokes evaluation reports rollout and one-step errors averaged over three random seeds, with logarithmic y-axes.The experiments use predictions of 1.5 s.
  • Architecture ablations: U-Net ablations compare 1 × 1 and 3 × 3 encoding and decoding kernels for Navier–Stokes velocity-function predictions.The results include both rollout and one-step errors.
  • Tabulated comparisons: Tables report summed mean-squared errors for rollout and one-step predictions under the standard and L2 training objectives.Results are averaged over three random seeds.

B.5. Parameter conditioning.

The parameter-conditioning experiments test whether a single Navier–Stokes surrogate can handle different buoyancy-force values and time windows. They compare conditioning choices and illustrate rollouts at multiple force strengths.

  • Experimental setup: Parameter-conditioned Navier–Stokes inputs use fields from the previous timestep, while rollout examples vary the buoyancy force value.The experiments use the same spatial resolutions and boundary conditions as the main Navier–Stokes setup.
  • Quantitative evaluation: Table 10 reports summed mean-squared errors across training-trajectory counts and time windows, averaged over 208 unseen buoyancy-force values.The table highlights the best model in each model class and compares conditioning choices.
  • Conditioning ablation: Figure 19 compares Addition and AdaGN for U-Net blocks across time windows using one-step errors.For low time windows, AdaGN seems beneficial.
  • Force-strength generalization: Example U-Netmod rollouts at f = 0.21, f = 0.33, and f = 0.48 compare predicted scalar and vector velocity fields with ground truth.Each example uses a 1.5 s prediction window and fields at the last timestep as model inputs.
Loading 2209.15616v2…