Layerwise LQR for Geometry-Aware Optimization of Deep Networks

This paper introduces Layerwise LQR (LLQR), a scalable optimization framework that reformulates second-order geometry-aware updates as a linear quadratic regulator problem to learn structured preconditioners that preserve cross-layer interactions without inverting the global curvature matrix, thereby improving training dynamics and final performance in deep networks.

The Gist

Imagine you are trying to navigate a massive, foggy mountain range to find the lowest valley (the best solution for an AI). This is what training a deep neural network is like.

Most standard methods, like Gradient Descent, are like a hiker who only looks at the slope directly under their feet. They take a step downhill based on how steep the ground is right there. It works, but if the valley is shaped like a long, narrow canyon (a common problem in AI), the hiker zig-zags back and forth, taking a very long time to reach the bottom.

Newton's Method is like a hiker with a perfect 3D map. They can see the entire shape of the canyon and take a direct, perfect step to the bottom. However, calculating that perfect map for a giant AI is so computationally expensive that it's impossible to do in real-time. It's like trying to draw a map of the entire world while you are still walking.

Other methods try to compromise by using a "rough sketch" of the map (approximations), but they often throw away important details about how different parts of the mountain connect to each other.

The Paper's Big Idea: "Layerwise LQR" (LLQR)

The authors of this paper propose a new way to navigate: Layerwise LQR. They use a clever trick from the world of optimal control (the math used to guide rockets and robots) to solve this problem.

Here is the analogy:

1. The "Rocket" Analogy (The LQR Connection)

Think of the neural network not just as a static map, but as a rocket flying through space.

• The Layers: Each layer of the network is a stage in the rocket's flight.
• The Goal: We want to steer the rocket (the AI) from its current position to the target (the best solution) with the least amount of fuel (error).
• The Physics: The paper shows that the math used to find the perfect "steering step" for a rocket is exactly the same as the math used to find the perfect "learning step" for an AI.

In rocket science, this is called a Linear Quadratic Regulator (LQR). It's a way to calculate the perfect path by looking at how the rocket moves forward (dynamics) and the cost of deviating from the path (loss).

2. The Problem with the "Perfect" Rocket

If you try to calculate the perfect path for a giant rocket (a huge AI) all at once, the math becomes too heavy. You need to know how every single part of the rocket affects every other part simultaneously. This is the "dense matrix" problem that makes Newton's method too slow.

3. The LLQR Solution: "Learning the Steering Wheel"

Instead of calculating the perfect path every single second, the authors suggest a smarter approach:

• Step 1: They set up the "perfect rocket physics" (the LQR problem) to understand exactly how the layers of the AI are connected. This captures the complex, 3D shape of the canyon that simple methods miss.
• Step 2: Instead of solving the whole rocket equation every time, they learn a "steering wheel" (a preconditioner). This steering wheel is a simplified tool that knows how to turn the rocket in the right direction based on the complex physics they just studied.
• Step 3: They train this steering wheel to be as good as possible at mimicking the perfect path, but they keep it simple (structured) so it's fast to use.

The Key Innovation:
Most other methods try to simplify the map before they start navigating. This paper says: "Let's first understand the full, complex physics of the mountain, and then build a simple, fast steering tool that respects those connections."

What They Found (The Results)

The authors tested this new "steering wheel" on standard AI tasks, like recognizing images (ResNets) and translating languages (Transformers).

• Faster Convergence: The AI learned faster. It didn't zig-zag as much in the "canyons."
• Better Final Score: Because it navigated more efficiently, it often ended up in a better spot (higher accuracy) than standard methods.
• Low Cost: The "steering wheel" didn't require a massive amount of extra computing power. It added only a small amount of time (about 3% slower on large datasets) but gave significant performance boosts.
• Grokking: In a specific phenomenon called "grokking" (where an AI suddenly understands a pattern after a long period of confusion), this method helped the AI "wake up" and learn much faster.

