# Layerwise LQR for Geometry-Aware Optimization of Deep Networks

arXiv: https://arxiv.org/abs/2605.04230 · HTML: https://mlsanity.com/paper/2605.04230/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 48 references that are in the map.
- First posted on arXiv 5 May 2026.

## Builds on (most cited first, 23 of 23)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Improved Regularization of Convolutional Neural Networks with Cutout](https://mlsanity.com/paper/1708.04552/) (2017)
- [Automatic differentiation in machine learning: a survey](https://mlsanity.com/paper/1502.05767/) (2015)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets](https://mlsanity.com/paper/2201.02177/) (2022)
- [Deep Equilibrium Models](https://mlsanity.com/paper/1909.01377/) (2019)
- [New insights and perspectives on the natural gradient method](https://mlsanity.com/paper/1412.1193/) (2014)
- [Scaling Neural Machine Translation](https://mlsanity.com/paper/1806.00187/) (2018)
- [Old Optimizer, New Norm: An Anthology](https://mlsanity.com/paper/2409.20325/) (2024)
- [Revisiting Natural Gradient for Deep Networks](https://mlsanity.com/paper/1301.3584/) (2013)
- [Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis](https://mlsanity.com/paper/1806.03884/) (2018)
- [Maximum Principle Based Algorithms for Deep Learning](https://mlsanity.com/paper/1710.09513/) (2017)
- [An Optimal Control Approach to Deep Learning and Applications to Discrete-Weight Neural Networks](https://mlsanity.com/paper/1803.01299/) (2018)
- [PipeFisher: Efficient Training of Large Language Models Using Pipelining and Fisher Information Matrices](https://mlsanity.com/paper/2211.14133/) (2022)
- [Fast and Furious Convergence: Stochastic Second Order Methods under Interpolation](https://mlsanity.com/paper/1910.04920/) (2019)
- [Optimizers Qualitatively Alter Solutions And We Should Leverage This](https://mlsanity.com/paper/2507.12224/) (2025)
- [AdaFisher: Adaptive Second Order Optimization via Fisher Information](https://mlsanity.com/paper/2405.16397/) (2024)
- [Structured Inverse-Free Natural Gradient: Memory-Efficient & Numerically-Stable KFAC](https://mlsanity.com/paper/2312.05705/) (2023)
- [No More Adam: Learning Rate Scaling at Initialization is All You Need](https://mlsanity.com/paper/2412.11768/) (2024)
- [Scalable K-FAC Training for Deep Neural Networks with Distributed Preconditioning](https://mlsanity.com/paper/2206.15143/) (2022)

## Built on by (most cited first, 1 of 1)

- [Navigating Potholes with Geometry-Aware Sharpness Minimization](https://mlsanity.com/paper/2605.16134/) (2026)

## Abstract (arXiv)

Geometry-aware optimizers such as Newton and natural gradient can improve conditioning in deep learning, but scalable variants such as K-FAC, Shampoo, and related preconditioners usually impose structural approximations early, often discarding cross-layer interactions induced by the network computation. We introduce Layerwise LQR (LLQR), a framework for learning structured inverse preconditioners under a global layerwise optimal-control objective. The starting point is an exact equivalence: the steepest-descent step under a broad class of divergence-induced quadratic models--including Newton, Gauss-Newton, Fisher/natural-gradient, and intermediate-layer metrics--can be written as a finite-horizon Linear Quadratic Regulator (LQR) problem. This formulation serves as a reference that exposes the layerwise dynamics and cost matrices encoding the original dense geometry. We then derive a scalable relaxation that learns diagonal, (E-)Kronecker-factored, or other structured inverse preconditioners by minimizing the LQR objective and reusing them across iterations. The resulting optimizer wraps standard methods while retaining a principled connection to second-order geometry, without forming or inverting the global curvature matrix. Experiments on ResNets and Transformers show that LLQR improves optimization dynamics and often translates these gains into improved final test performance, while adding only modest wall-clock overhead. It establishes LLQR as a practical framework for geometry-aware second-order methods and a reference for evaluating scalable approximations.
