Towards better generalization: Weight Decay induces low-rank bias for neural networks

## Abstract

Abstract We study the implicit bias towards low-rank weight matrices when training neural networks (NN) with Weight Decay (WD).
We prove that when a ReLU NN is sufficiently trained with Stochastic Gradient Descent (SGD) and WD, its weight matrix is approximately a rank-two matrix.
Empirically, we demonstrate that WD is a necessary condition for inducing this low-rank bias across both regression and classification tasks.
Our work differs from previous studies as our theoretical analysis does not rely on common assumptions regarding the training data distribution, optimality of weight matrices, or specific training procedures.
Furthermore, by leveraging the low-rank bias, we derive improved generalization error bounds and provide numerical evidence showing that better generalization can be achieved.
Thus, our work offers both theoretical and empirical insights into the strong generalization performance of SGD when combined with WD.

## 1 Introduction

We consider the learning problem of a neural network

$$ $\min_{\theta\in\Theta}F(\theta)\coloneqq\mathbb{E}_{z}\left[l(z,\phi_{\theta}) \right]\,,$ (1) $$

where $\theta$ represents the parameter of the neural network $\phi_{\theta}$, $\Theta\subset\mathbb{R}^{d}$ is the parameter space, $F:\Theta\rightarrow\mathbb{R}$ is called the population loss, and $l(z,\phi_{\theta})$ measures the loss function of $\phi_{\theta}$ over a random data point $z\in\mathcal{Z}$. The optimization problem above aims to fit the data distribution $\mathcal{Z}$ by adjusting the model parameter $\theta$.
In practice, the data distribution $\mathcal{Z}$ is unknown thus the expectation over $z$ cannot be calculated. An optimization of the empirical loss function $\hat{F}:\Theta\rightarrow\mathbb{R}$ is considered instead

$$ $\min_{\theta\in\Theta}\hat{F}(\theta)\coloneqq\frac{1}{N}\sum_{i=1}^{N}l(z_{i} ,\phi_{\theta})\,.$ (2) $$

Here $N$ is the total number of data points and $\mathsf{z}\coloneqq[z_{1},\ldots,z_{N}]\in\mathcal{Z}^{N}$ are independent and identically distributed (i.i.d.) samples from $\mathcal{Z}$.
A learning algorithm solving Equation 2 can be written as $A:\mathcal{Z}^{N}\rightarrow\Theta$ that maps the training data $\mathsf{z}$ to a parameter $A(\mathsf{z})$. Let $\theta^{\ast}$ denote a minimizer of Equation 2, the population loss of the algorithm $A$ can be decomposed into three parts:

$$ $F(A(\mathsf{z}))=\underbrace{\left(F(A(\mathsf{z}))-\hat{F}(A(\mathsf{z})) \right)}_{\text{generalization error}}+\underbrace{\left(\hat{F}(A(\mathsf{z}) )-\hat{F}(\theta^{\ast})\right)}_{\text{optimization error}}+\underbrace{\hat{ F}(\theta^{\ast})}_{\text{approximation error}}\,.$ (3) $$

The first term measures the generalization error due to finite number of training data; the second term quantifies the gap between the empirical loss of the learned parameter $A(\mathsf{z})$ and the minimum empirical loss; the last term indicates approximation power of the model hypothesis space $\{\phi_{\theta}:\theta\in\Theta\}$ in fitting the data distribution $\mathcal{Z}$.
When the dimension of the parameter space $\Theta$ increases, both the approximation error and optimization error decrease to zero due to universal approximation theory (Hornik, ; Chen and Chen, ; Yarotsky, ) and neural tangent kernel theory (Jacot et al., ) respectively.

Nevertheless, challenges remain that as classical generalization error analysis (Anthony et al., ; Feldman, ) suggests that the generalization error for a neural network with $\mathcal{O}(d)$ parameters grows at the order of $\mathcal{O}(\sqrt{d})$, regardless of any learning algorithms.

However, empirically it is observed that SGD can produce implicit regularization (Keskar et al., ; Zhu et al., ; Zhang et al., ) and thus achieve better generalization performance than Gradient Descent (GD) (Zhang et al., ).
It is found that different learning algorithms exhibit implicit biases in the trained neural network, including alignment of neural network weights (Du et al., ; Ji and Telgarsky, , ; Lyu and Li, ) and low-rank bias (Timor et al., ; Jacot, ; Xu et al., ). Although these implicit biases are not theoretically understood yet,
the generalization error is expected to be improved by leveraging these implicit biases. As a consequence, one can obtain “algorithm-dependent" generalization bounds that are usually smaller than the uniform bound.

In this work, we focus on the low-rank bias observed in many computer vision problems (Yu et al., ; Alvarez and Salzmann, ; Arora et al., ).
In particular, it is numerically shown that replacing the weight matrices by their low-rank approximation only results in small prediction error. The line of works (Ji and Telgarsky, ; Lyu and Li, ) generalizes the theory of Ji and Telgarsky () and proves the implicit bias effect of SGD algorithms with entropy loss. Galanti and Poggio (),Timor et al. () and Jacot () first studied the low-rank bias for neural network trained with WD for square loss. They show that when minimizing the $l_{2}$ norm of coefficients of a neural network that attains zero loss, the coefficient matrices are of low-rank for sufficiently deep neural networks. Similar low-rank bias is also empirically observed in other training settings, including Sharpness Aware Minimization (Andriushchenko et al., ) and normalized weight decay (Xu et al., ).

**Table 1: Assumptions and results of various papers on low-rank bias. The notation $M$ denotes an upper bound of the Frobenius norm of weight matrices, $B$ denotes the batch size, and $m_{l}$ denotes the number of patches in the input space. Here GF denotes Gradient Flow and N/A means there is no specific assumption in the corresponding column. Our paper considers the most general setting and proves a strict low-rank result.**
| Paper | NN Architecture | Data | Optimization | Optimizer | Convergence | Result |
| --- | --- | --- | --- | --- | --- | --- |
| Ji and Telgarsky () | Linear NN | Linearly<br>separable | Entropy loss | GF | N/A | Each layer has rank $\leq 1$ |
| Le and Jegelka () | Top $K$ linear layers | Linearly<br>separable | Entropy loss | GF | N/A | Top $K$ layers have rank $\leq 1$ |
| Ergen and Pilanci () | NN without bias | $1$-dimensional | Min $l_{2}$ with<br>data fitting | N/A | Global min | First layer rank $\leq 1$ |
| Ongie and Willett () | NN with bias | $d$-dimensional | Min $l^{2}$ with<br>data fitting | N/A | Global min | Linear layers have rank $\leq d$ |
| Timor et al. () | $L$-layer NN without bias | Separable<br>by $L^{\prime}$-layer NN | Min $l_{2}$ with<br>data fitting | N/A | Global min | Each layer has rank $\leq M^{L/L^{\prime}}$ |
| Xu et al. () | NN without bias | N/A | Square loss<br>normalized $l_{2}$ regularization | SGD | Global min | Each layer has rank $\leq 2$ |
| Galanti and Poggio () | NN without bias | N/A | Differentiable loss<br>$l_{2}$ regularization | SGD | Convergence of<br>NN weights | Each layer has rank $\leq Bm_{l}$ |
| This paper | Two-layers NN with bias | N/A | Square loss<br>$l_{2}$ regularization | SGD | N/A | Each layer has rank $\leq 2$ |

In this work, we study the low-rank bias of two-layer ReLU neural networks trained with mini-batch SGD and weight decay using square loss function. We demonstrate, both theoretically and empirically, that weight decay leads to the low-rank bias of ReLU neural networks. We then leverage the low-rank bias to obtain a smaller generalization error of the trained neural network model. Our contributions are summarized as follows:

- •
Low-rank bias for sufficiently trained neural network.
In theorem 2.6, we prove that the matrix weight parameters in the trained neural network is close to a rank two matrix and the distance between them depends on the batch size, the strength of the weight decay, and the batch gradient. In contrast to common assumptions on attaining global minimum and sufficient depth (Timor et al., ; Jacot, ), strictly zero batch gradient (Xu et al., ), and convergence of weight matrices (Galanti and Poggio, ), we only assume that the neural network is sufficiently trained so that the batch gradient is small( see Assumption 2.5). This result is demonstrated and supported by numerical evidence of both regression problem and classification problem in Section 4.
- •
Low-rank bias leads to better generalization.
In theorem 3.11, we show that the low-rank bias leads to better generalization and thus implicitly explain the generalization effect of SGD with weight decay. For a two-layer ReLU neural networks with input dimension $m$ and width $n$ trained on $N$ data pairs, we prove that the generalization error bound can be improved from $\mathcal{O}\left(\sqrt{\frac{mn\ln m\ln N}{N}}\right)$ to $\mathcal{O}\left(\sqrt{\frac{(m+n)\ln m\ln N}{N}}\right)$. Our result thus provides an alternative explanation to the generalization power of SGD algorithm through the lens of low-rank bias. The improved generalization error bound is also echoed with the recent work of Park et al. (), where they leverage the local contraction property of SGD.

