PaperlayerView exact revision on arXiv

Source-linked AI summary

Solving High-Dimensional PDEs with Latent Spectral Models

Haixu Wu, Tengge Hu, Huakun Luo, Jianmin Wang, Mingsheng Long

arXiv:2301.12664v3cs.LGphysics.comp-ph

TL;DR

High-dimensional PDEs pose costly computation and complex input-output mappings that are difficult for single-operator coordinate-space methods. LSM projects data into a compact latent space and learns multiple basis operators; it achieves consistent state-of-the-art results across seven benchmarks, with an 11.5% relative error reduction. Its theoretical analysis assumes a finite coordinate set.

  • Problem

    High-dimensional discretized PDE data makes solving computationally costly and complex input-output mappings difficult to approximate precisely.

    Source: §1 · p. 1

  • Method

    LSM uses attention-based hierarchical projection to map data into compact latent space and a neural spectral block that learns multiple basis operators.

    Source: §1 · p. 2

  • Results

    Across seven benchmarks, LSM achieves consistent state-of-the-art results, with an 11.5% relative error reduction.

    Source: §1 · p. 2

  • Takeaways & Limitations

    LSM is presented as a promising PDE solver for real-world applications, with a favorable trade-off between accuracy and efficiency.

    Source: §5 · p. 9

    The theoretical derivations assume a finite coordinate set, matching settings where PDE inputs are observed on finite grids, meshes, or point clouds.

    Source: §3.2 · p. 5

Abstract

from arXiv · show

Deep models have achieved impressive progress in solving partial differential equations (PDEs). A burgeoning paradigm is learning neural operators to approximate the input-output mappings of PDEs. While previous deep models have explored the multiscale architectures and various operator designs, they are limited to learning the operators as a whole in the coordinate space. In real physical science problems, PDEs are complex coupled equations with numerical solvers relying on discretization into high-dimensional coordinate space, which cannot be precisely approximated by a single operator nor efficiently learned due to the curse of dimensionality. We present Latent Spectral Models (LSM) toward an efficient and precise solver for high-dimensional PDEs. Going beyond the coordinate space, LSM enables an attention-based hierarchical projection network to reduce the high-dimensional data into a compact latent space in linear time. Inspired by classical spectral methods in numerical analysis, we design a neural spectral block to solve PDEs in the latent space that approximates complex input-output mappings via learning multiple basis operators, enjoying nice theoretical guarantees for convergence and approximation. Experimentally, LSM achieves consistent state-of-the-art and yields a relative gain of 11.5% averaged on seven benchmarks covering both solid and fluid physics. Code is available at https://github.com/thuml/Latent-Spectral-Models.

1. Introduction

High-dimensional PDEs are costly to solve and their complex input-output mappings are difficult to approximate; LSM addresses these challenges by projecting data into a compact latent space and learning multiple basis operators. It achieves state-of-the-art results across seven benchmarks covering solid and fluid physics.

1. Introduction

  • High-dimensional discretized PDE data raises computation costs and makes complex mappings difficult for rough deep models to approximate.The paper identifies efficient and precise approximation of these mappings as a key problem in PDE solving.Source: §1 · p. 1
  • Prior operator-learning methods learn mappings as a whole, while multiscale models still operate in coordinate space and retain high-dimensionality challenges.The paper argues that a single operator may not cover sufficiently complex high-dimensional mappings.Source: §1 · pp. 1–2
  • LSM uses an attention-based hierarchical projection network to reduce high-dimensional data into a compact latent space in linear time.Latent tokens act as physical prompts, highlighting physics properties while removing redundant coordinate information.Source: §1 · p. 2
  • Its neural spectral block learns multiple basis operators to approximate complex mappings and has theoretical guarantees for approximation and convergence.The design is inspired by classical spectral methods.Source: §1 · p. 2
  • LSM achieves an 11.5% relative error reduction against previous state-of-the-art models, averaged across seven PDE-solving benchmarks.The benchmarks cover representative PDEs in solid and fluid physics.Source: §1 · p. 2

