PaperlayerView exact revision on arXiv

Source-linked AI summary

Tensor neural networks for high-dimensional Fokker-Planck equations

Taorui Wang, Zheyuan Hu, Kenji Kawaguchi, Zhongqiang Zhang, George Em Karniadakis

arXiv:2404.05615v2math.NAcs.NA

TL;DR

High-dimensional Fokker–Planck equations pose approximation challenges on unbounded domains. This paper applies tensor neural networks with selected numerical supports and reports feasible, accurate approximations, including relative errors below 5% in different density regions.

  • Problem

    Fokker–Planck equations involve approximation on unbounded domains and in high dimensions, where grid-based numerical methods are challenging.

    Source: §1 · p. 1

  • Method

    The paper investigates tensor radial-basis-function and tensor-feedforward networks, selecting numerical supports and constraining radial-basis-function parameters.

    Source: §6 · p. 26

  • Results

    Examples and ablation studies illustrate the feasibility and accuracy of both tensor networks, with relative errors below 5% on different density regions in one reported result.

    Source: §4 · p. 21, §6 · p. 26

  • Takeaways & Limitations

    Radial basis functions provide exact derivatives and integration over the numerical support, significantly reducing storage requirements.

    Source: §6 · p. 26

    Efficiency is sensitive to the choice of numerical support.

    Source: §6 · p. 26

Abstract

from arXiv · show

We solve high-dimensional steady-state Fokker-Planck equations on the whole space by applying tensor neural networks. The tensor networks are a linear combination of tensor products of one-dimensional feedforward networks or a linear combination of several selected radial basis functions. The use of tensor feedforward networks allows us to efficiently exploit auto-differentiation (in physical variables) in major Python packages while using radial basis functions can fully avoid auto-differentiation, which is rather expensive in high dimensions. We then use the physics-informed neural networks and stochastic gradient descent methods to learn the tensor networks. One essential step is to determine a proper bounded domain or numerical support for the Fokker-Planck equation. To better train the tensor radial basis function networks, we impose some constraints on parameters, which lead to relatively high accuracy. We demonstrate numerically that the tensor neural networks in physics-informed machine learning are efficient for steady-state Fokker-Planck equations from two to ten dimensions.

1. Introduction

The introduction motivates tensor neural networks for high-dimensional Fokker–Planck equations on unbounded domains, outlining numerical-support selection, tensor-network design, and approximation in dimensions 2–10. It also previews ablation studies and a concluding discussion of the work’s limitations and future directions.

1. Introduction

  • High-dimensional benchmarks often use simple targets, leaving approximation challenges unaddressed for functions with complicated interactions among coordinates and multiple modes.The introduction also notes that meaningful solution values can be small in much of the domain, making them difficult to resolve.Source: §1 · p. 2, §1.1 · p. 3

1.1. Our proposed methodology and contributions

  • The work develops tensor radial-basis and tensor feedforward networks for steady-state Fokker–Planck equations in 2–10 dimensions.The authors report experiments across several dimensions and ablation studies.Source: §1.1 · pp. 3–4
  • Tensor radial-basis networks avoid auto-differentiation for derivatives, while effective initialization and parameter constraints address their more demanding training.Their tensor structure makes spatial and parameter derivatives simple to compute, reducing storage required by auto-differentiation packages.Source: §1.1 · p. 3
  • Training results in high dimensions are sensitive to numerical support, and the authors introduce an approximation-error metric focused on high-probability regions.They note that smaller supports may improve efficiency and accuracy in ten dimensions.Source: §1.1 · p. 3
  • An algorithm estimates a bounded numerical support from a small amount of SDE data offline, with no such data used during later training.The selection uses a few simulated trajectories and is reported to address support selection in six dimensions.Source: §1.1 · pp. 2–4
  • Tensor radial-basis networks guarantee density solutions, while feedforward solutions require accurate integration and may incur large errors or fail in high dimensions.For feedforward networks, the reported solutions are close to density solutions when integration is sufficiently close to 1.Source: §1.1 · p. 4
  • Section 5 presents ablation studies, while Section 6 summarizes the work and discusses its limitations and possible future directions.Source: §1.1 · p. 4

2. Literature review

The review surveys general approaches to high-dimensional PDEs and methods specific to Fokker–Planck equations, focusing on issues relevant to this work rather than aiming for a complete list. It discusses neural-network, derivative, kernel, and data-based strategies alongside their stated scope limits.

