Robust Implicit Regularization via Weight Normalization

TL;DR

The paper investigates implicit regularization in overparameterized diagonal linear networks trained with gradient flow and weight normalization (WN). By reparameterizing weights in polar form and introducing a learning-rate ratio, the authors prove that WN yields a robust bias toward sparse (L1-minimizing) solutions without requiring tiny initializations, extending prior results that relied on small initialization. They establish invariants, use a Liouville/Lojaowski framework to show convergence (linear for L=2) and exponential decay of the loss, and demonstrate a magnified implicit regularization via a magnification factor $\rho$. Empirical results across varying depths, initializations, and ground-truth sparsity confirm that WN-GD often achieves substantially lower reconstruction error and faster convergence than standard GD, supporting the practical relevance of the approach for sparse recovery in overparameterized settings.

Abstract

Overparameterized models may have many interpolating solutions; implicit regularization refers to the hidden preference of a particular optimization method towards a certain interpolating solution among the many. A by now established line of work has shown that (stochastic) gradient descent tends to have an implicit bias towards low rank and/or sparse solutions when used to train deep linear networks, explaining to some extent why overparameterized neural network models trained by gradient descent tend to have good generalization performance in practice. However, existing theory for square-loss objectives often requires very small initialization of the trainable weights, which is at odds with the larger scale at which weights are initialized in practice for faster convergence and better generalization performance. In this paper, we aim to close this gap by incorporating and analyzing gradient flow (continuous-time version of gradient descent) with weight normalization, where the weight vector is reparameterized in terms of polar coordinates, and gradient flow is applied to the polar coordinates. By analyzing key invariants of the gradient flow and using Lojasiewicz Theorem, we show that weight normalization also has an implicit bias towards sparse solutions in the diagonal linear model, but that in contrast to plain gradient flow, weight normalization enables a robust bias that persists even if the weights are initialized at practically large scale. Experiments suggest that the gains in both convergence speed and robustness of the implicit bias are improved dramatically by using weight normalization in overparameterized diagonal linear network models.

Figures (5)

• Figure 1: WN converges to minimal $\ell_1$-norm solutions from a wider range of initialization, and hence is more robust in the sense that it is not sensitive to the choice of initialization. This suggests that GF with WN could be used as an efficient alternative for $\ell_1$-minimization. Each data point in Figure \ref{['fig:WN_TestError_Initialization_0']} is an average over multiple random initializations at fixed scale. The improvement ratio (reconstruction error for GD divided by reconstruction error for WN-GD) can be huge; when the initialization scale $\alpha$ (defined in \ref{['eq:alpha']}) equals to $0.1$, such ratio is more than $10^5$!
• Figure 2: WN-GD yields significantly much smaller error than GD. The training loss converges to values close to zero.
• Figure 3: As the learning rate ratio ${\tilde{\eta}}$ decreases, the error decreases. Note that when ${\tilde{\eta}}=10$, the assumption of Theorem \ref{['theorem:optimality_constant_rate']} is violated, and we see that WN-GD is not better than GD.
• Figure 4: Comparison between $L=2$ and $L=3$. WN-GD is better in both cases, but $L=3$ requires smaller initialization.
• Figure 5: The setting is the same as in Figure \ref{['fig:exp1']}, except that in such setting we can recover ground truth vectors that are not necessarily non-negative.

Theorems & Definitions (54)

• Theorem 2.1: Theorem 2.1 from chou2021more
• Theorem 2.2: Magnification of implicit regularization
• proof
• Theorem 2.3: Convergence
• proof
• Theorem 2.4: A time-dependent learning rate
• proof
• Lemma 3.1: Re-scaled learning rate $\eta_{{\bf u}}$
• proof
• Lemma 3.2: Constant norm