PaperlayerView exact revision on arXiv

Source-linked AI summary

Improved architectures and training algorithms for deep operator networks

Sifan Wang, Hanwen Wang, Paris Perdikaris

arXiv:2110.01654v2cs.LGmath.NAphysics.comp-phstat.ML

TL;DR

DeepONets can learn PDE solution operators without paired input-output observations, but their training favors larger-magnitude functions and can suffer from gradient pathologies. The paper uses NTK-guided adaptive weighting and a more gradient-resilient architecture, reporting 10-50x predictive-accuracy improvements across PDE benchmarks.

  • Problem

    DeepONet training is biased toward larger-magnitude target functions, while effective normalization and loss balancing remain unclear for infinite-dimensional operator learning.

  • Method

    The paper analyzes DeepONet and physics-informed DeepONet training with NTK theory, adaptively weights loss terms, and introduces an architecture that better preserves input-signal propagation.

  • Results

    10-50x improvement in predictive accuracy is reported across benchmarks learning parametric PDE solution operators without paired input-output observations.

  • Takeaways & Limitations

    The developments provide training insights and improve DeepONet accuracy across PDE solution-operator benchmarks with different dynamics.

  • Takeaways & Limitations

    Adaptive NTK weighting can produce extremely large weights in the presence of vanishing gradients, potentially making training unstable.

Abstract

from arXiv · show

