PaperlayerView exact revision on arXiv

Source-linked AI summary

Latent Neural Operator for Solving Forward and Inverse PDE Problems

Tian Wang, Chuang Wang

arXiv:2406.03923v5cs.LGcs.NAmath.NA

TL;DR

Neural operators learn PDE solution maps from data, but attention over large sample sequences is computationally costly. LNO uses Physics-Cross-Attention to encode samples into a learnable latent space, learns the operator there, and decodes at query locations. It reports reduced GPU memory and training time alongside state-of-the-art accuracy across several forward and inverse benchmarks.

  • Problem

    Attention-based neural operators face quadratic computational complexity as input sequences grow, while simpler linear attention sacrifices precision.

    Source: §1 · p. 2, §3.1 · p. 4

  • Method

    LNO uses PhCA to encode inputs into a learnable latent space, applies Transformer layers there, and decodes outputs at query locations with inverse PhCA.

    Source: §1 · p. 2, §3.3 · p. 6

  • Results

    50% lower GPU memory and 1.8 times faster training accompany state-of-the-art accuracy on four of six forward benchmarks and a Burgers’ inverse benchmark.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Decoupled observation and prediction locations let LNO decode at unobserved positions, supporting interpolation and extrapolation for forward and inverse problems.

    Source: §1 · p. 2

    For time-dependent equations, intermediate states are decoded into geometric space to compute the loss; the latent space has no imposed physical prior.

    Source: §5 · p. 10

Abstract

from arXiv · show

Neural operators effectively solve PDE problems from data without knowing the explicit equations, which learn the map from the input sequences of observed samples to the predicted values. Most existing works build the model in the original geometric space, leading to high computational costs when the number of sample points is large. We present the Latent Neural Operator (LNO) solving PDEs in the latent space. In particular, we first propose Physics-Cross-Attention (PhCA) transforming representation from the geometric space to the latent space, then learn the operator in the latent space, and finally recover the real-world geometric space via the inverse PhCA map. Our model retains flexibility that can decode values in any position not limited to locations defined in the training set, and therefore can naturally perform interpolation and extrapolation tasks particularly useful for inverse problems. Moreover, the proposed LNO improves both prediction accuracy and computational efficiency. Experiments show that LNO reduces the GPU memory by 50%, speeds up training 1.8 times, and reaches state-of-the-art accuracy on four out of six benchmarks for forward problems and a benchmark for inverse problem. Code is available at https://github.com/L-I-M-I-T/LatentNeuralOperator.

1 Introduction

Neural operators learn PDE solution maps from data, but attention-based models face quadratic costs as sample sequences grow. LNO uses a learnable latent space to target the accuracy–efficiency trade-off, with reported efficiency gains and strong benchmark accuracy.

1 Introduction

  • Attention can model long-distance, quadratic relationships among sample points, but its computational complexity grows quadratically with input-sequence length.Linear attention speeds computation but sacrifices precision because of limited expressive power.Source: §1 · p. 2
  • LNO encodes input samples into a learnable latent space, learns the PDE operator there, and decodes predictions at query locations through inverse PhCA.Unlike approaches that switch between latent and geometric spaces in each layer, LNO transforms data only at the first and last layers.Source: §1 · p. 2
  • 50% lower GPU memory and 1.8 times faster training accompany state-of-the-art accuracy on four of six forward benchmarks and the Burgers’ inverse benchmark.The forward benchmarks are Pipe, Elasticity, Darcy, and Navier-Stokes.Source: §1 · p. 2
  • PhCA decouples input and output sample locations, allowing predictions at positions beyond the observations and supporting interpolation and extrapolation.The passage identifies this flexibility as useful for inverse problems.Source: §1 · p. 2
  • The paper presents LNO as a model for both forward and inverse PDE problems, addressing a setting where operator-learning methods still have room to improve prediction accuracy.Operator learning maps input-output function pairs and can generalize to different initial and boundary conditions.Source: §1 · p. 2

2 Related works

Related work contrasts physics-informed methods, which incorporate PDE knowledge, with data-driven operator learning, and reviews latent-space and inverse-problem approaches. LNO is positioned as an end-to-end learned mapping between geometric and latent spaces.

2.1 Physics-Informed Machine Learning

  • Physics-informed methods encode PDE knowledge in losses or architectures and can be precise, but require instance-specific optimization and are less flexible when equations are unknown or inexact.PINN incorporates the differential equation and initial and boundary conditions into its loss.Source: §2.1 · pp. 2–3