### 1.1 Discussion of related works

We have summarized relevant findings on low-rank bias in Table 1.
The studies of Ji and Telgarsky (, ); Lyu and Li () investigated the directional convergence of the parameter $\theta$ for homogenized neural network in classification tasks. Their results demonstrated that $\theta$ converges to a direction that maximized the margin of the classifiers. A comparable result is presented in Kumar and Haupt () for regression problems.
Our work aligns with these findings, showing that $\theta$ converges to a low-rank matrix.
For regression problems, Timor et al. () established that weight decay induces a low-rank bias. However, our contribution differs from Timor et al. () as their results are restricted to sufficiently deep neural networks.

To the best of the authors’ knowledge, the most closely related works to our results are Galanti and Poggio (), Xu et al. () and Park et al. ().
Galanti and Poggio () studies a fully connected neural network without bias vectors, trained with weight decay. They prove that when the norms of weight matrices converge, the normalized weight matrices are close to a matrix whose rank is bounded by the batch size and the number of patches in the input space.
Our work differs from theirs in several key aspects: we do not assume the convergence of the neural network but instead consider the case where the neural network is sufficiently trained (cf. Assumption 2.5). Additionally, we establish a stronger bound on the rank of weight matrices.

Xu et al. () extends the idea of normalized weight matrices in  Galanti and Poggio () and examines the mini-batch SGD training of neural networks with normalized WD. Rather than using the standard mini-batch SGD, they normalize all neural network weight matrices and train both the normalized weight matrices and the product of matrix norms together with mini-batch SGD.
In contrast, our work employs the SGD algorithm with conventional WD, eliminating the need for an additional normalization step. Furthermore,  Xu et al. () assumes that all batch gradients equal to zero, which is not practical in most regression problems. In comparison, we assume only that the batch gradients are small, and our assumption is numerically verified.

Park et al. () does not address low-rank bias; instead, they leverage the piecewise contraction property of SGD optimization.
They obtain a dimension-independent estimate of the generalization error due to this property. However, their proof relies on the assumption that the loss function is piecewise convex and smooth, and the generalization error estimate depends on the total number of pieces in the parameter space.
A straightforward calculation of the total number of pieces for a two-layer ReLU network leads to a generalization error of the order $\mathcal{O}\left(\sqrt{\frac{(\ln N+m)\ln N+\ln(1/\delta)}{N}}\right)$, which is essentially of the same order as our result when $m\geq n$.

## 2 Low-rank bias of mini-batch SGD

### 2.1 Setup

We consider a two-layer ReLU neural network $\phi_{\theta}:\mathcal{X}\rightarrow\mathbb{R}$

$$ $\phi_{\theta}(x)\coloneqq U\sigma(Vx+b)\,.$ (4) $$

Here $\mathcal{X}\subset\mathbb{R}^{n}$ is the input domain. The model parameter $\theta\coloneqq[U,V,b]$ consists of a row vector $U\in\mathbb{R}^{1\times m}$, a matrix $V\in\mathbb{R}^{m\times n}$, and a column vector $b\in\mathbb{R}^{m\times 1}$. As the width $m$ of the neural network increases, the neural network can approximate any continuous function Hornik ().

Note that the ReLU activation function $\sigma:\mathbb{R}\rightarrow\mathbb{R}$ is a piecewise linear function evaluated elementwisely on the vector $Vx+b$. For convenience, we rewrite the ReLU activation as a matrix operator $D(x,V,b)\in\{0,1\}^{m\times m}$, defined as follows:

$$ $D(x,V,b)=\text{diag}(\text{sign}(\sigma(Vx+b)))\,.$ (5) $$

In other word, $D(x,V,b)$ is a diagonal matrix whose diagonal entries are either one or zero, depending on the sign of the first layer output $\sigma(Vx+b)$. Consequently, we can rewrite the two-layer NN as a product of three transformations

$$ $\phi(x,\theta)=UD(x,V,b)(Vx+b)\,.$ (6) $$

We first show that $D(x,V,b)\in\{0,1\}^{m\times m}$ does not change its values for a small perturbation of $V$, consequently we have $\frac{\partial D(x,V,b)}{\partial V}=0$ almost surely in the parameter space $\Theta$.

###### Lemma 2.1 .

Consider a two-layer NN in Equation 4, for any fixed $(x,b)\in\mathbb{R}^{n}\times\mathbb{R}^{m}$,

$$ $\frac{\partial D(x,V,b)}{\partial V}=0\,,\quad\text{for all }V\in\mathbb{R}^{m \times n}/\mathcal{V}^{0}\,,$ $$

where $\mathcal{V}^{0}\in\mathbb{R}^{m\times n}$ is a measure zero set that depends on $x$ and $b$.

This result implies that the activation component of neural network model is not sensitive to small changes in the model parameter $V$, except on a measure zero set.

We show below that such a property also affects the gradient of neural network model towards a low-rank bias.

###### Lemma 2.2 .

Consider a two-layer NN in Equation 4 and fix $x\in\mathbb{R}^{n}$ and $b\in\mathbb{R}^{m}$,

$$ $\frac{\partial\phi(x;\theta)}{\partial V}=xUD(x,V,b)\quad\text{a.s. in }V\,.$ (7) $$

That is, $\frac{\partial\phi(x;\theta)}{\partial V}$ is a rank one matrix a.s. in $V$.

This result implies that gradient-based training algorithms only update the weight matrix $V$ with rank one matrix increments. Although this does not directly imply that the algorithm converges to a low-rank matrix, we show below this would happen when combined with WD.

###### Lemma 2.1 .

###### Lemma 2.2 .

### 2.2 Training dynamics of SGD with Weight Decay

Given a training dataset $\mathcal{S}=\{z_{i}\coloneqq(x_{i},y_{i}):\mathbb{R}^{n}\times\mathbb{R}\,,i
\in[N]\}$,
we consider the two-layer NN with following Mean Square Error (MSE) loss function with WD

$$ $\mathcal{L}_{\mathcal{S}}(\theta)=\frac{1}{2N}\sum_{i=1}^{N}|\phi(x_{i},\theta )-y_{i}|^{2}+\frac{\mu_{U}}{2}\|U\|_{F}^{2}+\frac{\mu_{V}}{2}\|V\|_{F}^{2}+ \frac{\mu_{b}}{2}\|b\|^{2}\,,$ (8) $$

where $\mu_{U},\mu_{V},\mu_{b}>0$ does not depend on $\theta$.

In practice, the training is implemented with mini-batch SGD method. That is, in each iteration, the gradient is updated over batch $\mathcal{S}^{\prime}\subset\mathcal{S}$ with batch size $|\mathcal{S}^{\prime}|=B$,

$$ $\frac{\partial\mathcal{L}_{\mathcal{S}^{\prime}}}{\partial V}=\frac{1}{B}\sum_ {i\in\mathcal{S}^{\prime}}(\phi(x_{i},\theta)-y_{i})\frac{\partial\phi(x_{i}, \theta)}{\partial V}+g(x_{i},y_{i})V\,.$ (9) $$

Note that we have chosen $\mu_{V}=\frac{1}{B}\sum_{i\in\mathcal{S}^{\prime}}g(x_{i},y_{i})$ for some fixed positive function $g:\mathbb{R}^{n}\times\mathbb{R}\rightarrow\mathbb{R}^{+}$.

#### 2.2.1 low-rank bias of the critical point

If the training converges to a limit $\theta^{\ast}$, then heuristically all batch gradient converges to zero. We first consider the following assumption.

###### Assumption 2.3 .

For a fixed batch size $2\leq B<|\mathcal{S}|$, the batch gradient $\frac{\partial\mathcal{L}_{\mathcal{S}^{\prime}}}{\partial V}(\theta^{\ast})=0$ for all batches $\mathcal{S}^{\prime}\subset\mathcal{S}$ of size $B$.

We can now prove the low-rank bias of training with weight decay.