2.1. General issues in high dimensions

  • DeepBSDE and related stochastic-representation approaches approximate solutions at one or a few points and may require extensive computation to cover a large region.The review distinguishes local approximations from solutions over a large region.Source: §2.1 · p. 4
  • Randomized calculus replaces approximate solutions and derivatives with stochastic approximations using Stein’s identity, avoiding auto-differentiation but adding approximation error and requiring no more than four dimensions for nontrivial solutions.The cited dimensional limit is reported for the method in.Source: §2.1 · p. 5
  • Earlier tensor-neural-network PDE studies considered problems in up to four dimensions, while this work applies tensor neural networks to higher-dimensional problems.The review mentions tensor neural networks for Schrödinger equations [13, 14], separable PINNs [12], and low-rank networks.Source: §2.1 · p. 5
  • Kernel networks, including additive kernels and mixtures of kernels, have been applied to accommodate complex underlying problem structures and increase approximation capacity; multiple-kernel learning algorithms for large-scale problems reduce computational complexity and memory usage.Source: §2.1 · p. 5

2.2. Solving Fokker-Planck equations with neural networks

  • For Fokker–Planck equations, prior work uses PINN formulations [9] and can incorporate sampled SDE solutions into training and data losses.For time-dependent equations, training at initial and terminal times is reported to improve efficiency.Source: §2.2 · p. 6
  • Other Fokker–Planck strategies include variational losses and transforming high-dimensional equations into finite-dimensional ODE systems.The variational formulation is described in relation to.Source: §2.2 · p. 6
  • Reviewed directions include computing a generalized potential from data—the logarithm of the invariant distribution for time-independent Fokker–Planck equations—applying a logarithmic solution transformation [56], and estimating SDE drift and diffusion before solving.Source: §2.2 · p. 6
  • RBF methods have been applied to Fokker–Planck equations, but the cited RBF neural-network work addresses only low-dimensional problems.The review mentions Kansa’s approach, Hermite interpolation, and Galerkin methods, as well as a Gaussian RBF network with normalization.Source: §2.2 · p. 6

3. Problem and methodology

This section presents a methodology for approximating steady-state Fokker–Planck densities with tensor networks trained on an estimated bounded support. It covers support selection, loss terms and parameter constraints, and training choices including model-capacity tradeoffs.

3. Problem and methodology

  • SDE trajectories are used to select the computational support, after which a physics-informed neural network approximates the density; no SDE simulation data are used for training.The paper describes the method as not data-driven, apart from data used to find the computational domain.Source: §3 · p. 7

3.1. Estimating Bounded Domain

  • A bounded hypercube is estimated from simulated SDE trajectories, then refined using a trained density approximation so that its integrated mass exceeds a chosen threshold.The initial simulation uses an Euler–Maruyama trajectory with a burn-in period; the refinement threshold is typically 0.95–1.0.Source: §3.1 · pp. 7–8

3.2. Tensor neural network (TNN)

  • The tensor network forms sums of products of nonnegative one-dimensional networks, using either fully connected feedforward components or additive radial-basis components.For the radial-basis version, the normalization constant is computed analytically from the selected basis functions.Source: §3.2 · pp. 9–11
  • Radial-basis shifts are constrained in relation to the hypercube half-edge length and penalized during training, helping concentrate the basis functions inside the bounded domain.The radial-basis network has 3Ndm parameters, where N is rank, d is dimension, and m is the number of basis functions.Source: §3.2 · p. 11

3.3. Loss function

  • The loss combines the Fokker–Planck residual with unit-mass normalization; nonnegativity is built into activations or kernels, while TFFN and TRBFN enforce boundary values by multiplication and penalization, respectively.Source: §3.3 · p. 11
  • The loss penalizes boundary-condition violations and parameter constraints from Equation (6) through separate terms in (9).Source: §3.3 · p. 11

3.4. Training settings

  • Both networks use LION for 10^5 epochs; TFFN uses float64 for better training results, whereas TRBFN uses float32 because the precisions perform the same.Source: §3.4 · p. 12
  • In Example 4.1, increasing N with small m (m < 10) yielded smaller losses and more accurate solutions than increasing m with small N (N < 10).Because the TRBFN parameter count is of order 3Ndm, the authors choose large N and small m for Section 4.Source: §3.4 · p. 12
  • Implementation requires balancing rank N against one-dimensional network size or basis-function count because prevalent Python packages use tensor implementations.Source: §3.4 · p. 12

4. Numerical results