2.2 Operator Learning

  • Operator-learning methods learn input-to-solution function maps from data without requiring explicit knowledge of the PDE.Prior approaches include kernel-integral neural operators, attention variants, and methods that reduce computation through linear attention or factorization.Source: §2.2 · p. 3
  • Prior latent-space methods use per-layer physics attention, cross-attention with orthogonal bases, or graph-based super-nodes; LNO instead learns geometric-to-latent mappings end-to-end without manual bases or extra losses.The related approaches named are Transolver, LSM, and UPT.Source: §2.2 · p. 3

2.3 Inverse Problem

  • Inverse PDE work spans sensing and reconstruction applications and methods including direct sampling, PINNs, neural inverse operators, learnable boundary conditions, and generative initial-condition models.LNO aims to unify forward and inverse solutions in latent space by decoupling observation and prediction positions.Source: §2.3 · p. 3

3 Method

LNO frames forward and inverse PDE tasks as operator learning, mapping sampled functions into a learnable latent space for operator modeling before decoding outputs at queried geometric positions. PhCA provides the cross-space transformations while separating location queries and keys from physics values.

3.1 Problem Setup

  • Forward problems map boundary or coefficient samples to solution values, while inverse problems use partial observations to infer boundary conditions or PDE coefficients.Source: §3.1 · p. 4

3.2 Latent Neural Operator

  • The LNO architecture lifts inputs, encodes them into a latent space, models the operator with self-attention layers, and decodes outputs to geometric space.Source: §3.2 · p. 4
  • Encoding reduces the processed sequence from N_in samples to M latent tokens, with M ≪ N_in, making feature extraction and conversion more efficient.Source: §3.2 · p. 5
  • The latent operator uses stacked Transformer blocks with self-attention, and decoding applies PhCA at output query positions to recover geometric-space values.Source: §3.2 · p. 5

3.3 Physics-Cross-Attention

  • LNO uses PhCA to transform representations between geometric and learnable latent spaces, with Transformer layers modeling the operator between encoding and decoding.Source: §3.2 · p. 4, §3.3 · p. 6
  • PhCA uses learnable latent-position queries and separates position-based queries and keys from physics values, allowing predictions at positions distinct from observations.Source: §3.2 · p. 5, §3.3 · p. 6

4 Experiments

The experiments evaluate LNO on six forward benchmarks and a Burgers inverse problem, then assess efficiency, resolution generalization, and design choices through ablations.

4.1 Accuracy for the forward problems

  • LNO achieves the best accuracy on Darcy, NS2d, Elasticity, and Pipe, while remaining close to SOTA on Airfoil and Plasticity.The comparison covers six forward-problem benchmarks.Source: Table 1 · p. 7, §4.1 · p. 7

4.2 Accuracy for the inverse problem

  • LNO outperforms DeepONet and GNOT in the inverse problem’s completer stage and is most accurate overall when used for both stages.Propagator performance declines as the completer’s subdomain reconstruction error increases.Source: §4.2 · p. 8, Table 2 · p. 8, Table 3 · p. 8

4.3 Efficiency

  • Compared with Transolver, LNO reduces average parameter count by 30% and GPU memory consumption by 50%, with 1.8x training acceleration.The comparison spans all six forward-problem benchmarks.Source: §4.3 · pp. 8–9, Table 4 · p. 8

4.4 Generalization

  • Trained on Darcy at 43 × 43 resolution, LNO consistently outperforms other methods when tested at different discretization resolutions.The tested resolutions range from 43 × 43 to 421 × 421.Source: §4.4 · p. 9, Table 5 · p. 9

4.5 Ablation Study

  • Replacing PhCA with Physics-Attention verifies that decoupling observation and prediction positions is necessary for LNO’s flexibility and accuracy.With Physics-Attention, observation and prediction positions must remain the same.Source: §4.5 · p. 9, Table 6 · p. 10
  • Scaled dot-product quadratic attention is more expressive than linear attention, and its latent-space use does not incur significant overhead.The comparison tests Galerkin-type, efficient, and Nystrom attention alternatives.Source: §4.5 · p. 9, Table 6 · p. 10
  • Shared MLPs in the PhCA encoder and decoder perform better on most forward benchmarks than separate MLPs, except on Plasticity.The ablation compares shared and non-shared weights.Source: §4.5 · p. 9, Table 7 · p. 10
  • Increasing latent sampling numbers improves performance to saturation and then slight decline on Airfoil, Elasticity, Plasticity, and Pipe, while Darcy and NS2d improve through M = 512.Inverse-problem completer and propagator trends are similar; the passage attributes differing forward trends to whether benchmarks include physics quantities as well as location information.Source: §4.5 · pp. 9–10, Figure 3 · p. 10