###### Theorem 2.4 .

Consider a two-layer NN in Equation 4 trained with mini-batch SGD as in Equation 9, under Assumption 2.3, the neural network parameters converges to a matrix $V^{\ast}$ with rank at most two. In particular, if $g(\cdot,\cdot)=\mu_{V}$ is a constant, then $V^{\ast}$ is of rank one.

###### Assumption 2.3 .

###### Theorem 2.4 .

#### 2.2.2 low-rank bias in a neighborhood of critical points

In practice, Assumption 2.3 may not be realistic as it is impossible to train all batch gradients to zero. Instead, we consider the following assumption that the batch gradients are close to zero.

###### Assumption 2.5 .

There exists $\varepsilon>0$ and a fixed batch size $2\leq B<|\mathcal{S}|$, such that the batch gradient $\|\frac{\partial\mathcal{L}_{\mathcal{S}^{\prime}}}{\partial V}(\theta^{\ast})
\|_{F}\leq\varepsilon$ for all batches $\mathcal{S}^{\prime}\subset\mathcal{S}$ of size $B$.

Under Assumption 2.5, we show that the coefficient matrix $V^{\ast}$ is close to a low-rank matrix $\tilde{V}^{\ast}$.

###### Theorem 2.6 .

Consider a two-layer NN in Equation 4 trained with mini-batch SGD as in Equation 9, under Assumption 2.5, the neural network parameters $V^{\ast}$ satisfies that $\|V^{\ast}-\tilde{V}^{\ast}\|_{F}\leq C\varepsilon$ for some matrix $\tilde{V}^{\ast}$ with rank at most two. Here the constant $C$ only depends on the batch size $B$ and the choice of $g(\cdot,\cdot)$. In particular, if $g(\cdot,\cdot)=\mu_{V}$ is a constant, then $\tilde{V}^{\ast}$ is of rank one and $C=\frac{2B}{\mu_{V}}$.

###### Assumption 2.5 .

###### Theorem 2.6 .

## 3 Generalization Error Analysis

### 3.1 Preliminaries

In this section, we derive the generalization error of a low-rank neural network by following the techniques developed in Anthony et al. (); Bartlett et al. ().
We first define a few notions that measures the complexity of a class of functions.

###### Definition 3.1 (Covering Number) .

For any $\varepsilon>0$ and a set $W$, we define the $l_{1}$ $\varepsilon$-covering number of $W$, denoted as $\mathcal{N}(\varepsilon,W,\|\cdot\|_{1})$, to be minimal cardinality of a $\varepsilon$-cover of $W$ under $l_{1}$ metric.

###### Definition 3.2 (Uniform Covering Number) .

For any $\varepsilon>0$, a class of functions $\mathcal{F}$, and an integer $k$, we define the uniform covering number $\mathcal{N}_{1}(\varepsilon,\mathcal{F},k)$ as

$$ $\mathcal{N}_{1}(\varepsilon,\mathcal{F},k)=\max\{\mathcal{N}(\varepsilon, \mathcal{F}|_{\mathsf{x}},\|\cdot\|_{1})\,,\mathsf{x}=[x_{1},\ldots,x_{k}]\in \mathcal{X}^{k}\}\,,$ $$

where $\mathcal{F}|_{\mathsf{x}}\coloneqq\{(f(x_{1}),\ldots,f(x_{k})):f\in\mathcal{F}\}$.

It can be shown that for a bounded function class $\mathcal{F}$, the generalization error for any $f\in\mathcal{F}$ can be uniformly bounded by the uniform covering number of $\mathcal{F}$.

###### Theorem 3.3 (Theorem 17.1 in Anthony et al. ( 1999 ) ) .

Suppose $\mathcal{F}$ is a set of functions defined on a domain $\mathcal{X}$ and mapping into the real interval $[0,1]$. Let $z,z_{1},\ldots,z_{N}$ be i.i.d. random variables on $\mathcal{Z}=\mathcal{X}\times[0,1]$, $\varepsilon$ any real number in $[0,1]$, and $N$ any positive integer. Then with probability at least $1-\delta$

$$ $\sup_{f\in\mathcal{F}}|\mathbb{E}_{z}\left[l(z,f)\right]-\frac{1}{N}\sum_{i=1} ^{N}l(z_{i},f)|\leq\varepsilon\,,$ $$

where $l(z,f)=|f(x)-y|^{2}$ is the square loss and the failure probability $\delta\leq 4\mathcal{N}_{1}(\varepsilon/16,\mathcal{F},2N)\exp(-N\varepsilon^{
2}/32)$.

We then introduce the pseudo-dimension, which is closely related to the uniform covering number.

###### Definition 3.4 (Pseudo-dimension) .

Let $\mathcal{F}$ be a class of functions from $\mathcal{X}$ to $\mathbb{R}$. The pseudo-dimension of $\mathcal{F}$, denoted as $\text{Pdim}(\mathcal{F})$, is the largest integer $m$ for which there exists $(x_{1},\ldots,x_{m},y_{1},\ldots,y_{m})\in\mathcal{X}^{m}\times\mathbb{R}^{m}$ such that for any $(b_{1},\ldots,b_{m})\in\{0,1\}^{m}$ there exists $f\in\mathcal{F}$ such that

$$ $\forall i,f(x_{i})>y_{i}\Leftrightarrow b_{i}=1\,.$ $$

###### Theorem 3.5 (Theorem 18.4 in Anthony et al. ( 1999 ) ) .

Let $\mathcal{F}$ be a nonempty set of real functions from $\mathcal{X}$ to $[0,1]$, then

$$ $\mathcal{N}_{1}(\varepsilon,\mathcal{F},k)\leq e(\text{Pdim}(\mathcal{F})+1) \left(\frac{2e}{\varepsilon}\right)^{\text{Pdim}(\mathcal{F})}\,.$ $$

Combining the above results together, we obtain the following generalization bounds.

###### Proposition 3.6 .

Let $\mathcal{F}$ be a nonempty set of real functions from $\mathcal{X}$ to $[0,1]$, $z,z_{1},\ldots,z_{N}$ be i.i.d. random variables on $\mathcal{Z}=\mathcal{X}\times[0,1]$, $\varepsilon$ any real number in $[0,1]$, and $N$ any positive integer. Then with probability at least $1-\delta$

$$ $\sup_{f\in\mathcal{F}}|\mathbb{E}_{z}\left[l(z,f)\right]-\frac{1}{N}\sum_{i=1} ^{N}l(z_{i},f)|\leq C\sqrt{\frac{\ln(1/\delta)}{N}}+C\sqrt{\frac{\text{Pdim}( \mathcal{F})\ln N}{N}}\,,$ $$

where $l(z,f)=|f(x)-y|^{2}$ is the $l_{2}$ loss and $C>0$ is an independent constant.

We now consider the class of two-layer ReLU neural network functions defined as follows:

$$ $\mathcal{F}(m,n)\coloneqq\{\phi_{\theta}(x)=U\sigma(Vx+b):\theta=[U,V,b]\in \mathbb{R}^{1\times m}\times\mathbb{R}^{m\times n}\times\mathbb{R}^{m\times 1} \}\,.$ (10) $$

We further assume that the neural network functions are uniformly bounded.

###### Assumption 3.7 .

There exists an independent constant $L>0$ such that $\sup\limits_{f\in\mathcal{F}(m,n),x\in\mathcal{X}}|f(x)|\leq L$.

The pseudo-dimension of two-layer ReLU neural networks was studied in Bartlett et al. ().
We use the following simplified version of Theorem 7 in Bartlett et al. ().

###### Theorem 3.8 (Theorem 7 in Bartlett et al. ( 2019 ) ) .

Let $\mathcal{F}$ denotes neural networks with $W$ parameters and $U$ ReLU activation functions in $L$ layers, then

$$ $\text{Pdim}(\mathcal{F})\leq\mathcal{O}\left(LW\ln(U)\right)\,.$ $$

A simple consequence for two-layer ReLU neural networks is

$$ $\text{Pdim}(\mathcal{F}(m,n))\leq Cmn\ln(m)\,,$ (11) $$

where $C>0$ is an independent constant. As a consequence, an algorithm independent generalization bound can be obtained by combining Equation 11 and Proposition 3.6.

###### Corollary 3.9 .

Let $\phi_{\theta}$ be any two-layer neural network from Equation 10, $z,z_{1},\ldots,z_{N}$ be i.i.d. random variables on $\mathcal{Z}=\mathcal{X}\times[0,1]$. Under Assumption 3.7, then with probability at least $1-\delta$,