2. Preliminaries

Spectral methods approximate PDE solutions with finite sums of basis functions and optimize their coefficients. Deep PDE models include equation-constraint and operator-learning approaches; LSM instead projects data into a compact latent space and decomposes complex mappings into multiple basis operators.

2.1. Spectral Methods

  • Spectral methods approximate a PDE solution as a finite sum of orthogonal basis functions, then optimize their coefficients to better satisfy the PDE.The number of basis functions is a hyperparameter, and the methods have approximation and convergence properties.Source: §2.1 · p. 2

2.2. Deep Models for PDEs

  • Equation-constraint methods optimize a parameterized solution against PDE and boundary conditions, but require explicit PDE formalizations that may be hard to obtain.The paper therefore focuses on operator learning, which does not require explicit PDE formalizations.Source: §2.2 · p. 3
  • Operator-learning methods approximate mappings between input-output pairs; prior designs include Fourier, multiwavelet, and attention-based operators.LSM differs by decomposing the complex nonlinear operator into several basis operators.Source: §2.2 · p. 3
  • Multiscale models add hierarchical processing, while LSM projects high-dimensional data into compact latent space to focus on essential physical information.The projection is intended to avoid redundant coordinate-space information.Source: §2.2 · p. 3

3. Latent Spectral Models

LSM reframes high-dimensional PDE solving as a compact latent-space problem: hierarchical attention projectors reduce coordinate data, and a neural spectral block approximates latent mappings with multiple basis operators.

3. Latent Spectral Models

  • LSM divides PDE solving into coordinate-to-latent projection, latent-space solving, and projection back to coordinate space.The framework uses the hierarchical projection network for the projections and the neural spectral block for latent solving.Source: §3 · p. 3

3.1. Hierarchical Projection Network

  • Learned latent tokens act as physical prompts to compress coordinate observations into C latent tokens, with projector computation linear in the coordinate-set size, O(|D|).The tokens are shared across input-output pairs and optimized to cover common PDE constraints.Source: §3.1 · p. 4
  • The patchified multiscale architecture addresses different physical states across scales and regions, sharing parameters within a scale but not across scales.Outputs are spliced across patches and successively upsampled from coarse to fine.Source: §3.1 · pp. 4–5

3.2. Neural Spectral Block

  • The neural spectral block decomposes complex mappings into multiple learned trigonometric basis operators, with a residual connection and parameters shared within each scale.The basis-operator parameters are independent across scales.Source: §3.2 · p. 5
  • Classical spectral approximation slows as dimension M grows; LSM instead formulates solving over independent latent tokens restricted to [0, π].The paper identifies high-dimensional spectral convergence as a motivation for solving in latent space.Source: §3.2 · p. 6
  • For Lipschitz functions on [0, π], the neural spectral block has a uniform convergence guarantee with a constant independent of f and N, plus universal approximation capacity.The theorem states the convergence speed, with K3 independent of f and N.Source: §3.2 · pp. 5–6

4. Experiments

LSM is evaluated across seven solid- and fluid-physics benchmarks with varied geometries, and achieves state-of-the-art performance across them. The experiments also examine component ablations, efficiency, resolution robustness, transferability, and learned representations.

4. Experiments

  • The seven benchmarks cover solid and fluid PDEs, varied geometries, and 2D or 3D data; the comparison includes fourteen established models.The benchmark geometries include regular grids, point clouds, and structured meshes; Elasticity-G is interpolated from Elasticity-P.Source: §4 · p. 6

