A Local Polyak-Lojasiewicz and Descent Lemma of Gradient Descent For Overparametrized Linear Models

TL;DR

This work addresses why gradient descent can efficiently train overparameterized two-layer linear networks with general losses. It develops a trajectory-aware optimization geometry by introducing a weight-dependent operator \(\mathcal{T}\) and proving local PL and Descent Lemma bounds that hold along GD iterates, enabling a linear convergence result with an adaptive step size. The key contributions include explicit expressions for the local constants \(\mu_t\) and \(K_t\), a convergence theorem that yields a rate depending on initialization through \(\alpha_1/\alpha_2\) and a width-dependent conditioning of \(\mathcal{T}\), and an adaptive step-size schedule that accelerates convergence. The results show that, with proper initialization and sufficient width, the overparameterized model can approach the convergence speed of the non-overparameterized model, supported by experiments validating the theory and practical step-size gains.

Abstract

Most prior work on the convergence of gradient descent (GD) for overparameterized neural networks relies on strong assumptions on the step size (infinitesimal), the hidden-layer width (infinite), or the initialization (large, spectral, balanced). Recent efforts to relax these assumptions focus on two-layer linear networks trained with the squared loss. In this work, we derive a linear convergence rate for training two-layer linear neural networks with GD for general losses and under relaxed assumptions on the step size, width, and initialization. A key challenge in deriving this result is that classical ingredients for deriving convergence rates for nonconvex problems, such as the Polyak-Łojasiewicz (PL) condition and Descent Lemma, do not hold globally for overparameterized neural networks. Here, we prove that these two conditions hold locally with local constants that depend on the weights. Then, we provide bounds on these local constants, which depend on the initialization of the weights, the current loss, and the global PL and smoothness constants of the non-overparameterized model. Based on these bounds, we derive a linear convergence rate for GD. Our convergence analysis not only improves upon prior results but also suggests a better choice for the step size, as verified through our numerical experiments.

Figures (5)

• Figure 1: Tightness of the theoretical upper bound versus reconstruction error $L(t)$ for different choices of step size in §\ref{['subsec:converg_over']}, shown in different colors. We run the simulations for nine different settings of initialization and data generation. For each setting, we repeat the simulation thirty times. The triangle lines represent the theoretical upper bound on the training loss in equation \ref{['eqn:bound_constant']} and equation \ref{['eqn:theoretical_bounds_loss']}. The solid lines represent the mean of the $\log_{10}$ of the reconstruction error $L(t)$. The shaded area is the mean of $\log_{10}L(t)$ plus and minus one standard deviation.
• Figure 2: Evolution of the loss and of the step size for different choices of the step size schedule under different initialization and data generation. We run the simulations thirty times. For each setting, we repeat the simulation thirty times. The solid lines represent the mean of $\log_{10}$ of the reconstruction error $L(t)$. The shaded area is the mean of $\log_{10}L(t)$ plus and minus one standard deviation.
• Figure 3: Plot of a toy example illustrating the loss function $\ell(x)$ and its overparametrized version $L(x_1,x_2)$, along with the corresponding local PL constant $\mu_{\mathrm{over}}(x_1,x_2)$ and smoothness constant $K_{\mathrm{over}}(x_1,x_2)$. The definitions of $\ell(x)$, $L(x_1,x_2)$, $\mu_{\mathrm{over}}(x_1,x_2)$, and $K_{\mathrm{over}}(x_1,x_2)$ are given in equation \ref{['eqn:toy1']}, equation \ref{['eqn:toy2']}, and equation \ref{['eqn:toy3']}.
• Figure 4: $\kappa(\mathcal{T}_0)$ under different choices of $p$ and $h$. We repeat the simulation thirty times and plot the average value of $\kappa(\mathcal{T}_0)$.
• Figure 5: Comparison of convergence rate of GD for the non-overparametrized model and overparametrized model. We run the simulations thrity times. The red line represents $\log_{10} \ell(t)$ and the blue line represents $\log_{10} L(t)$. The shaded area represents plus and minus one standard deviation of the reported loss.

Theorems & Definitions (43)

• Proposition 2.1: Non-existence of global PL constant and smoothness constant
• Theorem 3.1: Local Descent Lemma and PL condition for GD
• Theorem 3.2: Linear convergence of GD for \ref{['eqn:obj_over']}
• Lemma 3.1: Uniform spectral bounds on $\mathcal{T}_t, W(t)$.
• Lemma 3.2: Induction step to show $\mu_t, K_t$ is bounded and $L(t)$ converges linearly.
• Lemma 3.3: Mild overparametrization ensures $\alpha_1>0$
• Lemma 3.4: Lemma 1 in min2021explicit
• Theorem 3.3
• Lemma A.1: Inequality on the Frobenius norm
• Lemma A.2: Singular values of $\mathcal{T}$