$$ $|\mathbb{E}_{z}\left[l(z,\phi_{\theta})\right]-\frac{1}{N}\sum_{i=1}^{N}l(z_{i },\phi_{\theta})|\leq CL^{2}\sqrt{\frac{\ln(1/\delta)}{N}}+CL^{2}\sqrt{\frac{ mn\ln m\ln N}{N}}\,,$ $$

where $l(z,\phi_{\theta})=|\phi_{\theta}(x)-y|^{2}$ is the $l^{2}$ loss and $C>0$ is an independent constant.

The above estimate holds for any two-layer ReLU neural network functions thus it is an algorithm independent bound of the generalization error. In Theorem 2.4, we have shown that the neural network parameter $\theta$ is close to a low-rank parameter $\theta^{\ast}$ when trained with mini-batch SGD. This low-rank bias essentially implies the learned neural network function $\phi_{\theta^{\ast}}$ belongs to a neighborhood of a smaller function class than $\mathcal{F}(m,n)$. We show below that low-rank bias can improve the generalization error from $\mathcal{O}\left(\sqrt{\frac{mn\ln m\ln N}{N}}\right)$ to $\mathcal{O}\left(\sqrt{\frac{(m+n)\ln m\ln N}{N}}\right)$.

###### Definition 3.1 (Covering Number) .

###### Definition 3.2 (Uniform Covering Number) .

###### Theorem 3.3 (Theorem 17.1 in Anthony et al. ( 1999 ) ) .

###### Definition 3.4 (Pseudo-dimension) .

###### Theorem 3.5 (Theorem 18.4 in Anthony et al. ( 1999 ) ) .

###### Proposition 3.6 .

###### Assumption 3.7 .

###### Theorem 3.8 (Theorem 7 in Bartlett et al. ( 2019 ) ) .

###### Corollary 3.9 .

### 3.2 Generalization error of a low-rank NN

For a two-layer ReLU neural network $\phi_{\theta}$ with a rank-$k$ matrix $V\in\mathbb{R}^{m\times n}$, we reparametrize $V=V_{2}V_{1}$ with $V_{1}\in\mathbb{R}^{k\times n}$ and $V_{2}\in\mathbb{R}^{m\times k}$. In other words, we can rewrite $\phi_{\theta}(x)$ as a composition of a linear function $y=V_{1}x$ and a two-layer ReLU neural network $\tilde{\phi}_{\theta}(y)\in F(m,k)$ with input dimension $k$.
We thus define the linear function class

$$ $L(k,n)=\{f:\mathbb{R}^{n}\rightarrow\mathbb{R}^{k}:f(x)=Ax,A\in\mathbb{R}^{k \times n}\}\,.$ $$

We define the composition of functions from $L(k,n)$ and $F(m,k)$, i.e. the low-rank two-layer ReLU neural networks, as the following

$$ $\tilde{\mathcal{F}}(m,n,k)\coloneqq\{f\circ g:f\in F(m,k),g\in L(k,n)\}\,.$ (12) $$

We can easily calculate the pseudo-dimension of $\tilde{\mathcal{F}}(m,n,k)$ using the following property of pseudo-dimension.

###### Lemma 3.10 .

For any positive integers $m,n,k$,

$$ $\text{Pdim}(\tilde{\mathcal{F}}(m,n,k))\leq C(m+n)k\ln(m)\,.$ $$

We can now calculate the generalization error for neural networks trained with mini-batch SGD and WD.

###### Theorem 3.11 .

Let $\phi_{\theta^{\ast}}$ be a two-layer neural network in Equation 4 trained with mini-batch SGD as in Equation 9 with i.i.d. data $\mathsf{z}\in\mathcal{Z}^{N}$, and $z\in\mathcal{Z}$ is an independent copy of $z_{1}$.
Suppose Assumption 3.7 and Assumption 2.3 is satisfied, then with probability at least $1-\delta$, we have

$$ $|\mathbb{E}_{z}\left[l(z,\phi_{\theta^{\ast}})\right]-\frac{1}{N}\sum_{i=1}^{N }l(z_{i},\phi_{\theta^{\ast}})|\leq CL^{2}\sqrt{\frac{\ln(1/\delta)}{N}}+CL^{2 }\sqrt{\frac{(m+n)\ln m\ln N}{N}}\,,$ $$

where $l(z,\phi_{\theta})=|\phi_{\theta}(x)-y|^{2}$ is the $l^{2}$ loss and $C>0$ is an independent constant.

###### Lemma 3.10 .

###### Theorem 3.11 .

## 4 Numerical Experiments

In this section, we numerically demonstrate that WD leads to low-rank bias for two-layer neural networks.
Our experiments are conducted on two datasets: the California housing dataset for regression tasks and the MNIST dataset for classification. We investigate how the regularization coefficient $\mu_{V}$ and batch size $B$ affect the rank of $V$, providing evidence to support Theorem 2.6. Additionally, we validate Assumption 2.5, which suggests that sufficient training leads to small batch gradients. The rank of the coefficient matrix $V$ is measured by the stable rank $r_{\text{s}}(V)=\frac{\|V\|^{2}_{F}}{\|V\|_{2}^{2}}$. The generalization error in Equation 3 is computed as the difference between the testing Mean Squared Error (MSE) and the training MSE.

### 4.1 Experimental Setup

#### 4.1.1 Training and Initialization

All experiments are conducted using two-layer neural networks with network sizes and hyperparameters specified in Table 2. These networks are trained with mini-batch SGD and the loss function defined in Equation 8. Note that here we use the square loss instead of the cross entropy loss for classification task.
The weights matrices $U$ and $V$ are initialized with Kaiming initialization (He et al., ), while the bias $b$ is initialized from the zero vector. The learning rate is set to $0.0001$ with a decay rate of 0.95 every 200 epochs. For simplicity, we have chosen $g$ as a constant function thus $\mu_{V}$ remains constant during the training dynamics.

**Table 2: Hyperparameters for neural network training**
| Dataset | Width of NN | $\mu_{U}$ | $\mu_{b}$ | Epoch |
| --- | --- | --- | --- | --- |
| California Housing Prices | 8192 | 1e-4 | 1e-4 | 5000 |
| MNIST | 32768 | 1e-6 | 1e-6 | 2000 |

#### 4.1.2 Datasets

California Housing Prices.
The California housing dataset (Pace and Barry, ) comprises 20640 instances, each with eight numerical features. The target variable is the median house value for California districts, measured in hundreds of thousands of dollars. We fetch dataset by scikit-learn package (Pedregosa et al., ), and randomly samples 1800 instances for training and 600 for testing.

MNIST.
The MNIST dataset (Deng, ) consists of 70000 grayscale images of handwritten digits of size 28$\times$28. We use built-in dataset from torchvision package (Contributors, ) and randomly sampled 9000 images for training 1000 for testing. Pixel values were normalized from $[0,255]$ to $[-1,1]$ via a linear transform, with labels remaining in the range $[0,9]$.

### 4.2 California Housing Price

In Figure 1, we plot the stable rank of $V$ for various values of $\mu_{V}$ and the singular values of $V$ when $\mu_{V}=10^{-4}$ and $\mu_{V}=1$. It is observed that the stable rank is close to one only for large weight decay when $\mu_{V}=0.1$ and $\mu_{V}=1$, while $V$ is almost full rank for small weight decay $\mu_{V}=10^{-4}$, $10^{-3}$ and $10^{-2}$. This supports our theoretical predictions in Theorem 2.4 and Theorem 2.6 that larger $\mu_{V}$ shrinks the distance from a rank one matrix $V$, and thus leads to a low-rank bias.
In Figure 2, we plot the stable rank $r_{\text{s}}(V)$, training MSE, and generalization error during the training. It again aligns with our theory that WD leads to the low-rank bias. Furthermore, although WD leads to larger training MSE, the best generalization error was achieved in the large WD regime $\mu_{V}=0.05$. We also found the stable rank and generalization error have weak dependence on the batch size $B$; see more details in the Appendix A.2.
To validate Assumption 2.5, which suggests that sufficient training results in small batch gradient norms,
we analyze all random batches of the final epoch. We compute the Frobenius norm of these random batch gradients and present the results as histograms in Figure 3.
It is shown that the Frobenius norm of most batch gradients is smaller than 2, with the largest norm being less than 12, regardless of the values of $\mu_{V}$ (ranging from $10^{-4}$ to 1).
In comparison, the Frobenius norm of a random Gaussian matrix with zero mean and variance $0.01$ of the same size is approximately $26$.