4.1. Main Results

  • 11.5% average relative error reduction over the previous best method across seven benchmarks, with gains including 13.0% on Elasticity-G and 15.6% on Darcy.The benchmarks cover both solid and fluid physics; Table 2 records MSE, for which smaller values are better.Source: §4.1 · pp. 6–7, Table 2 · p. 7
  • Removing the neural spectral block raises estimation error by 13.8% on Elasticity-P and 13.3% on Darcy, while removing the projector causes serious performance drops.These ablations support the roles of latent-space solving and the neural spectral block.Source: §4.1 · p. 7, Table 3 · p. 7
  • Multiscale design fits Darcy's multiphase flow, while patchify is essential for Elasticity-P's local information; alternative projectors impair accuracy or efficiency.Average pooling removes coordinate information but, without latent tokens as physics prompts, fails to capture essential physical information.Source: §4.1 · pp. 7–8, Table 3 · p. 7
  • In the Airfoil showcase, LSM is the only model reported to precisely capture the shock wave; it also surpasses FNO and U-NO on two solid-physics showcases.Figure 4 shows Darcy and Airfoil for fluid physics, and Elasticity-P and Plasticity for solid physics; Plasticity is shown at T = 20.Source: §4.1 · p. 8, Figure 4 · p. 8

4.2. Model Analysis

  • On solid-physics benchmarks, LSM is more accurate than U-NO by 15.6% on Elasticity-P, 12.8% on Elasticity-G, and 26.5% on Plasticity.For fluid physics, LSM is reported more accurate and efficient than HT-Net, WMT, and U-Net; F-FNO remains lightweight with comparable running time and accuracy.Source: §4.2 · p. 8, Figure 5 · p. 8
  • The latent input tokens for Navier–Stokes show a process similar to the input features over time, while Darcy's latent representation transforms heterogeneous inputs and outputs.Figure 6 visualizes input features, latent input and output tokens, and output features in a 2D PCA projection of test-set data.Source: §4.2 · pp. 8–9, Figure 6 · p. 9
  • On Darcy, LSM consistently surpasses baselines from 32×32 to 1024×1024; transfer from Pipe to Airfoil is positive across limited-data settings.LSM performs best in both pre-trained and scratch cases; the benchmarks share Navier–Stokes equations but have distinct boundary conditions.Source: §4.2 · p. 9, Figure 7 · p. 9, Table 4 · p. 9

5. Conclusions and Future Work

The paper presents LSM as a solver for high-dimensional PDEs that projects data into a compact latent space and approximates complex mappings with a neural spectral block. It reports state-of-the-art performance across solid and fluid benchmarks, while identifying cross-PDE generalization as future work.

5. Conclusions and Future Work

  • LSM uses hierarchical projection to reduce high-dimensional data into compact latent space, then applies a neural spectral block to approximate complex mappings under theoretical guarantees.The conclusion characterizes the resulting solver as efficient.Source: §5 · p. 9
  • LSM achieves state-of-the-art results across solid- and fluid-physics benchmarks with a favorable accuracy-efficiency trade-off.The conclusion describes LSM as promising for real-world PDE-solving applications.Source: §5 · p. 9
  • Future work will explore LSM's generalization across different PDEs toward a foundation model.Source: §5 · p. 9

A. Proofs of Theorems 3.6

The proof reduces approximation of a Lipschitz function on [0, π] to trigonometric approximation of a periodic extension of its residual. It then transfers convergence back to the original function.

A. Proofs of Theorems 3.6

  • For Lipschitz f on [0, π], the proof defines the residual g(x) = f(x) − x and shows g remains Lipschitz.Source: A. Proofs of Theorems 3.6 · p. 12
  • The proof extends g to a 2π-periodic function while preserving the Lipschitz condition, allowing use of trigonometric approximation.The extension is built through an intermediate even function on [−π, π].Source: A. Proofs of Theorems 3.6 · pp. 12–13
  • After the periodic extension is approximated, the proof recovers an approximation to f on [0, π] by adding x to the approximation of g.The text states that the extended function converges and then relates its approximation to f.Source: A. Proofs of Theorems 3.6 · p. 13

B. Details for Benchmarks

