Layerwise LQR for Geometry-Aware Optimization of Deep Networks

• URL: http://arxiv.org/abs/2605.04230v1
• Date: Tue, 05 May 2026 19:16:00 GMT
• Title: Layerwise LQR for Geometry-Aware Optimization of Deep Networks
• Abstract summary: We introduce Layerwise LQR, a framework for learning structured inverse preconditioners under a global layerwise optimal-control objective.<n>We derive a scalable relaxation that learns diagonal, (E-)Kronecker-factored, or other structured inverse preconditioners.<n> Experiments on ResNets and Transformers show that LLQR improves optimization dynamics and often translates these gains into improved final test performance.
• Score: 38.84575015854741
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Hierarchical Grading in Large Language Models [0.0]
  We introduce Graded Large Language Models (GLLMs)<n>GLLMs are a framework that equips the representation space of a transformer with a grading.<n>The construction extends the theory of graded neural networks and graded transformers to autoregressive language models.
  arXiv  Detail & Related papers  (2026-07-23T20:14:48Z)
• Towards The Implicit Bias on Multiclass Separable Data Under Norm Constraints [3.421126849736858]
  We introduce NucGD, a geometry-aware geometry designed to enforce low rank structures through nuclear norm constraints.<n>To enable scalable training, we derive an efficient SVD-free update rule via asynchronous power iteration.<n>We empirically dissect the impact of gradient optimization dynamics, characterizing how varying levels of gradient noise induced by mini-batch sampling and momentum modulate the convergence toward the expected maximum margin solutions.
  arXiv  Detail & Related papers  (2026-03-24T05:51:50Z)
• GPU-friendly and Linearly Convergent First-order Methods for Certifying Optimal $k$-sparse GLMs [7.079949618914198]
  Branch-and-Bound (BnB) frameworks can certify optimality using perspective relaxations.<n>Existing methods for solving these relaxations are computationally intensive, limiting their scalability.<n>We develop a unified proximal framework that is both linearly convergent and computationally efficient.
  arXiv  Detail & Related papers  (2026-03-01T22:26:09Z)
• GeoRA: Geometry-Aware Low-Rank Adaptation for RLVR [10.820638016337869]
  We propose GeoRA, which exploits the anisotropic and compressible nature of RL update subspaces.<n>GeoRA mitigates optimization bottlenecks caused by geometric misalignment.<n>It consistently outperforms established low-rank baselines on key mathematical benchmarks.
  arXiv  Detail & Related papers  (2026-01-14T10:41:34Z)
• Noise-Adaptive Layerwise Learning Rates: Accelerating Geometry-Aware Optimization for Deep Neural Network Training [31.259303817974693]
  We introduce a noise-adaptive layerwise learning rate scheme on top of geometry-aware optimization algorithms.<n>Our method estimates gradient variance in the dual norm induced by the chosen LMO on the fly.<n>Our algorithm achieves a sharp convergence rate.
  arXiv  Detail & Related papers  (2025-10-15T18:42:13Z)
• Preconditioned Norms: A Unified Framework for Steepest Descent, Quasi-Newton and Adaptive Methods [50.070182958880146]
  We propose a unified framework generalizing descent, quasi-Newton methods, and adaptive methods through the novel notion of preconditioned matrix norms.<n>Within this framework, we provide the first systematic treatment of affine and scale invariance in the matrix- parameterized setting.<n>We introduce two new methods, $ttMuAdam$ and $texttMuAdam-SANIA$, which combine the spectral geometry of Muon with Adam-style preconditioning.
  arXiv  Detail & Related papers  (2025-10-12T19:39:41Z)
• Stable Nonconvex-Nonconcave Training via Linear Interpolation [51.668052890249726]
  This paper presents a theoretical analysis of linearahead as a principled method for stabilizing (large-scale) neural network training. We argue that instabilities in the optimization process are often caused by the nonmonotonicity of the loss landscape and show how linear can help by leveraging the theory of nonexpansive operators.
  arXiv  Detail & Related papers  (2023-10-20T12:45:12Z)
• A Unified Algebraic Perspective on Lipschitz Neural Networks [88.14073994459586]
  This paper introduces a novel perspective unifying various types of 1-Lipschitz neural networks. We show that many existing techniques can be derived and generalized via finding analytical solutions of a common semidefinite programming (SDP) condition. Our approach, called SDP-based Lipschitz Layers (SLL), allows us to design non-trivial yet efficient generalization of convex potential layers.
  arXiv  Detail & Related papers  (2023-03-06T14:31:09Z)
• Smooth over-parameterized solvers for non-smooth structured optimization [3.756550107432323]
  Non-smoothness encodes structural constraints on the solutions, such as sparsity, group sparsity, low-rank edges and sharp edges. We operate a non-weighted but smooth overparametrization of the underlying nonsmooth optimization problems. Our main contribution is to apply the Variable Projection (VarPro) which defines a new formulation by explicitly minimizing over part of the variables.
  arXiv  Detail & Related papers  (2022-05-03T09:23:07Z)
• Leveraging Non-uniformity in First-order Non-convex Optimization [93.6817946818977]
  Non-uniform refinement of objective functions leads to emphNon-uniform Smoothness (NS) and emphNon-uniform Lojasiewicz inequality (NL) New definitions inspire new geometry-aware first-order methods that converge to global optimality faster than the classical $Omega (1/t2)$ lower bounds.
  arXiv  Detail & Related papers  (2021-05-13T04:23:07Z)
• Cogradient Descent for Bilinear Optimization [124.45816011848096]
  We introduce a Cogradient Descent algorithm (CoGD) to address the bilinear problem. We solve one variable by considering its coupling relationship with the other, leading to a synchronous gradient descent. Our algorithm is applied to solve problems with one variable under the sparsity constraint.
  arXiv  Detail & Related papers  (2020-06-16T13:41:54Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.