Figure: Figure 1: California Housing Prices. Left: Stable rank $r_{\text{s}}(V)$ versus $\mu_{V}$. Right: Singular values of $V$ for $\mu_{V}=0.0001$ and $1$. Here we fix batch size $B=16$.
Refer to caption: x1.png

Figure: Figure 2: California Housing Prices. Left: Stable rank. Middle: Training MSE. Right: Absolute value of generalization error. Here we fix the batch size $B=16$. The sharp transition in the generalization error happens when it changes sign.
Refer to caption: x3.png

Figure: (a) $\mu_{V}=0.0001$
Refer to caption: x6.png

### 4.3 MNIST Datasets

To examine our theoretical results, we consider the square loss function thereby reformulating this classification problem as a regression problem.
In Figure 4 we plot the stable rank $r_{\text{s}}(V)$ for various values of $\mu_{V}$ and the singular values of $V$ when $\mu_{V}=10^{-5}$ and $\mu_{V}=1$. In particular, the stable rank changes from $600$ to approximately $1$ when $\mu_{V}$ increases from $10^{-5}$ to $1$.
In Figure 5, we plot the stable rank $r_{\text{s}}(V)$, training MSE and generalization error during training. It is observed that the best generalization error is achieved when $\mu_{V}=1$.
In Figure 6, we present training, testing, and generalization accuracy for different values of $\mu_{V}$ and batch size $B$. We observe that generalization accuracy tends to zero as we increase $\mu_{V}$, while it exhibits weak dependence on the batch size $B$.
Additionally, Figure 7 presents histograms of the Frobenius norm of batch gradients in the final epoch. The norm of most gradients are bounded by $50$ with the largest norm being $92$. In comparison, the Frobenius norm of a random Gaussian matrix with zero mean and $0.01$ variance of the same size is approximately $507$.
All these results are consistent with results from regression task.

Figure: Figure 4: MNIST. Left: Stable rank versus $\mu_{V}$. Right: Singular values of $V$ for $\mu_{V}=10^{-5}$ and $1$. Here we fix batch size $B=64$.
Refer to caption: x12.png

Figure: Figure 5: MNIST. Left: Stable rank. Middle: Training MSE. Right: Absolute value of generalization error. Here we fix the batch size $B=64$. The sharp transition in the generalization error happens when it changes sign.
Refer to caption: x14.png

Figure: Figure 6: MNIST. Left: Accuracy when $B=64$. Right: Accuracy when $\mu_{V}=1$.
Refer to caption: x17.png

Figure: (a) $\mu_{V}=0.0001$
Refer to caption: x19.png

## 5 Discussion

Bounding generalization error is a fundamental problem in learning theory. Investigating the implicit bias of different algorithms helps address this problem by offering algorithm-specific estimates (Neyshabur et al., ). Among these biases, the low-rank bias is of particular interest due to its potential for better compression and generalization in many classification tasks. In this paper, we provide a mathematical proof that weight decay induces a low-rank bias in two-layer neural networks without relying on unrealistic assumptions. Our work explicitly shows how the gradient structure of neural networks, combined with weight decay, results in a low-rank weight matrix. We then further demonstrate that low-rank bias leads to better generalization bound.
An interesting follow-up question is whether our analysis can be extended to deeper ReLU neural networks and other variants used in computer vision tasks, such as convolutional neural networks.
Another natural extension of our results involves adaptive regularized SGD as in Equation 9, where the regularization strength varies across different batches.
In the proof of Theorem 2.6, we show that the distance between trained weight matrix $V^{\ast}$ and a rank-two matrix $\tilde{V}^{\ast}$ is bounded by $\frac{2B}{|g(x_{i_{1}},y_{i_{1}})-g(x_{i_{2}},y_{i_{2}})|}$. This suggests that a smaller generalization error bound can be achieved if the function $g$ can distinguish more effectively between different data points.

## Appendix A Appendix

### A.1 Proofs of Lemmas, propositions and theorems

###### Proof of Lemma 2.1 .

Denote $D_{ii}$ as the $i$-th diagonal entry of $D$ for $i=1,\ldots,m$, then we have

$$ $D_{ii}(x,V,b)=1\Leftrightarrow v_{i}^{\top}x+b_{i}>0\,,\quad\text{and}\quad D_ {ii}(x,V,b)=0\Leftrightarrow v_{i}^{\top}x+b_{i}\leq 0\,,$ $$

where $v_{i}^{\top}\in\mathbb{R}^{n}$ denotes the $i$-th row of $V$. For a fixed $\hat{v}_{i}$ such that $\hat{v}_{i}^{\top}x+b>0$, one can construct a neighborhood $\mathcal{V}_{i}$ around $\hat{v}_{i}$ such that $v_{i}^{\top}x+b_{i}>0$ for all $v_{i}\in\mathcal{V}_{i}$. This implies that $D_{ii}(x,V,b)=1$ on $\mathcal{V}_{i}$ and thus $\frac{\partial D(x,V,b)}{\partial v_{i}}=0$ at $\hat{v}_{i}$. Analogously, $\frac{\partial D(x,V,b)}{\partial v_{i}}=0$ at $\hat{v}_{i}$ as long as $\hat{v}_{i}^{\top}x+b<0$. Therefore, we have

$$ $\frac{\partial D(x,V,b)}{\partial v_{i}}=0$ $$

for any $v_{i}$ except a measure zero set $\mathcal{V}^{0}_{i}\coloneqq\{v_{i}\in\mathbb{R}^{n}:v_{i}^{\top}x+b_{i}=0\}$. We thus conclude that $\frac{\partial D(x,V,b)}{\partial V}=0$ for all $V$ except on $\mathcal{V}^{0}\coloneqq\{V\in\mathbb{R}^{m\times n}:v_{i}\in\mathcal{V}^{0}_{
i}\,,\text{for some }i\in[m]\}$, which is a measure zero set.
∎

###### Proof of Lemma 2.2 .

Apply product rule to Equation 6, we have

$$ $\frac{\partial\phi(x;\theta)}{\partial V}=U\left(\frac{\partial D(x,V,b)}{ \partial V}\right)(Vx+b)+xUD(x,V,b)\,.$ $$

The conclusion holds by applying Lemma 2.1 and noting $xU\in\mathbb{R}^{n\times m}$ is a rank one matrix.
∎

###### Proof of Theorem 2.4 .

We consider two batches $\mathcal{S}_{1}^{\prime},\mathcal{S}_{2}^{\prime}\subset\mathcal{S}$ with $|\mathcal{S}_{1}^{\prime}|=|\mathcal{S}_{1}^{\prime}|=B$ such that they only differ on one data pair, meaning that $\mathcal{S}_{1}^{\prime}/\mathcal{S}_{2}^{\prime}=\{(x_{i_{1}},y_{i_{1}})\}$ and $\mathcal{S}_{2}^{\prime}/\mathcal{S}_{1}^{\prime}=\{(x_{i_{2}},y_{i_{2}})\}$ for some $i_{1}\neq i_{2}\in\mathcal{S}$.
Then under Assumption 2.3, we must have

$$ $\frac{1}{B}\sum_{i\in\mathcal{S}_{j}^{\prime}}(\phi(x_{i},\theta^{\ast})-y_{i} )\frac{\partial\phi(x_{i},\theta^{\ast})}{\partial V}+g(x_{i},y_{i})V^{\ast}=0 \,,\quad j=1,2\,.$ (13) $$

Consequently, we have

$$ $(\phi(x_{i_{1}},\theta^{\ast})-y_{i_{1}})\frac{\partial\phi(x_{i_{1}},\theta^{ \ast})}{\partial V}+g(x_{i_{1}},y_{i_{1}})V^{\ast}=(\phi(x_{i_{2}},\theta^{ \ast})-y_{i_{2}})\frac{\partial\phi(x_{i_{2}},\theta^{\ast})}{\partial V}+g(x_ {i_{2}},y_{i_{2}})V^{\ast}\,.$ (14) $$

We further choose $i_{1}$ and $i_{2}$ such that $g(x_{i_{1}},y_{i_{1}})\neq g(x_{i_{2}},y_{i_{2}})$. Then we can solve for $V^{\ast}$,