The numerical section evaluates TRBFNs on five steady-state Fokker–Planck examples and TFFNs on the first three, using a specific drift to compare with exact solutions. It examines how support choice relates to numerical accuracy, including results for six- and ten-dimensional tests.

Example 4.1 (2D, ring potential)

  • At 10^5 epochs, TRBFN(1000, 3) reaches loss 10^-4 versus 10^-2 for TRBFN(5, 720) in the 2D ring-potential example.The lower loss suggests a better approximation.Source: Figure 2 · p. 14, §4 · p. 14
  • TRBFN(1000, 3) achieves at least one order of magnitude higher accuracy than TRBFN(5, 720), indicating greater approximation capability at larger rank.Source: §4 · p. 14, Table 3 · p. 15
  • TFFN(1024, [1 8 8 8 1]) matches the accuracy magnitude of TRBFN(1000, 3), with relative errors no larger than 0.05% on Γ10^-2.In the same example, larger-rank TFFNs improve accuracy, while rank-5 TFFN accuracy is two orders of magnitude lower than rank-5 TRBFN accuracy.Source: §4 · p. 14, Table 3 · p. 15

Example 4.2 (4D, solution with a single mode)

  • On the initial 4D support, TFFN(128, [1 8 8 1]) and TFFN(64, [1 8 8 1]) have relative errors greater than 10% on Γ10^-2.Source: §4 · p. 16
  • In the 4D single-mode example, TRBFN relative average errors are below 1% across the tested regions except for TRBFN(64, 27) on Γ10^-2; larger rank and fewer basis functions typically reduce error.Source: §4 · p. 15, Table 4 · p. 16
  • After refining the 4D support with ϑ = 0.999, both tested TFFNs achieve relative errors below 7% across Γ10^-2, Γ5×10^-2, and Γ10^-1.Source: Table 4 · p. 16, §4 · p. 16

Example 4.3 (6D, solution with a single mode)

  • In the 6D single-mode example, TRBFNs have smaller relative errors than TFFNs, while increasing TFFN rank and network depth or width reduces errors.Both network types capture key features of the exact solution.Source: Table 5 · p. 17, §4 · p. 17

Example 4.4 (6D, solution with multiple modes)

  • The 6D multiple-mode example tests a density with four separated modes arising from three-way interactions among (x1, x2, x3) and (x4, x5, x6).The training batch size is 40000.Source: §4 · pp. 17–18

Example 4.4

  • In Example 4.4, initial-support relative errors were under 30% across density regions, and the numerical solution captured the exact solution’s modes.The loss oscillated over epochs during TRBFN(800, 6) training.Source: §4 · p. 18, Table 6 · p. 19, Figure 6 · p. 19
  • In Example 4.4, refining support with threshold ϑ = 0.97 yielded errors below 5% in Γ10^-3 and below 10% in Γ2×10^-4.The refinement used Algorithm 2 on the trained network; Γ = [-2, 2]6.Source: §4 · p. 19, Table 7 · p. 20

Example 4.5 (10D, solution with two modes)

  • For bounded domain N10, relative error is over 99.99%; to obtain a smaller bounded domain Ω, anisotropic half-edge lengths were chosen using simulated SDE trajectories.Source: §4 · p. 20
  • On the anisotropic support in Example 4.5, density-region errors were around 75%, although projected contours showed that TRBFN captured the accurate solution’s modes.The trained network’s integrals were also evaluated over supports with half-edge lengths from 0.3 to 2.4.Source: §4 · p. 20, Figure 7 · p. 21
  • In Example 4.5, errors were below 50% on one tested support and below 5% on another, demonstrating 10D capability but leaving support selection unresolved.The authors state that finding a smaller domain containing most probability mass remains future work.Source: Table 10 · p. 21, §4 · p. 21

5. Comparison and ablation studies

The section compares tensor neural networks with a data-driven method and tests how optimization, boundary enforcement, radial basis choices, and arithmetic precision affect performance.

5.1. Comparison with a data-driven method

  • Both tensor neural networks perform similarly to the data-driven method, while avoiding SDE simulations during training.The comparison is for Example 4.3.Source: §5.1 · p. 22

5.2. Using different optimization methods

  • LION minimizes TRBFN loss most effectively, suggesting it may achieve the best accuracy among the tested optimizers.The loss comparison is shown in Figure 8; Table 13 reports average relative errors.Source: §5.2 · p. 23, Table 13 · p. 23