5 Conclusions

The paper explores PDE characterization in a learnable latent space and validates LNO for forward and inverse problems. It identifies limitations in time-dependent processing and the absence of physical priors in the latent space.

5 Conclusions

  • LNO explores PDE characterization in a learnable latent space and is validated for forward and inverse problems.The paper presents PhCA for encoding inputs into and decoding outputs from that space.Source: §5 · p. 10
  • For time-dependent equations, intermediate states are decoded into geometric space to compute the loss.Additional training strategies may enable the dynamic evolution process to run directly in latent space, potentially enhancing efficiency.Source: §5 · p. 10
  • The latent space has no imposed prior; exploring physical priors such as symmetry is left for future work.The paper identifies possible precision improvement as a motivation for this exploration.Source: §5 · p. 10

A Scaling

Scaling experiments examine how LNO performance varies with Transformer depth and token width. Performance peaks at depth 4 or 8, while benefits from width gradually saturate at dimension 256.

A Scaling

  • The scaling experiments vary model depth, defined as the number of stacked Transformer blocks, and model width, the token dimension.Figures 4 and 5 show these effects across six forward benchmarks.Source: A Scaling · p. 14, Figure 4 · p. 14, Figure 5 · p. 14
  • LNO performance initially improves and then declines as depth increases, with optimal performance at 4 or 8 Transformer blocks.The authors report that simply adding blocks does not necessarily improve accuracy.Source: A Scaling · p. 14, Figure 4 · p. 14
  • LNO performance continues benefiting from greater token width, but gradually saturates when token dimension reaches 256.Figure 5 covers six forward-problem benchmarks.Source: A Scaling · p. 14, Figure 5 · p. 14

B Observation Ratio and Locations for the Inverse PDE Task

The inverse-task experiments examine observation intervals and locations. LNO retains accuracy with sparse observations, but spatial undersampling is more damaging than temporal undersampling and is poorly compensated by temporal oversampling.

B Observation Ratio and Locations for the Inverse PDE Task

  • The inverse problem remains accurately solvable when temporal and spatial sampling intervals are both 16, at approximately 0.4% observation ratio.Figure 8 shows LNO as both completer and propagator across sampling intervals.Source: B Observation Ratio and Locations for… · p. 15, Figure 8 · p. 15
  • Increasing the spatial sampling interval compromises accuracy more than increasing the temporal interval, and temporal oversampling hardly compensates for spatial undersampling.These findings concern fixed observation situations.Source: B Observation Ratio and Locations for… · p. 15, Figure 7 · p. 15, Figure 8 · p. 15
  • Figure 6 illustrates a two-stage strategy for differing observation regions: interpolate within the subdomain, then extrapolate from it to the whole domain.The figure visualizes this strategy for Burgers’ equation.Source: Figure 6 · p. 15

C Significance

The paper reports means and standard deviations from repeated trials for forward and inverse experiments. The supplied table captions specify the trial counts and evaluation scopes, but not outcome values.

C Significance

  • Forward and inverse experiments use repeated random-seed runs, with means and standard deviations reported for significance.Figure 8 and Figure 9 are cited for these results.Source: C Significance · p. 16
  • Table 8 reports mean and standard deviation across five independent trials for forward-problem models.The caption identifies the table’s scope but supplies no performance values.Source: Table 8 · p. 16
  • Table 9 reports mean and standard deviation across three independent trials for inverse-problem models, including completer observation rates and a propagator setting.Relative MAEs are expressed as percentages.Source: Table 9 · p. 16

D Full Implementation Details

This section reports the implementation and training configurations for LNO on forward and inverse problems. Experiments used one or two NVIDIA RTX 3090 GPUs.

D Full Implementation Details

  • Experiments were conducted on one or two NVIDIA RTX 3090 GPUs.Source: D Full Implementation Details · p. 16
  • Forward models were trained for 500 epochs with relative L2 error, AdamW, and OneCycleLr, using an initial learning rate of 10^-3.Transolver was reproduced following its appendix B.3 implementation, and data splits matched appendix B.1.Source: D Full Implementation Details · p. 16
  • Inverse-problem training used mean squared error and StepLr, with 4096 training, 128 validation, and 128 test samples.DeepONet used four Transformer blocks with 128-dimensional branch tokens and a three-layer trunk; GNOT used four blocks, 96-dimensional tokens, and one expert.Source: D Full Implementation Details · p. 16
  • Table 10 presents LNO hyperparameters and training configurations for the forward and inverse problems in Sections 4.1 and 4.2.Source: Table 10 · p. 16

Related papers

Loading 2406.03923v5…