$$ $V^{\ast}=-\frac{(\phi(x_{i_{1}},\theta^{\ast})-y_{i_{1}})\frac{\partial\phi(x_ {i_{1}},\theta^{\ast})}{\partial V}-(\phi(x_{i_{2}},\theta^{\ast})-y_{i_{2}}) \frac{\partial\phi(x_{i_{2}},\theta^{\ast})}{\partial V}}{g(x_{i_{1}},y_{i_{1} })-g(x_{i_{2}},y_{i_{2}})}\,.$ $$

We apply Lemma 2.2 and conclude that $V^{\ast}$ is at most rank two if $V^{\ast}$ is not in the measure zero set.

If our choice of data pairs $(x_{i_{1}},y_{i_{1}})$ and $(x_{i_{2}},y_{i_{2}})$ do not exist. This implies that $g(x_{i},y_{i})=c>0$ is identical for all $(x_{i},y_{i})\in\mathcal{S}$. Since Equation 14 still holds for our choice of $\mathcal{S}_{1}^{\prime}$ and $\mathcal{S}_{2}^{\prime}$ we have

$$ $(\phi(x_{i_{1}},\theta^{\ast})-y_{i_{1}})\frac{\partial\phi(x_{i_{1}},\theta^{ \ast})}{\partial V}=(\phi(x_{i_{2}},\theta^{\ast})-y_{i_{2}})\frac{\partial \phi(x_{i_{2}},\theta^{\ast})}{\partial V}\,.$ $$

Since $\mathcal{S}_{1}^{\prime}$ and $\mathcal{S}_{2}^{\prime}$ are arbitrary, and so are $i_{1}$ and $i_{2}$, we conclude that

$$ $(\phi(x_{i_{1}},\theta^{\ast})-y_{i_{1}})\frac{\partial\phi(x_{i_{1}},\theta^{ \ast})}{\partial V}=R$ $$

for some rank one matrix $R$.
Then Equation 13 implies that $V^{\ast}=-\frac{1}{c}R$ is a rank one matrix.
∎

###### Proof of Theorem 2.6 .

The proof is similar to that of Theorem 2.4. We choose two $\mathcal{S}^{\prime}_{1},\mathcal{S}^{\prime}_{2}\subset\mathcal{S}$ the same way. Then under Assumption 2.5, we must have

$$ $\|\frac{1}{B}\sum_{i\in\mathcal{S}_{j}^{\prime}}(\phi(x_{i},\theta^{\ast})-y_{ i})\frac{\partial\phi(x_{i},\theta^{\ast})}{\partial V}+g(x_{i},y_{i})V^{\ast} \|_{F}\leq\varepsilon\,,\quad j=1,2\,.$ (15) $$

Consequently, we have

$$ $\|(\phi(x_{i_{1}},\theta^{\ast})-y_{i_{1}})\frac{\partial\phi(x_{i_{1}},\theta ^{\ast})}{\partial V}+g(x_{i_{1}},y_{i_{1}})V^{\ast}-\left((\phi(x_{i_{2}}, \theta^{\ast})-y_{i_{2}})\frac{\partial\phi(x_{i_{2}},\theta^{\ast})}{\partial V }+g(x_{i_{2}},y_{i_{2}})V^{\ast}\right)\|_{F}\leq 2B\varepsilon\,.$ (16) $$

If there exists $i_{1}$ and $i_{2}$ such that $g(x_{i_{1}},y_{i_{1}})\neq g(x_{i_{2}},y_{i_{2}})$, then we can derive that

$$ $\|V^{\ast}-\tilde{V}^{\ast}\|_{F}\leq C\varepsilon\,,\text{ where }C=\frac{2B} {|g(x_{i_{1}},y_{i_{1}})-g(x_{i_{2}},y_{i_{2}})|}\,,$ $$

Here $\tilde{V}^{\ast}\coloneqq\frac{(\phi(x_{i_{1}},\theta^{\ast})-y_{i_{1}})\frac{
\partial\phi(x_{i_{1}},\theta^{\ast})}{\partial V}-(\phi(x_{i_{2}},\theta^{
\ast})-y_{i_{2}})\frac{\partial\phi(x_{i_{2}},\theta^{\ast})}{\partial V}}{g(x
_{i_{1}},y_{i_{1}})-g(x_{i_{2}},y_{i_{2}})}$ is a matrix with rank at most two due to Lemma 2.2.

If for any $i_{1}$ and $i_{2}$, $g(x_{i_{1}},y_{i_{1}})=g(x_{i_{2}},y_{i_{2}})$, then $g(x_{i},y_{i})=c$ for all $(x_{i},y_{i})\in\mathcal{S}$.
We first notice that Equation 15 can be simplified to the following.

$$ $\|cV^{\ast}+\frac{1}{B}\sum_{i\in\mathcal{S}^{\prime}}(\phi(x_{i},\theta^{\ast })-y_{i})\frac{\partial\phi(x_{i},\theta^{\ast})}{\partial V}\|_{F}\leq \varepsilon\,,\quad\text{for any }\mathcal{S}^{\prime}\subset\mathcal{S}\text{ with }|\mathcal{S}^{\prime}|=B\,.$ (17) $$

From Equation 16, we can derive that

$$ $\|(\phi(x_{i_{1}},\theta^{\ast})-y_{i_{1}})\frac{\partial\phi(x_{i_{1}},\theta ^{\ast})}{\partial V}-\left((\phi(x_{i_{2}},\theta^{\ast})-y_{i_{2}})\frac{ \partial\phi(x_{i_{2}},\theta^{\ast})}{\partial V}\right)\|_{F}\leq 2B \varepsilon\,.$ $$

Without loss of generality, we may define $\tilde{V}^{\ast}\coloneqq(\phi(x_{i_{1}},\theta^{\ast})-y_{i_{1}})\frac{
\partial\phi(x_{i_{1}},\theta^{\ast})}{\partial V}$ for any $i_{1}$, then the above equation implies that $(\phi(x_{i},\theta^{\ast})-y_{i})\frac{\partial\phi(x_{i},\theta^{\ast})}{
\partial V}$ is within an $2B\varepsilon$-ball of $\tilde{V}^{\ast}$ for all $i$, so is the average for any minibatch $\mathcal{S}^{\prime}\subset\mathcal{S}$ with $|\mathcal{S}|=B$. In other words,

$$ $\|\frac{1}{B}\sum_{i\in\mathcal{S}^{\prime}}(\phi(x_{i},\theta^{\ast})-y_{i}) \frac{\partial\phi(x_{i},\theta^{\ast})}{\partial V}-\tilde{V}^{\ast}\|_{F} \leq 2B\varepsilon\,.$ $$

Combined this equation with Equation 17, we apply triangle inequality and derive that

$$ $\|cV^{\ast}+\tilde{V}^{\ast}\|\leq(2B+1)\varepsilon\,,$ $$

which further implies that

$$ $\|V^{\ast}-(-\frac{1}{c}\tilde{V}^{\ast})\|_{F}\leq C\varepsilon\,,$ $$

where $C=\frac{2B}{c}>0$ and $(-\frac{1}{c}\tilde{V}^{\ast})$ is a rank one matrix.

∎

###### Proof of Proposition 3.6 .

Apply Theorem 3.3, we choose $\delta=4\mathcal{N}_{1}(\varepsilon/16,\mathcal{F},2N)\exp(-N\varepsilon^{2}/32)$, which implies that

$$ $\sqrt{\frac{32\ln(4/\delta)}{N}}\leq\varepsilon\leq\sqrt{\frac{32\ln(4/\delta) }{N}}+\sqrt{\frac{32\ln\mathcal{N}_{1}(\varepsilon/16,\mathcal{F},2N)}{N}}\,.$ $$

Apply Theorem 3.5, we have further that

$$ $\varepsilon\leq C\sqrt{\frac{\ln(1/\delta)}{N}}+C\sqrt{\frac{\ln\text{Pdim}( \mathcal{F})+\text{Pdim}(\mathcal{F})\ln(1/\varepsilon)}{N}}\,.$ $$

To eliminate the dependence on $\varepsilon$ of second term on the RHS, we plug in $\sqrt{\frac{32\ln(4/\delta)}{N}}\leq\varepsilon$ and obtain that

$$ $\varepsilon\leq C\sqrt{\frac{\ln(1/\delta)}{N}}+C\sqrt{\frac{\ln\text{Pdim}( \mathcal{F})+\text{Pdim}(\mathcal{F})\ln(N)}{N}}\,,$ $$

which completes the proof.
∎

###### Proof of Lemma 3.10 .

