Research map / Architectures and efficiency / RNNs and normalization
Kronecker-Factored Approximate Curvature for Physics-Informed Neural Networks
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Extends KFAC to Gauss-Newton matrices of PINN losses with differential operators by viewing Taylor-mode derivatives as weight-sharing layers, enabling scalable natural-gradient-style optimization of PINNs.
Problem
Physics-informed neural networks are hard to train with first-order methods, and existing second-order (ENGD/Gauss-Newton) preconditioners need cubic-cost solves in the full parameter space, limiting them to networks with a few thousand parameters.
Approach
Interpret Taylor-mode autodiff of differential operators as a forward network with shared weights, then apply KFAC-expand for linear weight-sharing layers to Gauss-Newton matrices of PINN interior losses, yielding Kronecker-factored approximations that include differential-operator contributions.
Key results
- KFAC-based optimizers are competitive with full ENGD on small problems while ENGD runs out of memory beyond roughly 10,000 parameters.
- KFAC scales to networks with up to about 1.3 million parameters (100d Poisson) and consistently outperforms SGD, Adam and LBFGS in the reported experiments.
- KFAC achieves performance comparable to Hessian-free on high-dimensional Poisson and Fokker-Planck problems.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use KFAC instead of first-order and quasi-Newton optimizers | PINN training · 2d Poisson | relative L2 error | Helped |
| use KFAC for larger networks beyond ENGD's memory limit | PINN training · 2d Poisson | relative L2 error | Mixed |
| use KFAC rather than SGD, Adam, or LBFGS | PINN training · (4+1)d heat equation | relative L2 error | Helped |
| use KFAC rather than SGD, Adam, or LBFGS | PINN training · 5d, 10d, and 100d Poisson equations | relative L2 error | Helped |
| use KFAC rather than SGD, Adam, or LBFGS | PINN training · (9+1)d Fokker-Planck equation | relative L2 error | Helped |
| resample KFAC's training batch less frequently | KFAC · (4+1)d heat equation | relative L2 error | Helped |
| tune other methods' batch sizes, sampling frequencies, and hyperparameters | other optimizers · (4+1)d heat equation, 116864-parameter MLP | relative L2 error | No clear effect |
Surprising findings
- KFAC improves significantly when batches are sampled less frequently, suggesting it benefits from multiple iterations on a fixed batch.
Limitations
- Implementation currently supports only MLPs and the Poisson and heat equations.
- Adaptive damping and other KFAC heuristics were not fully explored, so performance may be improvable.
- Uses fixed-batch reuse and sparse batch resampling choices that differ between optimizers.
Experimental setup
Datasets: Poisson equation collocation points (synthetic, manufactured solutions) heat equation collocation points (synthetic) Fokker-Planck collocation points (synthetic)
Benchmarks: 2d Poisson equation (4+1)d heat equation 5d Poisson equation 10d Poisson equation 100d Poisson equation (9+1)d Fokker-Planck equation
Metrics: relative L2 error wall-clock time to reach error
Baselines: ENGD Hessian-free SGD Adam LBFGS
Builds on: KFAC ENGD KFAC-expand Taylor-mode automatic differentiation forward Laplacian
Closest papers by meaning
- Kronecker-factored Curvature Approximations for Recurrent Neural Networks (ICLR 2018 · similarity 0.92)
- Global curvature for second-order optimization of neural networks (ICML 2025 · similarity 0.87)
- The Challenges of the Nonlinear Regime for Physics-Informed Neural Networks (NeurIPS 2024 · similarity 0.86)
- PRDP: Progressively Refined Differentiable Physics (ICLR 2025 · similarity 0.85)
- Near-optimal Sketchy Natural Gradients for Physics-Informed Neural Networks (ICML 2025 · similarity 0.85)
- Deconstructing the Goldilocks Zone of Neural Network Initialization (ICML 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.