/WAN_PDE

Weak Adversarial Networks for high-dimensional PDE with Neural ODEs

Primary LanguageJupyter NotebookMIT LicenseMIT

Weak Adversarial Networks with Neural ODEs

  • Implemented Weak Adversarial Nets for solving high-dimenstional PDE.
  • Formulated time dependent solution as an ODE problem, instead of numerical integration. Backpropagation by adjoint method of neural ODEs.

Results:

  1. Both Time dependent and independent PDEs converged well.
  2. Re-formulating as ODE problem serves a better approximation than numerical approximation methods. The typical method of backpropagation through ODEs is adjoint method as mentioned in Neural Ordinary Differential Equations paper (NIPS 2018) by Chen et al.
  3. Faster and better convergence, but at the cost of greater memory consumption during backpropagation by Tracker.jl.

Advantages:

  • Faster Convergence of time dependent PDEs leading upto one-third number of iterations required for converging.
  • We can select from a wide variety of ODE solvers based on order of error tolerance vs speed to converge trade-off.

Elliptic Dirichlet (20 dimensions)

Note that graph is shown about x1 and x2 axis only while keeping remaining x_i = 1 for i=3, 4, ...20

Approximated Function:

Approx Function

True Function:

True Function

Time Dependent PDEs

Nonlinear diffusion-reaction of the form

ut − ∆u − u2 = f(x, t), in Ω × [0, T]

u(x, t) = g(x, t), on ∂Ω × [0, T]

u(x, 0) = h(x), in Ω

Where Ω = (−1, 1)d ⊂ Rd where d=5 in this case. T (time) spans [0,1]

Approximated Function:

Approximate Function

True Function:

True Functions

Absolute Difference b/w approximated function and true function:

Absolute Difference

Possible directions:

Current limitation of this approach of solving PDE consumes about 2.5x more memory while backpropagating via adjoint method when compared to using reverse mode AD with numerical integration methods.

Possible Solution to current limitation switch autodiff backend from Tracker.jl to Zygote.jl. This can be done when DiffEqFlux.jl and Flux.jl switches over to Zygote.jl. Another possible direction would be to see the performance of Reverse Mode AD of Zygote.jl vs Adjoint method.

Zygote.jl uses source code transformation for reverse mode AutoDiff. In general, this leads to faster and much less memory consumption while backpropagating when compared to Tracker.jl's Auto-diff.

References

  1. Yaohua Zang, Gang Bao, Xiaojing Ye, and Haomin Zhou. Weak Adversarial Networks for High-dimensional Partial Differential Equations. Journal of Computational Physics 2020, Article 109409. [arxiv]
  2. Ricky T. Q. Chen, Yulia Rubanova* Jesse Bettencourt, and David Duvenaud. Neural Ordinary Differential Equations. Advances in Neural Information Processing Systems 2018, pages 6571 - 6583. [arxiv]
  3. Lev Semenovich Pontryagin, EF Mishchenko, VG Boltyanskii, and RV Gamkrelidze. The mathematical theory of optimal processes. 1962.
  4. Mike Innes. Flux: Elegant machine learning with Julia. Journal of Open Source Software, 3(25), 602. [paper]
  5. Chris Rackauckas, Mike Innes, Yingbo Ma, Jesse Bettencourt, Lyndon White, and Vaibhav Dixit. DiffEqFlux.jl - A Julia Library for Neural Differential Equations. arXiv preprint arXiv:1902.02376. [paper]