We can treat any $f\in\tilde{\mathcal{F}}(m,n,k)$ as a three layer neural network with width $m$ and the total number of parameters is $nk+mk+m=\mathcal{O}\left((m+n)k\right)$. The conclusion can be obtained by applying Theorem 7 in Bartlett et al. [].
∎

###### Proof of Theorem 3.11 .

This is a direct consequence of Lemma 3.10, Proposition 3.6 and Theorem 2.4.
∎

###### Proof of Lemma 2.1 .

###### Proof of Lemma 2.2 .

###### Proof of Theorem 2.4 .

###### Proof of Theorem 2.6 .

###### Proof of Proposition 3.6 .

###### Proof of Lemma 3.10 .

###### Proof of Theorem 3.11 .

### A.2 Numerical Experiments with different batch size

In Figure 8, we plot the stable rank $r_{\text{s}}(V)$ and generalization error for various batch size $B$ in the strong weight decay regime $\mu_{V}=1$.
Despite small variations, the stable rank and generalization error remain approximately constant across different batch sizes. Similar results are observed for MNIST datasets in Figure 9.

Figure: Figure 8: California Housing Prices. Left: Stable rank $r_{\text{s}}(V)$ versus batch size. Right: Absolute value of generalization error. Here we fix $\mu_{V}=1$.
Refer to caption: x25.png

Figure: Figure 9: MNIST. Left: Stable rank $r_{\text{s}}(V)$ versus batch size. Right: Absolute value of generalization error. Here we fix $\mu_{V}=1$.
Refer to caption: x27.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: Please address the issues from the “weaknesses” section.
soundness: 2
strengths: The paper considers an interesting and fundamental question for understanding generalization in overparameterized networks. It provides some nice insights on this question.
summary: The paper studies the implicit bias towards rank minimization and its implications on generalization. They show that mini-batch SGD with WD converges to low-rank solutions in two-layer networks (under certain assumptions). They also provide a generalization bound for low-rank networks, which might explain how rank minimization helps generalization.
weaknesses: - The authors claim that the assumption about the small norm of the gradients simultaneously for all batches might be too strong. But even in Theorem 2.6 they make this assumption approximately up to some epsilon, and it’s not clear whether the epsilon in this assumption is small. 
- The experiments compute this epsilon, but I don’t understand whether they imply that this epsilon is sufficiently small. For example, in Figure 3 they show that epsilon is at most 12. If we use $\mu_V=1$, then the bound from Theorem 2.6 becomes $2B \cdot 12 = 24B$. I don’t see in the paper what the value of $B$ is in this experiment (maybe it’s $16$ as in Figures 1 and 2?). Then, the authors compare it to the Frobenius norm of a random Gaussian matrix with variance $0.01$, which is $26$. I don’t understand why the authors chose $0.01$, and I don’t see why $26$ is large compared to $24B$. 
- Lemmas 2.1 and 2.2 hold for almost all matrices V. However, in the proofs of Theorems 2.4 and 2.6, they use this property for the convergence point $V^*$. The convergence point is not a random point, and hence the fact that almost all matrices satisfy a certain property does not imply that this property holds for the convergence point. Essentially, the authors assume that at convergence there are no hidden neurons with input $0$, but they don’t state this assumption explicitly. 
- Regarding the generalization bounds:
    - Corollary 3.9, as well as Theorem 3.11, depend on $L^2$. If we assume w.l.o.g. that the inputs are in the unit ball, then $L$ can be upper bounded by the product of the spectral norms of the layers. Then, the authors should compare their bound to well-known sample complexity bounds for neural networks. (see, e.g., Golowich, Rakhlin, Shamir, “Size-independent sample complexity of neural networks”, and the reference therein).
    - The generalization bound holds only for low-rank matrices, and not for the “approximately low-rank” case. That is, it does not work with Theorem 2.6. Moreover, the authors claimed in section 2.2.2 that the assumption in Theorem 2.4 might not be feasible, and hence Theorem 2.6 is the more realistic result.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: See the questions asked in the weaknesses above.
soundness: 3
strengths: The paper studies an important problem of understanding the regularization effect of weight decay and SGD optimization on learning with Neural Networks that can memorize the training data. This is an active area of research, and the paper has correctly identified a gap in the literature. The paper is also written clearly, and little background knowledge (beyond what is discussed) is needed to understand the results. I like the brief mention of the extension in the discussion, which hints at a data/batch-adaptive regularization strategy. Such a result/algorithm would be interesting.
summary: The paper studies the implicit bias towards low rankness in two-layered RELU networks trained with SGD and weight decay. In the regime where the network nearly interpolates the data/achieves stationarity across all batches, the authors show that the network's weight matrix is close to a rank one matrix in the Frobenius norm. Along with existing tools in learning theory, this implies an improved generalization bound for two-layered RELU networks. The authors also conducted experiments on small datasets to verify their predictions.
weaknesses: 1. **Theorems 2.4 and 2.6**: I have several concerns about the theorems, which I believe are the main results of this paper. Section 2.2.1 seems redundant, and does not serve any purpose, given Theorem 2.4 is only a special case of Theorem 2.6. Why do the authors choose to discuss Theorem 2.4? Beyond the setting where $\epsilon=0$, why does the upper bound 2.6 not gradual decline as we reduce $\mu_V$ down to zero? Intuitively $V^\star$ should be approximable by matrices of higher and higher rank, even when weight decay is driven to zero. Also, a minor but important point here is that both the theorems assume $\mu_V>0$. The upper bound in 2.6 does not look tight when batch sizes are high. Why would that be the case? Should the rank-minimizing effect decrease when we use GD instead of SGD? If so, why? Also, ensuring $\epsilon$ is small for all batch sizes is not impossible in the interpolation regime. In that regime how does the result compare to the more general (albeit without SGD optimization and weight decay) result of Ongie and Willett? Finally, the results seem too weak, as the authors do not show the effect of depth, and the proofs themselves are not very insightful about what will happen when depth increases. 

2. **Generalization results**: The calculations for the generalization error of low-rank neural networks seem incremental. They are likely not novel, as all the tools needed to derive these results already exist in the learning theory literature. Having said that, I can not confirm this, as I could not find the exact result in existing literature myself. More importantly, the results only seem to make sense in the interpolating regime. This is because Theorem 3.11 assumes Assumption 2.3. Is it not possible to give a result in terms of epsilon from Assumption 2.5 to understand the effect on approximate sationarity?  