The appendix describes the solid, fluid, and porous-flow benchmark tasks and their input-output data. Showcases report boundary performance across several geometries and long-term predictions for Navier–Stokes flow.

B.1. Solid Material

  • Elasticity-P and Elasticity-G predict inner stress from material structure; they differ in representation, using 972-point clouds versus 41 × 41 grids.Elasticity-G is interpolated from Elasticity-P.Source: §B.1 · p. 14
  • Plasticity predicts mesh-point deformation over 20 future time steps from the die shape in a 101 × 31 structured mesh.The task models a plastic material struck from above by an arbitrary-shaped die.Source: §B.1 · p. 14

B.2. Navier-Stokes Equation

  • The Navier–Stokes dataset predicts 10 future frames from 10 past frames on a 64 × 64 field, with viscosity set to 10^-5.It simulates incompressible, viscous flow on the unit torus.Source: §B.2 · p. 15

B.3. Darcy Flow

  • Darcy flow uses the diffusion coefficient as input and the solution as output on an 85 × 85 grid, with external force fixed at 1.The dataset models two-dimensional flow in a unit box.Source: §B.3 · p. 15

C. More Showcases

  • The seven-benchmark showcases report strong predictions, with coordinate-wise errors shown for comparison.Figure 8 includes all seven benchmarks.Source: C. More Showcases · p. 16, Figure 8 · p. 16
  • LSM significantly surpasses other baselines at boundaries across four geometries, demonstrating capability in learning physical constraints.The cited cases are Figure 8(a), (b), (c), and (f).Source: C. More Showcases · p. 16, Figure 8 · p. 16
  • On Navier–Stokes, LSM predicts future velocity accurately through T = 20 while other methods’ performance drops from T = 18 to T = 20.Figure 8(g) shows the comparison over time.Source: C. More Showcases · p. 16, Figure 8 · p. 16

D. Full Ablations

The appendix reports ablations of LSM’s components and neural spectral block, then compares performance across input resolutions. It reports stable Darcy performance across resolutions and a larger Navier–Stokes advantage at higher resolutions.

D. Full Ablations

  • The full ablations assess projection, multiscale design, patchify, and the spectral block, with efficiency measured on 256 × 256 inputs at batch size 1.The ablation table includes component replacements and removals; “/” denotes out-of-memory.Source: D. Full Ablations · p. 16, Table 6 · p. 17
  • Replacing the neural spectral block with MLP or FNO makes performance close to removing the block; latent-space FFT with FNO can damage it further.The authors attribute this concern to latent tokens being independent.Source: D. Full Ablations · p. 17, Table 7 · p. 17
  • Replacing LSM’s trigonometric basis with polynomial basis decreases performance, and the authors leave exploration of other bases and their theory to future work.The detailed neural spectral block ablations are reported in Table 7.Source: D. Full Ablations · p. 17, Table 7 · p. 17

E. Performance Under Various Resolutions

F. Additional Experiments on Burger’s Equation

The paper adapts LSM and selected baselines to one-dimensional inputs for experiments on Burger’s equation. LSM performs well across the tested resolutions.

F. Additional Experiments on Burger’s Equation

G. Hyperparameter Sensitivity

The sensitivity experiments vary latent-token count, basis-operator count, and scale count one at a time. The authors report stable performance over selected ranges and choose settings intended to trade off efficiency and performance.

G. Hyperparameter Sensitivity

  • Performance is stable across choices of latent-token count C when N = 24 and K = 5.The authors suggest this may come from the equivalence of different latent tokens.Source: G. Hyperparameter Sensitivity · p. 18
  • Increasing basis-operator count N generally improves results but adds computation and optimization difficulty; performance drops slightly at N = 40 and N = 48.N = 0 is equivalent to removing the neural spectral block, which causes a serious performance drop.Source: G. Hyperparameter Sensitivity · p. 18
  • Adding scales K generally improves performance, but too many scales are unimplementable because of input-resolution limits.The experiment fixes C = 4 and N = 24.Source: G. Hyperparameter Sensitivity · p. 18
  • Overall, LSM is stable for C = 3–5, N = 24–40, and K = 5–7; the chosen settings are C = 4, N = 24, and K = 5.The authors describe these settings as an efficiency-performance trade-off.Source: G. Hyperparameter Sensitivity · p. 19, Table 11 · p. 19