Operator learning techniques have recently emerged as a powerful tool for learning maps between infinite-dimensional Banach spaces. Trained under appropriate constraints, they can also be effective in learning the solution operator of partial differential equations (PDEs) in an entirely self-supervised manner. In this work we analyze the training dynamics of deep operator networks (DeepONets) through the lens of Neural Tangent Kernel (NTK) theory, and reveal a bias that favors the approximation of functions with larger magnitudes. To correct this bias we propose to adaptively re-weight the importance of each training example, and demonstrate how this procedure can effectively balance the magnitude of back-propagated gradients during training via gradient descent. We also propose a novel network architecture that is more resilient to vanishing gradient pathologies. Taken together, our developments provide new insights into the training of DeepONets and consistently improve their predictive accuracy by a factor of 10-50x, demonstrated in the challenging setting of learning PDE solution operators in the absence of paired input-output observations. All code and data accompanying this manuscript are publicly available at \url{https://github.com/PredictiveIntelligenceLab/ImprovedDeepONets.}

1 Introduction

Operator-learning methods can approximate maps between infinite-dimensional function spaces and support fast surrogates for complex physical processes. This work addresses training challenges in DeepONets, including magnitude bias, uncertain normalization effects, and loss-balancing issues.

  • Operator-learning techniques approximate maps between infinite-dimensional Banach spaces and can support fast surrogate modeling of complex physical processes.
  • Physics-informed DeepONets enable self-supervised learning of parametric PDE solution operators under varied initial or boundary conditions, random inputs, and geometries.
  • Existing neural-network initialization and normalization practices may require revision for operator learning in infinite-dimensional function spaces.
  • DeepONets exhibit a bias toward approximating target functions with larger magnitudes, motivating output-function normalization approaches.
  • The paper uses NTK theory to analyze training dynamics, adaptively re-weights examples, and introduces an architecture resilient to vanishing gradients.
  • The proposed developments improve DeepONet predictive accuracy by 10-50x across challenging PDE solution-operator benchmarks without paired input-output observations.

2 Preliminaries

NTK theory characterizes neural-network training through a kernel whose eigenvalues determine convergence rates. DeepONets represent operators by combining branch features of input functions with trunk features of evaluation coordinates, while physics-informed variants enforce PDE and boundary constraints without paired observations.

  • 2.1 The Neural Tangent Kernel (NTK) of a fully-connected neural network: The NTK provides a theoretical framework for convergence and generalization in wide, over-parameterized neural networks.
  • 2.1 The Neural Tangent Kernel (NTK) of a fully-connected neural network: In the infinite-width, infinitesimal-learning-rate regime, neural-network training behaves like a linearized model governed by a deterministic NTK.
  • 2.1 The Neural Tangent Kernel (NTK) of a fully-connected neural network: NTK eigenvalues determine the convergence rate of the total training error, with positive eigenvalue components decaying at rates e^-λ_i t.
  • 2.2 DeepONets and Physics-informed DeepONets: DeepONets learn nonlinear operators by processing input functions with a branch network and evaluation coordinates with a trunk network.
  • 2.2 DeepONets and Physics-informed DeepONets: The branch and trunk feature embeddings are merged through an inner product to produce the predicted output function.
  • 2.2 DeepONets and Physics-informed DeepONets: Physics-informed DeepONets impose PDE constraints at randomly sampled collocation points and boundary conditions at sampled domain-boundary points.
  • 2.2 DeepONets and Physics-informed DeepONets: Physics-informed DeepONets can learn parametric PDE solution operators entirely without paired input-output observations.

3 Methods

The methods analyze DeepONet training bias through NTK theory, introduce adaptive loss weighting, and propose an architecture designed to reduce vanishing-signal problems. Experiments on the anti-derivative operator show that output magnitude biases accuracy, while NTK weighting and architectural changes improve training behavior.

  • 3.1 Gradients pathologies in DeepONets: DeepONets can produce poorer approximations for multi-scale input parameters, motivating analysis of output-magnitude bias in an anti-derivative benchmark.The benchmark learns the solution operator mapping a forcing term to its ODE solution.
  • 3.1 Gradients pathologies in DeepONets: The benchmark samples 10^4 forcing terms from a Gaussian random field with output scales k ∈ [0.01, 100] and uses one randomly selected solution observation per pair.Inputs use 100 fixed sensors, and solutions are generated with RK45 integration.
  • 3.1 Gradients pathologies in DeepONets: Large-output-scale samples are predicted more accurately, revealing a positive correlation between test error and output scale k despite balanced training data.The model uses three-layer ReLU MLP branch and trunk networks with 100 neurons per hidden layer.
  • 3.2 NTK analysis of physics-informed DeepONets: Large-value input functions create larger loss changes and back-propagated gradients, causing their mean-square errors to dominate training.The gradient mechanism is examined by monitoring each individual loss term rather than only the aggregate loss.
  • 3.2 NTK analysis of physics-informed DeepONets: NTK analysis relates the eigenvalues of the physics-informed DeepONet kernel to system stiffness and the convergence rate of total training error.The generalized NTK framework yields an ODE governing physics-informed DeepONet outputs.
  • 3.3 NTK-guided weights for physics-informed DeepONets: Algorithm 1 adaptively weights individual loss terms; with α = 1, NTK weights equalize their maximum convergence rate and counteract gradient-magnitude imbalance.The weighting rule penalizes large gradients and amplifies small ones, but can become unstable when gradients vanish.
  • 3.3 NTK-guided weights for physics-informed DeepONets: Among weighting schemes, NTK weights with α = 1 achieve the best predictive accuracy and smallest standard deviation across output scales in the anti-derivative test.The comparison uses identical architecture and hyperparameter settings across weighting schemes.
  • 3.4 An improved DeepONet architecture: The proposed architecture embeds branch and trunk inputs with encoders and merges those embeddings by point-wise multiplication in every hidden layer.This modifies standard DeepONets to improve signal propagation and represent nonlinearities.

4 Results

The experiments evaluate physics-informed DeepONets across weighting schemes and architectures using mean relative L2 test error. Modified architectures and adaptive or NTK-guided weighting improve accuracy across advection, Burgers’ equation, and Stokes flow benchmarks.

  • Predictive accuracy is assessed by the mean relative L2 error over test examples, comparing weighting schemes and DeepONet architectures.The experiments use physics-informed DeepONets trained with the listed weighting schemes and architectures.
  • Architectures: Two encoders inject embedded input-sample and coordinate features into every hidden layer of the branch and trunk networks before their outputs are merged.The modified architecture retains conventional DeepONet output merging while exchanging information earlier in the network.
  • 4.1 Advection equation: The advection benchmark uses positive random-field coefficient functions, with the DeepONet learning the associated spatio-temporal PDE solution operator.Training and testing use sampled input functions and numerical solutions generated on a uniform grid.
  • 4.1 Advection equation: A modified DeepONet with α = 1/2 improves advection accuracy by 6x over the baseline and reaches a relative L2 error of 0.73%.The improved model better captures sharp gradients than the conventional physics-informed DeepONet.
  • 4.2 Burgers’ equation: For Burgers’ equation, NTK weights with α = 1 reduce test error to 3.69%, at least 5x more accurate than the baseline physics-informed DeepONets.A representative best model reaches 1.19% relative L2 error versus 30.72% for the worst model and captures the viscous shock wave.
  • 4.3 Stokes flow: In Stokes flow, the modified DeepONet is evaluated after 2 × 10^5 gradient-descent iterations, while adaptive weights become especially large for the continuity equation.The continuity residual weight is two orders of magnitude greater than the other two residual-weight distributions.

5 Discussion

The work improves DeepONet training by calibrating loss-term weights and introducing an architecture that better preserves input-signal propagation. These developments yield large predictive-accuracy improvements, while physics-informed DeepONets retain theoretical and methodological open questions.

  • 5 Discussion: The proposed training algorithm assigns weights to individual loss terms to calibrate convergence rates and back-propagated gradients.The procedure is based on an analysis of DeepONet training dynamics using Neural Tangent Kernel theory.
  • 5 Discussion: A conventional physics-informed DeepONet completely fails on a representative Stokes test example, with relative L2 errors of 74.32%, 82.64%, and 99.63% for u, v, and p.These errors are reported for the velocity components and pressure, respectively.
  • 5 Discussion: The new DeepONet architecture better preserves input signals through successive network layers and outperforms conventional DeepONets on abstract nonlinear operators.The authors report consistent empirical improvement for this architecture.
  • 5 Discussion: 10-50x improvement in predictive accuracy is reported across benchmarks learning parametric-PDE solution operators with different dynamics.The result concerns the combined training and architectural developments.
  • 5 Discussion: The approximation properties of physics-informed DeepONets trained with PDE constraints remain insufficiently understood.The authors contrast this gap with established operator universal-approximation results and a-priori error estimates for physics-informed neural networks.
  • 5 Discussion: The proposed training algorithm analyzes full-batch gradient descent, while the effects of mini-batch gradient descent with momentum remain poorly understood.

A Nomenclature

The nomenclature appendix summarizes the symbols and notation used throughout the work.

  • A Nomenclature: Table 8 summarizes the main symbols and notation used in the work.
  • A Nomenclature: The table provides a reference for notation used across the paper.
  • A Nomenclature: The nomenclature is presented in a dedicated table labeled Table 8.

B Hyper-parameter settings

The appendix identifies the default hyper-parameter settings and the physics-informed DeepONet architectures used for the benchmarks.

  • B Hyper-parameter settings: Table 9 lists default hyper-parameter settings for each benchmark unless otherwise stated.
  • B Hyper-parameter settings: The benchmark settings and architectures are documented in separate appendix tables.
  • B Hyper-parameter settings: Table 10 lists the physics-informed DeepONet architectures used for each benchmark unless otherwise stated.

C Computational cost

The computational-cost appendix reports training times for physics-informed DeepONet models across benchmarks and architectures, using a single NVIDIA RTX A6000 graphics card.

  • C Computational cost: Training costs are reported in hours for physics-informed DeepONet models with different architectures and weighting schemes.
  • C Computational cost: All reported training timings were obtained on a single NVIDIA RTX A6000 graphics card.
  • C Computational cost: Table 11 reports computational costs across the benchmarks and architectures used in the work.

D Proof of Lemma 3.2

The proof establishes that H(θ) is a Gram matrix and bounds its maximum entry using inner-product and norm inequalities.

  • H(θ) is identified as a Gram matrix through vectors v_i defined from the training data and network output.The vectors are defined as v_i = T^(i)(u^(i)(x_i), G_θ(u^(i))(y_i)).
  • The proof invokes the defining property of a Gram matrix to justify the relevant matrix-entry relation.
  • The maximum diagonal entry is bounded above and below through the matrix infinity norm and Cauchy–Schwarz inequality.The displayed chain compares max_k H_kk(θ), ||H(θ)||∞, ⟨v_i,v_j⟩, and ||v_k||².

E Proof of Lemma 3.4

The proof recalls the loss definition, considers the corresponding gradient flow, and introduces the matrix K governing that flow.

  • The proof begins by recalling the definition of the loss function from equation 3.6.
  • It then considers the gradient flow associated with the loss.
  • Using Definitions 3.1 and 3.3, the proof introduces K as an N* × N* matrix with specified entries.

F Antiderivative

The appendix presents training-loss and predictive-error diagnostics for antiderivative and advection-equation DeepONets under different architectures and weighting schemes.

  • Antiderivative: Figure 15 compares weighted training-loss convergence for antiderivative DeepONets over 40,000 Adam iterations.Because all losses are weighted, their magnitudes are not informative.
  • Advection equation: Figure 16 compares training-loss convergence for advection-equation DeepONets across architectures and weighting schemes over 3 × 10^5 Adam iterations.The plotted losses are unweighted.
  • Advection equation: Figure 17 reports averaged test-set relative L2 error for physics-informed advection DeepONets across fixed weights λ_bc = λ_ic = λ in [10^-2, 10^2].
  • Advection equation: Figure 18 shows predicted solutions from the best-trained physics-informed DeepONet for three advection test examples.

H Burger’s equation

The appendix reports training-loss convergence, relative L2 errors, and predicted solutions for physics-informed DeepONets applied to Burgers’ and Stokes equations.

  • Training convergence: Figures 19 and 24 compare training-loss convergence across architectures and weighting schemes for Burgers’ and Stokes equations over 2 × 10^5 Adam iterations.The plotted losses are unweighted.
  • Burgers’ equation: Figure 20 reports averaged test-set relative L2 error for Burgers’ equation across fixed weights λ_bc = λ_ic = λ in [10^-2, 10^2].
  • Burgers’ equation: Figures 21–23 show predicted solutions for three Burgers’ test examples at ν = 0.01, 0.001, and 0.0001.
  • Stokes equation: Figures 25–27 show predicted Stokes solutions for one test example each.
Loading 2110.01654v2…