Layerwise LQR for Geometry-Aware Optimization of Deep Networks

TL;DR

The paper introduces Layerwise LQR, a scalable framework for learning structured inverse preconditioners that improve deep network optimization by leveraging second-order geometry without global matrix inversion.

Contribution

It formulates layerwise preconditioning as a finite-horizon LQR problem, enabling scalable learning of structured preconditioners that enhance optimization in deep networks.

Findings

LLQR improves optimization dynamics in ResNets and Transformers.

It often leads to better final test performance.

The method adds modest computational overhead.

Abstract

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.