stat.MLMar 11, 2026

Dual Space Preconditioning for Gradient Descent in the Overparameterized Regime

Authors: Reza Ghane, Danil Akhtiamov, Babak Hassibi

Organizations: Electrical Engineering Department, California Institute of Technology, Pasadena, CA 91125 USA · Computing + Mathematical Sciences Department, California Institute of Technology, Pasadena, CA 91125 USA

Abstract

In this work, we study the convergence properties of the Dual Space Preconditioned Gradient Descent, encompassing optimizers such as Normalized Gradient Descent and Gradient Clipping. We consider preconditioners of the form ∇K\nabla K, where K:Rd×k→RK: \mathbb{R}^{d \times k} \to \mathbb{R} is convex and apply ∇K(⋅)\nabla K(\cdot) to train an over-parameterized linear model with a convex loss of the form ℓ(XW−Y)\ell(X W - Y), for weights W∈Rd×kW \in \mathbb{R}^{d \times k}, labels Y∈Rn×kY \in \mathbb{R}^{n \times k} and data X∈Rn×dX \in \mathbb{R}^{n \times d}. Under the aforementioned assumptions, we prove that the iterates of the full-batch preconditioned gradient descent converge at an exponential rate to a point W∞∈Rd×kW_{\infty} \in \mathbb{R}^{d \times k} satisfying XW∞=YXW_{\infty} = Y. We also study the implicit bias of Dual Space Preconditioned Gradient Descent. First, we demonstrate analytically and empirically that, for general K(⋅)K(\cdot), W∞W_\infty depends on the chosen constant step size, hindering a precise characterization of the implicit bias. We also provide an approximate implicit bias property for general preconditioners, namely, ∥W0−W∞∥F≤c∥W0−WGD,∞∥F\|W_0 - W_{\infty}\|_F \le c \|W_0 - W_{\text{GD}, \infty}\|_F for a constant c>0c>0 and WGD,∞W_{\text{GD}, \infty} denoting the convergence point of GD initialized at W0W_0. Furthermore, for preconditioners of the form K(G)=h(∥G∥F)K(G) = h(\|G\|_F), known as {\it isotropic preconditioners}, and for the stochastic variation of the algorithm with arbitrary batch-size, we prove linear convergence to WGD,∞W_{\text{GD}, \infty}. Finally, in the experiments, we demonstrate faster convergence on a nonlinear model obtained using the smoothed matrix elastic-net as a preconditioner.

Figures & tables

Explore similar work

May 29, 2026

Local linear convergence of gradient methods for overparameterized Gaussian mixtures

We study the problem of learning Gaussian mixture models under overparameterization. Prior work has shown that while overparameterization is essential for avoiding spurious local optima and enables global recovery of the ground-truth model using the gradient-EM (expectation-maximization) algorithm, it can dramatically slow down the local rate of convergence. Under certain assumptions on the mixture weights, we show that a standard divergence measure minimized by statistical learning procedures possesses a manifold of slow growth on which the well-known Polyak stepsize reduces the loss geometrically, and design a gradient-based method that converges to minimizers at a locally linear rate. Additionally, we show that our method converges to nearly optimal solutions -- up to a natural misspecification threshold -- for mixtures with arbitrary weights. At a high level, the method alternates between several "short" gradient descent steps that approach the manifold and "long" Polyak steps that contract the distance to minimizers. Our results suggest that slow convergence is not an intrinsic challenge of overparameterization, but can be overcome by exploiting the favorable structure of the loss landscape.
Mar 2, 2026

Training Dynamics of Softmax Self-Attention: Fast Global Convergence via Preconditioning

We study the training dynamics of gradient descent in a softmax self-attention layer trained to perform linear regression and show that a simple first-order optimization algorithm can converge to the globally optimal self-attention parameters at a geometric rate. Our analysis proceeds in two steps. First, we show that in the infinite-data limit the regression problem solved by the self-attention layer is equivalent to a nonconvex matrix factorization problem. Second, we exploit this connection to design a novel "structure-aware" variant of gradient descent which efficiently optimizes the original finite-data regression objective. Our optimization algorithm features several innovations over standard gradient descent, including a preconditioner and regularizer which help avoid spurious stationary points, and a data-dependent spectral initialization of parameters which lie near the manifold of global minima with high probability.
May 7, 2026

Pro-KLShampoo: Projected KL-Shampoo with Whitening Recovered by Orthogonalization

Optimizers that exploit the matrix structure of gradients are central to modern LLM pre-training, with two distinct frontiers: explicit Kronecker-factored preconditioning -- most recently KL-Shampoo, which estimates the preconditioner via KL divergence minimization -- and orthogonalization of the gradient momentum, exemplified by Muon and analyzed as steepest descent under the spectral norm. The two routes are typically developed in isolation. We make a structural observation about KL-Shampoo's Kronecker preconditioners: their eigenvalue spectra exhibit a spike-and-flat shape -- a few dominant eigenvalues followed by an approximately uniform tail -- across layers and training stages, holding exactly under a rank-ρρ signal-plus-noise gradient model. We exploit this structure by restricting one of KL-Shampoo's Kronecker factors to a parametric family aligned with the spike-and-flat shape: full spectral structure on a tracked rr-dimensional subspace, single shared eigenvalue across the remaining n−rn-r directions. On these directions, we apply orthogonalization. An identity shows that this orthogonalization recovers the algebraic form of full KL-Shampoo's preconditioner. On four pre-training scales (GPT-2 124M / 350M, LLaMA 134M / 450M), Pro-KLShampoo consistently outperforms KL-Shampoo at every subspace rank we test in validation loss, peak per-GPU memory, and wallclock time to reach each loss level.