Research map / Optimization / Optimizers
Layerwise LQR for Geometry-Aware Optimization of Deep 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.
Proposes Layerwise LQR (LLQR), which rewrites divergence-induced steepest-descent steps (Newton, Gauss-Newton, NGD) as finite-horizon LQR and learns structured inverse preconditioners (diagonal, Kronecker, E-KFAC) by minimizing the LQR objective, used as a wrapper around SGDM/AdamW.
Problem
Scalable geometry-aware (Newton/natural-gradient) optimization of deep networks: existing K-FAC/Shampoo-style methods impose block structure on the curvature early, discarding cross-layer couplings.
Approach
Show the dense geometry-aware step equals a layer-coupled LQR problem exactly, then relax by learning structured inverse preconditioners U_i directly under the LQR objective (periodically refit, EMA-smoothed), so block structure is imposed on the learned inverse rather than on the curvature model.
Key results
- ResNet-50 on ImageNet (100 epochs): LLQR+E-KFAC NGD reaches 78.05±0.05 top-1 vs 77.42±0.12 for SGDM baseline at ~1.03x time per epoch.
- ResNet-18 CIFAR-100 with SGDM: LLQR E-KFAC NGD 79.96±0.10 vs 78.97±0.22 base; AdamW 77.07 vs 76.09.
- IWSLT14 De-En with AdamW: LLQR+E-KFAC improves BLEU from 34.24±0.12 to 34.51±0.05 at 1.16x slowdown.
Surprising findings
- Diagonal LLQR structure yields no noticeable convergence acceleration on CIFAR, suggesting it is insufficiently expressive, while Kronecker structures help.
- Relaxed LLQR with diagonal degrees of freedom on Rosenbrock stays better aligned with the Newton direction than diag(H)^-1 preconditioning.
Limitations
- Adds memory and compute overhead from storing and refitting the learned preconditioner, larger for Kronecker blocks.
- Relaxation is not exact: block-structured inverse preconditioners approximate the dense Riccati solution; multiple U solutions cause gap between exact LLQR and relaxed update.
- Diagonal structure gives little convergence benefit on CIFAR.
Experimental setup
Datasets: CIFAR-10 CIFAR-100 ImageNet IWSLT14 De-En algorithmic grokking datasets
Benchmarks: CIFAR-10 CIFAR-100 ImageNet IWSLT14 De-En Rosenbrock function
Metrics: top-1 accuracy BLEU wall-clock time ratio iterations to 95% test accuracy (grokking) cosine similarity to Newton direction
Baselines: SGDM AdamW K-FAC E-KFAC AdaFisher Shampoo Adam
Builds on: K-FAC Newton-LQR connection (stagewise Newton for MLPs) natural gradient descent Riccati recursion
Closest papers by meaning
- On The Concurrence of Layer-wise Preconditioning Methods and Provable Feature Learning (ICML 2025 · similarity 0.86)
- BPQP: A Differentiable Convex Optimization Framework for Efficient End-to-End Learning (NeurIPS 2024 · similarity 0.85)
- Learning Provably Improves the Convergence of Gradient Descent (NeurIPS 2025 · similarity 0.85)
- A Global Geometric Analysis of Maximal Coding Rate Reduction (ICML 2024 · similarity 0.85)
- A Layer-Wise Natural Gradient Optimizer for Training Deep Neural Networks (NeurIPS 2024 · similarity 0.85)
- Structured Preconditioners in Adaptive Optimization: A Unified Analysis (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.