Summary

The paper introduces LLQR, a method that treats training an AI like guiding a rocket. Instead of guessing the path or using a rough sketch, it uses advanced control theory to understand the full complexity of the AI's structure, then builds a lightweight, smart "steering tool" that uses that understanding to guide the AI to the solution much faster and more accurately than before. It bridges the gap between the "perfect but slow" math and the "fast but dumb" math we usually use.

Technical Summary

Technical Summary: Layerwise LQR for Geometry-Aware Optimization of Deep Networks

1. Problem Statement

Geometry-aware optimizers, such as Newton's method and Natural Gradient Descent (NGD), offer superior conditioning and convergence properties by utilizing second-order information (e.g., Hessian or Fisher Information matrices). However, these methods are computationally prohibitive for large-scale deep learning because the curvature matrices are dense and couple parameters across all layers via the chain rule. Directly solving the update equation $H\Delta\theta = -g$ is infeasible.

Existing scalable approximations, such as K-FAC, Shampoo, and related preconditioners, address this by imposing structural constraints (e.g., block-diagonal, Kronecker-factored) on the curvature matrix early in the derivation. While this makes inversion tractable, it discards cross-layer interactions before the optimization problem defining the update is even solved. The paper argues that this premature structural imposition limits the ability of these optimizers to capture the true geometry of the loss landscape, particularly the inter-layer couplings induced by the network's computation graph.

2. Methodology: Layerwise LQR (LLQR)

The authors propose Layerwise LQR (LLQR), a framework that reframes the geometry-aware update step as a finite-horizon Linear Quadratic Regulator (LQR) problem. This approach separates the network's dynamics from the choice of descent geometry, allowing for a scalable relaxation that learns structured preconditioners while retaining the layer-coupled objective.

Core Theoretical Insight:
The paper establishes an exact equivalence between the steepest-descent step under a broad class of divergence-induced quadratic models (including Newton, Gauss-Newton, Fisher/natural-gradient, and intermediate-layer metrics) and a finite-horizon LQR problem.

• Dynamics: The forward pass of the neural network defines linear perturbation dynamics: $\delta x_{i+1} = A_i \delta x_i + B_i \delta \theta_i$, where $A_i$ and $B_i$ are Jacobians of the layer maps.
• Cost: The chosen divergence (e.g., KL divergence for NGD, Bregman gap for Newton) defines the quadratic cost matrices ($Q_i, R_i, M_i$) associated with state and control perturbations.
• Exact Solution: The exact geometry-aware update can be recovered by solving this LQR problem via backward Riccati recursions, which compute local gain matrices and adjoints without forming the global dense Hessian.

Scalable Relaxation:
While the exact Riccati solution is still computationally expensive for large networks due to Jacobian-dependent quantities, the authors introduce a scalable relaxation. Instead of solving for the exact update $\delta \theta$, they parameterize the update as a preconditioned gradient:
$$\Delta \theta_i = -U_i \nabla_{\theta_i} L(\theta)$$
where $U = \text{diag}(U_0, \dots, U_{N-1})$ is a learned structured inverse preconditioner (e.g., diagonal, Kronecker-factored, or E-KFAC).

Crucially, the block structure is imposed on the learned preconditioner $U$, not on the curvature matrix itself. The preconditioner is learned by minimizing the LQR objective (Eq. 15) over a minibatch. This allows the optimizer to approximate the dense, layer-coupled geometry using structured blocks, effectively trading expressivity for scalability while maintaining a principled connection to the original second-order geometry.

Algorithmic Implementation:
The method wraps standard optimizers (e.g., SGDM, AdamW). Periodically (every $n$ iterations), the algorithm:

1. Linearizes the network dynamics ($A_i, B_i$) and forms local cost blocks ($Q_i, R_i, M_i$) based on the chosen divergence.
2. Solves an inner optimization problem to update the preconditioner $U$ using a standard optimizer (e.g., SGDM) to minimize the relaxed LQR objective.
3. Applies an Exponential Moving Average (EMA) to stabilize $U$.
4. Uses the updated $U$ to precondition gradients for subsequent outer-loop steps.

3. Key Contributions

1. Layerwise Optimal-Control Formulation: The paper demonstrates that steepest descent under a broad class of divergence-induced quadratic models can be written exactly as a finite-horizon LQR problem. This provides a new theoretical reference for geometry-aware updates that explicitly separates network dynamics from the metric choice.
2. Scalable Relaxation via Learned Preconditioners: The authors propose learning structured inverse preconditioners directly by minimizing the LQR objective. This yields a family of optimizers that can utilize diagonal, Kronecker-factored, or E-KFAC structures while preserving the layer-coupled objective induced by the original dense model.
3. Practical Optimizer Wrapper: The relaxed LLQR update is implemented as a wrapper for modern architectures (ResNets, Transformers) that reuses learned preconditioners across iterations, avoiding explicit curvature inversion and adding modest computational overhead.
4. Empirical Validation: Extensive experiments show that LLQR improves optimization dynamics and final test performance on image classification (CIFAR, ImageNet) and machine translation (IWSLT14) benchmarks. It also accelerates "grokking" in Transformers.

4. Experimental Results

• Toy Validation: On the Rosenbrock function, the exact LQR solution (via Riccati recursion) perfectly matches Newton's method. The relaxed LLQR with block-diagonal preconditioners converges faster than standard gradient descent and tracks the Newton trajectory more closely than diagonal-Hessian approximations, validating the ability of the method to capture inter-layer couplings.
• CIFAR-10/100: On ResNet-18, LLQR with E-KFAC structure consistently improves Top-1 accuracy over baselines (SGDM, AdamW) with only a modest increase in wall-clock time (e.g., $\times 1.03$ to $\times 1.15$). Diagonal preconditioners showed less improvement, suggesting Kronecker structures are necessary to capture curvature.
• ImageNet: Training ResNet-50 for 100 epochs, LLQR+E-KFAC with NGD achieved 78.05% Top-1 accuracy compared to 77.42% for the SGDM baseline, with a computational overhead of only $\approx 1.03\times$.
• Transformers (IWSLT14): LLQR+E-KFAC improved BLEU scores from 34.24 to 34.51 on German-to-English translation with a $1.16\times$ slowdown.
• Grokking: In algorithmic datasets, LLQR consistently accelerated the onset of grokking (sudden generalization) in terms of iteration count and wall-clock time compared to baselines.
• Efficiency Comparison: When compared against AdaFisher and other second-order methods under matched wall-clock budgets, LLQR achieved higher accuracy, demonstrating that richer preconditioner structures (E-KFAC) can be made practical at scale.

5. Significance and Claims

The paper positions LLQR as a practical framework for geometry-aware second-order methods that bridges the gap between theoretical optimality and scalability.

• Principled Approximation: Unlike methods that approximate the curvature matrix first, LLQR derives the update objective from the dense geometry and then restricts the preconditioner class. This ensures the learned preconditioner is optimized in the presence of cross-layer couplings encoded by the LQR dynamics.
• Flexibility: The framework is divergence-agnostic (supporting Newton, NGD, etc.) and structure-agnostic (supporting diagonal, Kronecker, E-KFAC).
• Efficiency: By amortizing the cost of learning the preconditioner and applying it inversion-free, LLQR shifts expressive preconditioning from a theoretically attractive but often impractical option into a computationally viable regime for large-scale deep learning.

The authors acknowledge limitations, noting that LLQR introduces memory and compute overhead for storing and refitting the preconditioner $U$. However, they argue this cost is controllable via implementation knobs (update frequency, chunk size) and is justified by the performance gains and the ability to use richer structures than standard diagonal approximations.