5.3. Enforcing boundary condition or not

  • TFFN is highly sensitive to boundary conditions, while TRBFN works without an explicit boundary condition and can improve when one is imposed.The TFFN test reports 99.99% relative error; Table 14 reports the TRBFN boundary-penalty comparison.Source: §5.3 · pp. 23–24, Table 14 · p. 24

5.4. Using different RBFs in TRBFN

  • TRBFN using only compactly supported radial basis functions achieves higher accuracy than the tested mixture of compactly supported and positive functions.The authors attribute the difference to the compactly supported functions being zero outside their support; Table 15 reports the comparison.Source: §5.4 · p. 25, Table 15 · p. 25

5.5. Using different arithmetic precision

  • TRBFN achieves the same relative error in single and double precision, whereas TFFN has smaller relative error in double precision.The authors suggest numerical integration error in computing Z(Θ) may explain the TFFN difference; TRBFN computes Z(Θ) analytically.Source: Table 16 · p. 25, §5.5 · p. 26

6. Summary and discussion

The section summarizes tensor radial-basis and tensor-feedforward networks for high-dimensional steady-state Fokker-Planck equations, including numerical-support selection and constraints on radial-basis parameters. It reports feasibility and accuracy across examples and identifies methodological limitations and future work.

6. Summary and discussion

  • The authors designed a procedure to estimate numerical supports and applied both network types to steady-state Fokker-Planck equations.Source: §6 · p. 26
  • Constraints on radial-basis parameters improve approximations, while radial-basis functions provide exact derivatives and support integration with reduced storage versus auto-differentiation.Source: §6 · p. 26
  • Tensor-feedforward networks are straightforward to implement but more sensitive to numerical-support choice, even in six or fewer dimensions.Source: §6 · p. 26
  • Examples and ablation studies illustrate the feasibility and accuracy of both tensor-network approaches.Source: §6 · p. 26
  • Hypercube supports can be large for slowly decaying or separated densities, and larger supports may require more training points.Source: §6 · p. 26
  • Kernel choice is somewhat arbitrary, and the general constraints imposed on radial-basis parameters remain problem-dependent.Source: §6 · p. 26
  • Approximation-error estimates and efficiency require theoretical study; combining the approach with DeepBSDE is suggested for accuracy at selected high-probability points.Source: §6 · p. 26

Declaration of competing interest

The authors declare no known competing financial interests or personal relationships that could have influenced the reported work.

Declaration of competing interest

Appendix A. Universal approximation of tensor neural networks

Appendix A presents an argument for universal approximation by tensor neural networks. Appendix B derives stability and generalization-error results for the least-square formulation under stated assumptions.

Appendix A. Universal approximation of tensor neural networks

  • Appendix A uses integral approximation over a bounded domain and numerical integration to support a universal-approximation argument for tensor-product kernels.Source: Appendix A · p. 33

Appendix B. Stability of the least-square formulation

  • Appendix B establishes stability for the least-square formulation under assumptions including a unique solution and positive-definite matrix D.Source: Appendix B · p. 34
  • The appendix gives a generalization-error estimate for network solutions, with terms tied to the trained loss and Monte Carlo integration error.Source: Appendix B · pp. 34–35

Appendix C. Hyper-parameters in Section 3.1 for numerical examples

The appendix reports the hyperparameters used in the numerical examples.

Appendix C. Hyper-parameters in Section 3.1 for numerical examples

Appendix C.1. Hyperparameters in Algorithm 2

Appendix C.1 reports Algorithm 2 hyperparameters and integration results for trained TRBFNs in Examples 4.2–4.4.

Appendix C.1. Hyperparameters in Algorithm 2

  • The appendix reports Algorithm 2 hyperparameters and describes an integration mean involving the trained network pN.The integration expression uses intervals centered at O with half-edge length r.Source: Appendix C.1 · p. 35
  • Table C.18 lists elements in Sn and their integration results with the trained TRBFN for Example 4.2.Source: Table C.18 · p. 35
  • Table C.19 lists elements in Sn and their integration results with the trained TRBFN for Example 4.3.Source: Table C.19 · p. 35
  • Table C.20 lists elements in Sn and their integration results with the trained TRBFN for Example 4.4.Source: Table C.20 · p. 35

Appendix C.2. Weights for TRBFN

The appendix reports penalty-term weights used in the TRBFN loss function, including a separate stated setting for Example 4.4.

Appendix C.2. Weights for TRBFN

Appendix D. Training Hyper-parameters

Appendix D lists the training batch sizes for TRBFNs and TFNNs across the Section 4 examples.

Appendix D. Training Hyper-parameters

Loading 2404.05615v2…