H. Training Stability

LSM ranks consistently across solid and fluid physics benchmarks, with reported efficiency and stability comparisons. The appendix details its training setup, multiscale operations, task-specific configurations, and evaluation protocol.

H. Training Stability

  • LSM shows comparable training stability to FNO while maintaining state-of-the-art performance across benchmarks.Figure 9 plots test-set MSE against training epochs.Source: H. Training Stability · p. 19, Figure 9 · p. 19
  • Across five repeats, LSM performance standard deviations were within 0.0001 for Elasticity-P, Elasticity-G, Plasticity, Darcy, and Airfoil, and 0.0002 for Navier-Stokes and Pipe.Source: H. Training Stability · p. 19

I. Implementation Details

  • Experiments use PyTorch on one NVIDIA RTX 3090 24GB GPU, with each method's final-epoch performance recorded as its final result.Source: §I · p. 19

I.1. Model Configurations

  • LSM configurations use zero padding to address division in the model configurations.Table 12 is identified as the model-configuration table.Source: §I.1 · p. 20, Table 12 · p. 20

I.2. Model Architecture

  • Patchify splits coordinates into equal-sized non-overlapping regions; downsampling pools and convolves features, while upsampling interpolates coarser features and fuses scales with local convolution.Interpolation is bilinear for 2D data and trilinear for 3D data; de-patchify splices the local patches.Source: §I.2 · p. 20

I.3. Benchmark Construction

  • Plasticity uses 3D operations for boundary-condition inputs and mesh displacement over time, whereas autoregressive Navier-Stokes uses 2D LSM to predict frames step by step.The neural spectral block is unchanged across the 2D and 3D versions; baselines are trained with L2 loss for comparison fairness.Source: §I.3 · p. 20

J. Efficiency

  • LSM achieves consistent state-of-the-art performance across seven tasks, while the top-five methods differ between solid and fluid benchmarks except LSM and F-FNO.LSM ranks first in both groups; F-FNO ranks fifth in both.Source: J. Efficiency · p. 21
  • At 512 × 512 and 1024 × 1024 inputs, LSM is clearly faster than U-NO (rank 2nd) and U-Net (rank 3rd).The comparison fixes a regular-grid input, one channel, and batch size 1, and records parameters, GPU memory, and running time across resolutions.Source: J. Efficiency · p. 21

K. Limitations

LSM's potential limitation is generality across different PDEs, while the paper presents it as a solver with applications across physics domains and limited-data transfer settings. The authors frame its latent-space approach as a supplement to operator learning and report no potential ethical risk or negative social impact for this work.

K. Limitations

  • A zero-shot solver generalizing across different PDEs remains an open problem; coefficient perturbations can seriously change PDE properties, placing a universal solution beyond this paper's scope.The authors leave the question of whether a universal solution exists for future work.Source: K. Limitations · p. 21

L. Societal Impacts

  • The authors identify airfoil design, civil-engineering load-bearing tests, and weather forecasting as potential applications, and note transferability in limited-data scenarios.They describe that transferability as important for fast adaptation to new scenarios.Source: L. Societal Impacts · p. 21
  • LSM's latent-space approach is presented as a new technology roadmap and a potential supplement to the operator-learning community.Source: L. Societal Impacts · p. 21
  • The paper reports no potential ethical risk or negative social impact, noting that its datasets come from public tools and follow their licenses.Source: L. Societal Impacts · p. 21
Loading 2301.12664v3…