Research map / Optimization / Optimizers
Improving Energy Natural Gradient Descent through Woodbury, Momentum, and Randomization
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.
Improves energy natural gradient descent (ENGD) for PINNs by applying the Woodbury identity (kernel-space ENGD-W), adding SPRING momentum, and testing a GPU-efficient randomized Nyström sketch on several Poisson, heat, and Fokker-Planck problems.
Problem
Training physics-informed neural networks to high accuracy is hard; natural gradient methods like ENGD are accurate but cost O(P^3) in the number of parameters, and existing scalable alternatives (K-FAC, matrix-free CG) have drawbacks.
Approach
Reformulate ENGD in sample space via the Woodbury/push-through identity so the solve costs O(N^2 P) instead of O(P^3); add SPRING momentum transferred from VMC; and use a GPU-efficient randomized Nyström sketch-and-solve (Cholesky-only, no SVD) for large batches.
Key results
- ENGD-W with Woodbury reaches the same sub-1e-3 L2 error as original ENGD up to 75x faster on 5d Poisson and takes more than 30x more steps in the same time budget.
- SPRING achieves an L2 error an order of magnitude lower than KFAC on the 9+1d log-Fokker-Planck problem.
- GPU-efficient Nyström is 13x to 65x faster than the standard Nyström for sketch sizes of 20% to 80% of N=5000 in the benchmark.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use the Woodbury kernel-space form of ENGD instead of the original parameter-space form | ENGD on 5D Poisson · 5D Poisson | iteration speed: more than 30 times | Helped |
| Use Woodbury ENGD to reach the same sub-10^-3 L2 error as original ENGD | ENGD on 5D Poisson · 5D Poisson | speedup at the same L2 error: up to 75× faster | Helped |
| Add SPRING momentum to ENGD-W | ENGD-W on PINN problems · 5D and 100D Poisson; 4+1D Heat; 9+1D log-Fokker-Planck | final L2 error | Helped |
| Use SPRING rather than ENGD-W | ENGD-W on 9+1D log-Fokker-Planck · 9+1D log-Fokker-Planck | L2 error: an order of magnitude lower | Helped |
| Use GPU-efficient Nyström instead of standard Nyström | Nyström approximation · Nyström benchmark, sketch size 20% of N | runtime (s): 0.1321 vs. 0.0099; 13.29 | Helped |
| Use GPU-efficient Nyström instead of standard Nyström | Nyström approximation · Nyström benchmark, sketch size 40% of N | runtime (s): 0.6246 vs. 0.0253; 24.70 | Helped |
| Use GPU-efficient Nyström instead of standard Nyström | Nyström approximation · Nyström benchmark, sketch size 60% of N | runtime (s): 2.1137 vs. 0.0461; 45.90 | Helped |
| Use GPU-efficient Nyström instead of standard Nyström | Nyström approximation · Nyström benchmark, sketch size 80% of N | runtime (s): 4.7717 vs. 0.0725; 65.41 | Helped |
| Use 10% sketch-size Nyström randomization rather than exact Woodbury computation (ablation) | ENGD-W on 5D Poisson · 5D Poisson, large batches | training speed: 2× | Mixed |
| Use randomized Nyström variants instead of exact SPRING computation (ablation) | SPRING on 100D Poisson · 100D Poisson | performance | Hurt |
| Use a 10% sketch to approximate the regularized kernel matrix (ablation) | Randomized ENGD-W and SPRING · 5D Poisson and 100D Poisson | accuracy: 10% | Hurt |
| Use exact Woodbury computation rather than randomized approximation in late training (ablation) | ENGD-W and SPRING · PINN problems with large batches | accuracy | Helped |
Surprising findings
- Randomized Nyström gives little or no speedup over ENGD-W in high dimensions and for late-stage training, because the regularized kernel's effective dimension stays above 50% of the batch size.
- Once the Woodbury identity is used, randomization loses its advantage, contrary to earlier reports (McKay et al.) that randomization helped the original ENGD.
Limitations
- Randomized Nyström gives diminishing returns in late training and fails to scale to high-dimensional problems.
- Woodbury form still scales quadratically in batch size N, and the authors use a fixed Nyström rank rather than an adaptive one.
- Results rely on extensive per-problem hyperparameter tuning (more than 4500 runs in total).
Experimental setup
Datasets: synthetic PDE collocation points (5d Poisson, 100d Poisson, 4+1d heat, 9+1d log-Fokker-Planck)
Benchmarks: 5d Poisson 100d Poisson 4+1d Heat equation 9+1d log-Fokker-Planck 10d Poisson
Metrics: L2 error wall-clock time to reach target L2 error
Baselines: ENGD SGD Adam Hessian-free optimization K-FAC MinSR standard Nyström approximation
Builds on: energy natural gradient descent MinSR SPRING K-FAC for PINNs Hessian-free optimization
Closest papers by meaning
- Near-optimal Sketchy Natural Gradients for Physics-Informed Neural Networks (ICML 2025 · similarity 0.92)
- Challenges in Training PINNs: A Loss Landscape Perspective (ICML 2024 · similarity 0.89)
- ANaGRAM: A Natural Gradient Relative to Adapted Model for efficient PINNs learning (ICLR 2025 · similarity 0.88)
- Dual Cone Gradient Descent for Training Physics-Informed Neural Networks (NeurIPS 2024 · similarity 0.88)
- Gradient Alignment in Physics-informed Neural Networks: A Second-Order Optimization Perspective (NeurIPS 2025 · similarity 0.88)
- ConFIG: Towards Conflict-free Training of Physics Informed Neural Networks (ICLR 2025 · similarity 0.88)
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.