Research map / Optimization / Optimizers
Structured Inverse-Free Natural Gradient Descent: Memory-Efficient & Numerically-Stable KFAC
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.
Proposes SINGD, an inverse-free, structured-Kronecker-factor natural gradient method that generalizes INGD, recovers KFAC-like updates (IKFAC), and trains stably in half precision with lower memory.
Problem
KFAC needs dense Kronecker factors and matrix inversions, making it memory-hungry and numerically unstable in low (half) precision, which limits its use in modern mixed-precision training of transformers and CNNs.
Approach
Replace matrix inversion with subtraction in a matrix-logarithm space using a truncated matrix exponential (inverse-free), show this recovers KFAC to first order (IKFAC), and impose sparse structures (diagonal, block-diagonal, hierarchical, Toeplitz, low-rank triangular) on the Kronecker factors via a unified projection that preserves closure under the update.
Key results
- In BFP-16, SINGD-Diag reaches 2.67 GiB peak memory and 23.8 min training time versus KFAC in FP-32 at 4.00 GiB and 83.2 min on the VGG CIFAR-100 setup.
- SINGD (including IKFAC and INGD as special cases) outperforms AdamW in many mixed-precision cases across transformers, CNNs and GNNs, while KFAC performs unstably in half precision.
- Diagonal Kronecker factors can match dense (INGD) performance on several models while lowering memory to roughly AdamW's level.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use diagonal instead of dense Kronecker factors (ablation) | SINGD / INGD · ImageWoof-10 (GCViT) | test error | No clear effect |
| Use hierarchical instead of dense Kronecker factors (ablation) | SINGD / INGD · ImageWoof-10 and CIFAR-100 models | test error | No clear effect |
| Use hierarchical instead of block-diagonal Kronecker factors (ablation) | SINGD · ImageWoof-10 and CIFAR-100 models | test error | Helped |
| Use diagonal Kronecker factors instead of dense factors (ablation) | SINGD / INGD · ImageWoof-10 (GCViT) | peak memory: 16.0 -> 15.8 | No clear effect |
| Use diagonal Kronecker factors instead of dense factors (ablation) | SINGD / INGD · ImageWoof-10 (GCViT) | training time: 203 -> 200 | Hurt |
| Use structured diagonal factors instead of dense factors (ablation) | SINGD / INGD · ImageWoof-10 (GCViT) | test error | Helped |
| Remove Riemannian momentum (ablation) | IKFAC · CIFAR-100 (VGG) | peak memory | Helped |
| Replace KFAC matrix inversions with inverse-free updates | IKFAC · BFP-16 training | numerical stability | Helped |
| Use SINGD instead of AdamW | SINGD · ImageWoof-10, CIFAR-100, and Cora | test error | Helped |
| Use SINGD instead of AdamW | SINGD · ImageNet-100 (GCViT) | test error | Helped |
| Use structured Kronecker factors (ablation) | SINGD · ImageWoof-10 (GCViT) | peak memory: 15.8 (1.02\,x) vs 15.7 (1.00\,x) | No clear effect |
Surprising findings
- Diagonal Kronecker factors perform as well as dense ones on several models, despite far fewer parameters.
- KFAC is unstable in BFP-16 even though it is the standard second-order baseline.
Limitations
- KFAC cannot be run natively in half precision, so its comparison requires FP-32 for matrix inverses.
- The KFAC equivalence (IKFAC) holds only up to first-order accuracy in the step size.
Experimental setup
Datasets: CIFAR-100 ImageWoof-10 Cora ImageNet-100
Benchmarks: CIFAR-100 ImageWoof-10 Cora ImageNet-100
Metrics: test error peak memory (GiB) training time (min)
Baselines: KFAC AdamW SGD INGD IKFAC
Builds on: KFAC INGD Bayesian learning rule Riemannian momentum
Closest papers by meaning
- Fast Convergence of Natural Gradient Descent for Over-Parameterized Neural Networks (NeurIPS 2019 · similarity 0.88)
- Fast Approximate Natural Gradient Descent in a Kronecker Factored Eigenbasis (NeurIPS 2018 · similarity 0.87)
- A Layer-Wise Natural Gradient Optimizer for Training Deep Neural Networks (NeurIPS 2024 · similarity 0.86)
- Regularized Adaptive Momentum Dual Averaging with an Efficient Inexact Subproblem Solver for Training Structured Neural Network (NeurIPS 2024 · similarity 0.85)
- Provable Acceleration of Nesterov's Accelerated Gradient for Asymmetric Matrix Factorization and Linear Neural Networks (NeurIPS 2024 · similarity 0.85)
- On The Concurrence of Layer-wise Preconditioning Methods and Provable Feature Learning (ICML 2025 · 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.