3. **The experiments are too simplistic and unsurprising**: It seems that the empirical results in are on less diverse tasks than [[1]](https://arxiv.org/abs/2408.11804), which looks at more practical deep learning tasks with more complex architectures showing a low-rank bias across them. While [[1]](https://arxiv.org/abs/2408.11804) does not give a theoretical result, they provide a more comprehensive picture of the effect of weight decay. Moreover, the authors do not verify their assumptions or even plot their predicted upper bound in their experiments, which would have verified (or shown slackness of) their theorem.  

4. **Missing related works**: The paper is missing several related works about the effect of weight decay on neural network training. I encourage the authors to go through section 5 of [[1]](https://arxiv.org/abs/2408.11804).

---

confidence: 5
contribution: 1
initial_rating: 1
presentation: 2
questions: 1. The paper states that the results do not assume weight convergence, yet Assumptions 2.3 and 2.5 appear to imply it. Could you clarify how you define convergence in this context and why Assumptions 2.3 and 2.5 are not considered convergence assumptions?

2. Could you elaborate on how your approach to analyzing low-rank bias fundamentally differs from that in Galanti and Poggio (2021) and Xu et al. (2023)? Given the apparent similarities in proof technique and assumptions, what distinct contributions does this paper offer in Sec. 2.2 beyond those works that is not just weakening Assumption 2.3 by taking $\epsilon$ into account?

3. Could you provide practical examples or empirical evidence to support the feasibility of the term 
$\frac{C \epsilon}{\| V^* \|}$? Given that previous analyses (such as in Galanti et al. 2.3) can simply show that this term is typically large how do you justify its practicality in your bounds?

4. Given that neural network weight matrices in practice do not typically converge to rank-2 matrices, how do you envision the practical implications of your theoretical rank assumptions?
soundness: 2
strengths: I believe that the overall goal of the paper is interesting. The authors attempt to study the implicit complexity of neural networks trained with SGD and regularization and to connect it to the test performance of the resulting neural network. I think it's an interesting problem to explore how the rank of the weights contributes to the model's performance.
summary: The paper investigates how weight decay (WD) induces a low-rank bias in neural networks, leading to better generalization. The authors prove that, under sufficient training with weight decay and stochastic gradient descent (SGD), a two-layer ReLU neural network’s weight matrix approaches a rank-two structure. Empirical evidence supports this claim across regression and classification tasks, showing that weight decay encourages low-rank matrices even without conventional assumptions about data distribution or specific network architectures.

This low-rank bias contributes to reduced generalization error. The authors theoretically derive improved error bounds by leveraging the low-rank property and confirm these bounds with experiments on the California housing and MNIST datasets. Specifically, larger weight decay values reduce the rank of the weight matrix, decreasing generalization error. This study offers insights into the regularizing effect of SGD with weight decay, suggesting that low-rank bias is an implicit mechanism for improved generalization in neural networks. The paper's findings could extend to other architectures and optimization methods in deep learning.
weaknesses: I believe the paper is quite misleading. The authors claim they do not rely on any assumptions about the convergence of the weights, but this is inaccurate.

1. Table 1 compares their results on low-rank bias with previous work. However, the authors' findings are very similar to those of Galanti and Poggio (2021) and Xu et al. (2023). The analysis follows a nearly identical approach: it involves writing out the formula for the training step, assuming small gradient steps (as in Assumptions 2.3 and 2.5 of the current paper) across all batches, examining the difference between gradient steps for two batches that differ by one sample, and then using this gap to suggest that the learned weights are close to a low-rank weight matrix. The authors never acknowledge that this proof technique was introduced in Galanti and Poggio (2021) and Xu et al. (2023). The main differences across these works are that Xu et al. (2023) used a slightly different optimization setting (e.g., training with SGD, WD and WN) and an assumption similar to 2.3 to justify that the weights converge to a low-rank matrix, while Galanti and Poggio (2021) relaxed the assumption, considering a scenario similar to 2.5. Both previous papers offer results that hold for a wider range of architectures than the current one.

2. The authors claim their results hold without assuming the convergence of the weights, but this is incorrect. Both Assumptions 2.3 and 2.5 are essentially convergence assumptions. While Assumption 2.3 makes this more explicit, Assumption 2.5 implies the same concept. Furthermore, the authors do not justify the practicality of the quantity $ C\epsilon/\|V^*\|$, which is, unfortunately, unrealistic. To illustrate, following the analysis in Galanti et al. 2023 [https://arxiv.org/abs/2206.05794], we can consider the last equation in their paper's page 10 and easily derive that $\|V^*\| \approx \|\frac{\partial L}{\partial V}\|/\mu_V = \epsilon/\mu_V$. This implies $C\epsilon/\|V^*\| \approx 2B$, which is a constant larger than 1. Therefore, the bound in Theorem 2.6 is trivially true.

This issue also applies to Galanti and Poggio (2021), which was addressed in their recent paper Galanti et al. 2023 [https://arxiv.org/abs/2206.05794]. It appears that the authors of the current paper overlooked this work, which similarly aims to demonstrate that SGD + WD induces a low-rank bias in deep learning. In this paper, the authors show that the rank of weight matrices is bounded by a function dependent on batch size, regularization parameter, and learning rate. These results hold for most architectures of interest without requiring assumptions about the data. By unrolling the training process of SGD + WD, they avoid assumptions like 2.3 and 2.5, making only the weaker assumption that the norms of the weights converge.

3. The result in 3.11 is also not very new. The result is based on a VC-style generalization bound applied to a neural network whose weight matrices are low rank. When combined with UV decomposition, we can represent the network as a network with 3 layers where each one of them includes a small number of parameters. I admit that I don't recall a specific paper that explicitly runs this derivation, but it's fairly trivial if I am being honest.

4. The narrative presented in this paper is somewhat simplistic. By relying on a very strong assumption (2.3 or 2.5), it derives an unrealistic result about matrix rank, which then leads to a very tight generalization bound for these specific solutions. However, in practice, the weight matrices of neural networks do not converge to rank-2 matrices. Furthermore, the rank of these matrices does not fully explain neural networks' success and only has a marginal effect on the performance (we can easily train neural networks that generalize well on common datasets without weight decay). A more compelling direction would be to explore the "in-between" scenario: understanding the actual behavior of rank in practice without assuming strong convergence assumptions (similar to what is done in Galanti et a. 2023) and its marginal impact on performance.

---

confidence: 5
contribution: 1
initial_rating: 1
presentation: 1
questions: None.
soundness: 1
strengths: **None**
summary: TLDR: this paper contains **zero** contribution -- all of the results presented are either trivial or wrong.
weaknesses: *This paper beyond any redemption!* I will just point out a few issues that I am most frustrated with.


The introduction is quite a wild ride. I feel that someone without intimate knowledge of the area will immediately be confused. Notably, the term "implicit regularization/bias" was introduced without a proper explanation. And the problem as defined in eq. (1-3) is lacking in context (I think a few sentences is enough, but jumping directly into equation certainly looks very awkward).


Dubious choice of references:
1. Some well-known work in the space were not mentioned. For example, [1, 2] are two of the earliest and most-cited works in the analysis of GD's implicit bias, but these are not referenced in this paper.
2. On line 50, NTK is definitely not the only reason why optimization error goes to zero as model size grows. There are plenty of other works covering this topic, e.g. [3, 4]. The authors should consider attributing a more diverse set of works.
3. On line 58, the author claimed that "these implicit biases are not theoretically understood yet" while the references on the previous line all provide theoretical contributions to this question.
4. Table 1 needs to be better explained. I struggle to understand what it exactly means.


**The proof of Theorem 2.4 is broken**. First, line 665 does not work when the function $g$ is constant, which corresponds to the standard practice of a fixed, data-independent weight decay factor. Secondly, assumption 2.3 contradicts the claim the results do not rely on the convergence of weights.


**The entirety of Section 3 contain no new results at all** All of the stated results are either direct citations or trivial applications of said existing results. The proofs for this section do not even span half of a page, yet this is advertised as a major technical section.


The experiments in Section 4 is very lacking. Results on MNIST have very little implication on generalization since plenty of ''small'' models work well for this. This is especially damning since there is very little theoretical contributions that need numerical verification.


In summary, this paper contains no meaningful contribution and authors fail to demonstrate a deep understanding of this topic. **I recommend rejection without a shred of doubt.**


[1] Soudry D, Hoffer E, Nacson MS, Gunasekar S, Srebro N. The implicit bias of gradient descent on separable data. Journal of Machine Learning Research. 2018.

[2] Ji Z, Telgarsky M. The implicit bias of gradient descent on nonseparable data. In Conference on learning theory. 2019.

[3] Belkin M, Hsu D, Ma S, Mandal S. Reconciling modern machine-learning practice and the classical bias–variance trade-off. Proceedings of the National Academy of Sciences. 2019.

[4] Du S, Lee J, Li H, Wang L, Zhai X. Gradient descent finds global minima of deep neural networks. In International conference on machine learning. 2019.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: I would personally appreciate some clarification on the weaknesses. Moreover, I would appreciate discussing with the authors the following, hoping there was no misunderstanding:

1) It is not clear to me what is the interest of studying the case for which $\mu_V = \frac{1}{B} \sum_{j \in S'} g(x_j,x_j)$ ? I understand that for a constant $g$ this leads to the typical case, but what is the gain in considering this more complicated case in the analysis? In the proof this leads to rank-2, but I cannot think of an example in which this kind of choice for the regularization parameter is done in applications.
2) The analysis doesn't rule out the possibility of $V^*$ being in $\mathcal V_{0,i}$ for some $i$, so to me it is essentially just the analysis of a deep-linear model with $L=2$ layers, which has been done already multiple times in literature.
soundness: 2
strengths: The paper is overall well written. The claims, proofs and motivation of the work are clearly presented. Moreover, the problem of implicit biases and their implications in generalization theory is a very significant topic.
summary: The authors prove an implicit bias toward low-rank (one or two) matrices in training two-layer RELU neural networks with stochastic gradient descent and weight decay. By exploiting this implicit bias towards the set of rank-two matrices, the authors were able to derive tighter generalization bounds.
weaknesses: The proposed work doesn't study any dynamical conditions, thus to kind of convergence to the claimed points was presented.
I would say the results are more on the line of studying fixed points of the SGD dynamics. While I can intuitively believe assumption 2.5, the authors don't present any theoretical guarantee or even empirical evidence suggesting that the dynamic in practice drives onto regions in parameter space for which that assumption is satisfied.
Given that the second part of the paper relies on the first one as a motivation, I believe also that the soundness of the results in the second part decreases (please see above and weaknesses).
Also, the experimental setting is too restricted.