SGD and Weight Decay Secretly Minimize the Rank of Your Neural Network
Extracted source text · arXiv 2206.05794
\documentclass{article} % For LaTeX2e
\usepackage{iclr2023_conference,times}
% Optional math commands from https://github.com/goodfeli/dlbook_notation.
\input{math_commands.tex}
\usepackage{hyperref}
\usepackage{url}
\usepackage{amsmath,amssymb,amsthm,bbm}
\usepackage{thm-restate,enumitem}
\usepackage{tabularx}
\usepackage{rotating}
\newtheorem{theorem}{Theorem}
\newtheorem{lemma}{Lemma}
\newtheorem{proposition}{Proposition}
\newtheorem{corollary}{Corollary}
\newtheorem{definition}{Definition}
\newtheorem{remark}{Remark}
\newcommand{\expnumber}[2]{{#1}\mathrm{e}{#2}}
\newcommand{\cC}{\mathcal{C}}
\newcommand{\cD}{\mathcal{D}}
\newcommand{\cF}{\mathcal{F}}
\newcommand{\cG}{\mathcal{G}}
\newcommand{\cH}{\mathcal{H}}
\newcommand{\cN}{\mathcal{N}}
\newcommand{\cP}{\mathcal{P}}
\newcommand{\cU}{\mathcal{U}}
\newcommand{\cX}{\mathcal{X}}
\newcommand{\cY}{\mathcal{Y}}
\renewcommand{\P}{\mathbb{P}}
\newcommand{\B}{\mathbb{B}}
\newcommand{\bI}{\mathbb{I}}
\newcommand{\EE}[1]{\E\left[ #1\right]}
\newcommand{\PP}[1]{\P\left[ #1\right]}
\newcommand{\I}[1]{\bI\left[ #1\right]}
\newcommand{\tiX}{\tilde{X}}
\newcommand{\tiS}{\tilde{S}}
\newcommand{\ty}{\tilde{y}}
\newcommand{\tcG}{\tilde{\cG}}
\newcommand{\tcP}{\tilde{\cP}}
\newcommand{\tcU}{\tilde{\cU}}
\newcommand{\tf}{\tilde{f}}
\newcommand{\tig}{\tilde{g}}
\newcommand{\tih}{\tilde{h}}
\newcommand{\clD}{\cD_{\cC}}
\newcommand{\one}{\mathbbm{1}}
\newcommand{\rank}{\textnormal{rank}}
\newcommand{\vect}{\textnormal{vec}}
\newcommand{\threepartdef}[6]
{
\left\{
\begin{array}{lll}
#1 & \mbox{if } #2 \\
#3 & \mbox{if } #4 \\
#5 & \mbox{if } #6
\end{array}
\right.
}
\newcommand{\twopartdef}[4]
{
\left\{
\begin{array}{lll}
#1 & \mbox{if } #2 \\
#3 & \mbox{if } #4 \\
\end{array}
\right.
}
\title{Weight Decay and SGD Provably Induce Low-Rank Bias in Neural Networks}
% Authors must not appear in the submitted version. They should be hidden
% as long as the \iclrfinalcopy macro remains commented out below.
% Non-anonymous submissions will be rejected without review.
\author{Antiquus S.~Hippocampus, Natalia Cerebro \& Amelie P. Amygdale \thanks{ Use footnote for providing further information
about author (webpage, alternative address)---\emph{not} for acknowledging
funding agencies. Funding acknowledgements go at the end of the paper.} \\
Department of Computer Science\\
Cranberry-Lemon University\\
Pittsburgh, PA 15213, USA \\
\texttt{\{hippo,brain,jen\}@cs.cranberry-lemon.edu} \\
\And
Ji Q. Ren \& Yevgeny LeNet \\
Department of Computational Neuroscience \\
University of the Witwatersrand \\
Joburg, South Africa \\
\texttt{\{robot,net\}@wits.ac.za} \\
\AND
Coauthor \\
Affiliation \\
Address \\
\texttt{email}
}
% The \author macro works with any number of authors. There are two commands
% used to separate the names and addresses of multiple authors: \And and \AND.
%
% Using \And between authors leaves it to \LaTeX{} to determine where to break
% the lines. Using \AND forces a linebreak at that point. So, if \LaTeX{}
% puts 3 of 4 authors names on the first line, and the last on the second
% line, try using \AND instead of \And before the third author name.
\newcommand{\fix}{\marginpar{FIX}}
\newcommand{\new}{\marginpar{NEW}}
%\iclrfinalcopy % Uncomment for camera-ready version, but NOT for submission.
\begin{document}
\maketitle
\begin{abstract}
We analyze deep ReLU neural networks trained with mini-batch Stochastic Gradient Descent (SGD) and weight decay. We show, both theoretically and empirically, that when training a neural network using SGD with weight decay and small batch size, the resulting weight matrices are expected to be of small rank. Our analysis relies on a minimal set of assumptions and the neural networks may be arbitrarily wide or deep, and may include residual connections, as well as batch normalization layers and convolutional layers. Furthermore, we study the source of SGD noise and prove that when training with weight decay, the only solutions of SGD at convergence are zero functions.
\end{abstract}
\section{Introduction}\label{sec:intro}
Over the past few years, neural networks have challenged machine learning theory with several mysteries. We consider two of them here that are, to our knowledge, so far unsolved. The first is about the mathematical reasons for a bias in Stochastic Gradient Descent (SGD) towards neural networks with low-rank weight matrices, a fact which has been empirically observed (e.g., \cite{DBLP:journals/corr/abs-2201-12760,DBLP:journals/corr/abs-2103-10427}). The second is regarding the role and origin of SGD noise (e.g., \cite{DBLP:journals/corr/abs-2110-06914}). The origin of the SGD noise represents a puzzling and rather neglected question: {\em why does SGD, under certain conditions, never stop making updates?} In this paper, we report a simple observation that solves both mysteries and, somewhat surprisingly, shows that they have the same mathematical root. The observation is that $\frac{\partial f_W}{\partial W^i}$ is a matrix of rank $\leq 1$, where $f_W$ is a deep ReLU network with a scalar output and weight matrices $\{W^i\}_{i \in I}$ ($I$ is some finite set of indices)\footnote{When $f_W$ outputs a $k$-dimensional vector, each output component of the tensor corresponding to $\frac{\partial f}{\partial W^i}$ is a matrix of rank at most $1$.}.
{\color{red}Why is low-rank/SGD noise important?}
\subsection{Related Work}\label{sec:related_work}
% Stochastic gradient descent (SGD) is one of the standard workhorses for optimizing deep models~\citep{bottou-91c}. Though initially proposed to remedy the computational bottleneck of gradient descent (GD), recent studies suggest SGD also induces a crucial implicit regularization, which prevents the overparameterized models from converging to the minima that cannot generalize well~\cite{DBLP:journals/corr/ZhangBHRV16,https://doi.org/10.48550/arxiv.1711.04623,https://doi.org/10.48550/arxiv.1609.04836,zhu}. Empirical studies suggest that (i) SGD outperforms GD~\cite{zhu}, (ii) small batch SGD generalizes better than large batch SGD~\citep{NIPS2017_a5e0ff62,https://doi.org/10.48550/arxiv.1609.04836}, and (iii) gradient descent with additional noise cannot compete with SGD~\cite{zhu}. However, despite many efforts, the potentially important, yet implicit, regularization induced by SGD noise has never been fully understood. Furthermore, we are unaware of any study of the noise in the output of the network during SGD training, a noise which seems to be present even at very long training times, suggesting that true convergence to a single set of weight matrices matrices for all data points is never reached.
A prominent thread in the recent literature is devoted to the study of the implicit bias in neural networks towards rank minimization. Most of the interest was on the matrix factorization problem, which corresponds to training a depth-2 linear neural network with multiple outputs w.r.t. the square loss. As described by~\cite{DBLP:journals/corr/abs-2201-12760},~\cite{https://doi.org/10.48550/arxiv.1705.09280} initially conjectured that the implicit regularization in matrix factorization can be characterized by the nuclear norm of the corresponding linear predictor. This conjecture was formally refuted by~\cite{DBLP:journals/corr/abs-2012-09839}. \cite{DBLP:journals/corr/abs-2005-06398} conjectured that the implicit regularization in matrix factorization can be explained by rank minimization, and also hypothesized that some notion of rank minimization may be key to explaining generalization in deep learning. \cite{DBLP:journals/corr/abs-2012-09839} established evidence that the implicit regularization in matrix factorization is a heuristic for rank minimization. Beyond factorization problems, \cite{DBLP:journals/corr/abs-2006-06657} showed that in linear networks of output dimension 1, gradient flow (GF) w.r.t. exponentially-tailed classification losses converges to networks where the weight matrix of every layer is of rank 1.
However, with nonlinear neural networks, things are less clear. Empirically, a series of papers~\citep{NIPS2014_2afe4567,10.5555/3294771.3294853,s21165599,8099498,pmlr-v80-arora18b} showed that replacing the weight matrices by low-rank approximations results in only a small drop in accuracy. This suggests that the weight matrices in practice may be close to low-rank matrices. However, whether they provably behave this way remains unclear. \cite{DBLP:journals/corr/abs-2201-12760} show that for nonlinear ReLU networks, GF does not minimize rank. They also argue that ReLU networks of sufficient depth can have low-rank solutions under $\ell_2$ norm minimization. This interesting result, however, applies to layers added to a network that already solves the problem and may not have any low-rank bias. It is not directly related to the mechanism described in this paper, which applies to all layers in the network but only in the presence of regularization and SGD, unlike the phenomenon described in \cite{DBLP:journals/corr/abs-2201-12760}.
\subsection{Contributions}
We provide a mathematical analysis of the origin of SGD noise and the implicit rank-minimization of SGD. We analyze deep ReLU networks trained with mini-batch stochastic gradient descent and $L_2$ regularization (i.e., weight decay). Our analysis is fairly generic: the neural networks may include fully-connected layers, residual connections, convolutional layers, and batch normalization layers. In addition, the loss function is only assumed to be differentiable. The central contributions are:
\begin{itemize}[leftmargin=0.5cm]
\item In Thm.~\ref{thm:batches} we prove that mini-batch SGD has an implicit bias towards networks with low-rank matrices. This theorem connects batch size, weight decay, optimization and the rank of the learned matrices.
\item In particular, in Sec.~\ref{sec:lowrank}, we theoretically show that the batch size and weight decay have a crucial effect on the rank of the learned matrices. In Sec.~\ref{sec:rank_exp} we provide an extensive empirical study of the relationship between the rank and certain hyperparameters.
\item In Props.~\ref{prop:propA}-\ref{prop:propB}, we study the source of SGD noise. The result shows that SGD noise must be always present, even asymptotically, regardless of the batch size, as long as we incorporate weight decay. Thus, SGD noise is a generic property; this also means that, generically, there is never convergence of SGD. These predictions are later validated empirically in Sec.~\ref{sec:sgd_noise_exp}.
\end{itemize}
\section{Problem Setup}\label{sec:setup}
We consider the problem of training a model for a standard classification or regression problem. Formally, the target task is defined by a distribution $P$ over samples $(x,y)\in \cX\times \cY$, where $\cX \subset \R^{c_{1} \times h_1 \times w_1}$ is the instance space (e.g., images), and $\cY \subset \R^k$ is a label space. %To simplify the presentation, we use one-hot encoding for the label space, that is, the labels are represented by the unit vectors in $\R^C$, and $\cY_C=\{e_c: c=1,\ldots,C\}$ where $e_c \in \R^C$ is the $c$th standard unit vector in $\R^C$; with a slight abuse of notation, sometimes we will also write $y=c$ instead of $y=e_c$.
%For a pair $(x,y)$ distributed by $P$, we denote by $P_c$ the class conditional distribution of $x$ given $y=c$ (i.e., $P_c(\boldsymbol{\cdot})=\P[x \in \boldsymbol{\cdot} \mid y=c]$).
We consider a parametric model $\cF \subset \{f':\cX\to \R^k\}$, where each function $f_W \in \cF$ is specified by a vector of parameters $W \in \R^{N}$. A function $f_W \in \cF$ assigns a prediction to an input point $x \in \cX$, and its performance on the distribution $P$ is measured by the {\em expected risk}
\begin{equation}
\label{eq:loss}
L_P(f_W)~:=~\E_{(x,y)\sim P}[\ell(f_W(x),y)],
\end{equation}
where $\ell:\R^k \times \cY \to [0,\infty)$ is a non-negative, differentiable, loss function (e.g., MSE or cross-entropy losses). For simplicity, throughout the paper, we focus on the case where $k=1$.
Since we do not have direct access to the full population distribution $P$, the goal is to learn a predictor, $f_W$, from some training dataset $S = \{(x_i,y_i)\}_{i=1}^m$ of independent and identically distributed (i.i.d.) samples drawn from $P$ along with regularization to control the complexity of the learned model. Namely, we intend to minimize the {\em regularized empirical risk}
\begin{equation}
\label{eq:loss_emp}
L^{\lambda}_S(f_W)~:=~\frac{1}{m}\sum^{m}_{i=1}\ell(f_W(x_i),y_i) + \lambda \|W\|^2_2,
\end{equation}
where $\lambda > 0$ is predefined hyperparameter, controlling the amount of regularization. In order to minimize this objective, we typically use mini-batch SGD, which is described below in more detail.
% Finally, the performance of the trained model is evaluated using the train and test accuracy rates, which are computed as follows:
% $err_S(f_W) = \sum^{m}_{i=1}\ell(f_W(x_i),y_i)$ and $err_P(f_W) = \E_{(x,y)\sim P}[\ell(f_W(x),y)]$.
%We assume that the $m_0$ samples are taken directly from the classes $P_c$, following the observation in~\cite{Papyan24652} that the data should be balanced in order to achieve a high degree of neural collapse.
{\bf Optimization.\enspace} In this work, we minimize the regularized empirical risk $L^{\lambda}_{S}(f_{W})$ by applying stochastic gradient descent (SGD) for a certain number of iterations $T$. Formally, we initialize $W_0$ using a standard initialization procedure and iteratively update $W_t$. At each iteration, we sample a subset $\tiS = \{(x_{ij},y_{ij})\}^{B}_{j=1} \subset S$ uniformly at random and update $W_{t+1} \leftarrow W_t - \eta_t \cdot \nabla_W L^{\lambda}_{\tiS}(f_{W_t})$, where $\eta_t > 0$ is our learning rate at iteration $t$. The algorithm returns $W_T$.
\paragraph{Notation.} Throughout the paper, we use the following notations. For an integer $k\geq 1$, $[k]=\{1,\ldots,k\}$. For any real vector $z$, $\|z\|$ denotes its Euclidean norm. For a finite set $A$, we denote by $U[A]$ the uniform distribution over $A$. We denote by $\bI:\{\textnormal{True},\textnormal{False}\} \to \{0,1\}$ the indicator function. For a given tensor $x \in \R^{c \times h\times w}$, we denote $\vect(x) = (x_{1,1,1},\dots,x_{1,1,w},x_{1,2,1},\dots,x_{1,2,w}, \dots ,x_{c \times h \times w})$ the vectorized form of $x$. We define tensor slicing as follows $x_{a:b} = (x_a,\dots,x_b)$.
\subsection{Architectures and Training}
In this work, the function $f_W$ represents a neural network, consisting of a set of layers of weights interlaced with activation units. We employ a fairly generic definition of a neural network, that includes convolutional layers and residual connections. %To include normalization layers, we extend $f_W$'s input to be $(x;\tiX)$, where $x\in \R^{n}$ and $\tiX = \{x_{ij}\}^{B}_{j=1} \subset \R^{n}$ is a batch of samples (see details below). For simplicity, we denote by $f_W(x)$ a network without normalization layers. When $\tiX$ is clear from context, we simply write $\tf_W(x) := f_W(x;\tiX)$.
{\bf Network architecture.\enspace} Formally, $f_W$ is as a directed acyclic graph (DAG) $G=(V,E)$, where $V= \{v_1,\dots,v_L\}$ consists of the various layers of the network and each edge $e_{ij}=(v_i,v_j) \in E$ specifies a connection between two layers. Each layer is a function $v_{i}:\R^{c_{1} \times h_1 \times w_1} \to \R^{c_{i} \times h_i \times w_i}$ and each connection holds a transformation $C^{ij} : \R^{c_{j} \times h_j \times w_j} \to \R^{c_i \times h_i \times w_i}$. The layers are divided into three categories: {\bf (i)} the input layer $v_1$, {\bf (ii)} the output layer $v_L$ and {\bf (iii)} intermediate layers. In this setting, we do not have connections directed towards the input layer nor connections coming out of the output layer (i.e., $\forall~i \in [L]:~(v_L,v_i), (v_0,v_i) \notin E$). Given an input $x$ and batch $\tiX$, the output of a given layer $v_i$ is evaluated as follows $v_i(x) := \sigma(\sum_{j \in \textnormal{pred}(i)} C^{ij}(v_{j}(x)))$, except for the output layer $v_L$ that computes $f_W(x) := v_L(x) := \sum_{j \in \textnormal{pred}(L)}C^{Lj} (v_{j}(x))$, where $\textnormal{pred}(i)$ is the set of indices $j$, such that $(v_i,v_j) \in E$. Here, $\sigma$ is an element-wise activation function (e.g., ReLU, Leaky ReLU, sigmoid).
In this setting, the transformations $C^{ij}$ could be trainable, $(v_i,v_j) \in E_T$ (e.g., convolutional layers), or constant, $(v_i,v_j) \notin E_T$ (e.g., residual connections or a pooling layer). Throughout the analysis, we consider paths within the graph $G$, denoted by $\pi = (\pi_0,\dots,\pi_T)$, where $\pi_0=1$, $\pi_T = L$ and for all $i=0,\dots,T-1:~ (v_{\pi_{i}},v_{\pi_{i+1}}) \in E$.
{\bf Convolutional layers.\enspace} The networks include convolutional layers~\cite{6795724,10.5555/646469.691875} (see also~\cite{Goodfellow-et-al-2016}) $C^{ij}:\R^{c_i \times h_i \times w_i} \to \R^{c_j \times h_j \times w_j}$ with kernel sizes $(k_1,k_2)$, padding $p$ and stride $s$. Each convolutional layer is parameterized by a tensor $Z \in \R^{c_{i}\times c_j \times k_1 \times k_2}$ and computes the following tensor as output
\begin{equation}\label{eq:conv}
y_{c,i,j} = \sum^{c_{1}}_{c'=1} \vect(Z_{c',~:~,~:~})^{\top} \cdot \vect(\textnormal{Pad}_p(x)_{c',is~:~(i+1)s+k_1 , js~:~(j+1)s+k_2}),
\end{equation}
where $(c,i,j) \in [c_2]\times [h_2] \times [w_2]$. Here, $\textnormal{Pad}_p$ takes a tensor $x \in \R^{c \times h \times w}$ and returns a new tensor $x' \in \R^{c \times (h+2p) \times (w+2p)}$, where the first and last $p$ rows and columns of each $x'_{c,~:~,~:~}$ are zeros and the middle $1 \times h \times w$ matrix is equal to $x_{c,~:~,~:~}$. The formula for the output dimensions are $h_2 = \left(\lfloor(h_1 - k_1 +2p)\rfloor/s+1\right)$ and $w_2 = \left(\lfloor(w_1 - k_2 +2p)\rfloor/s+1\right)$. For a given convolutional layer $C^{ij}$ with weights $Z^{ij}$, we define a matrix $W^{ij}$ that computes $\forall~ x \in \R^{c_i \times h_i \times w_i}:~ W^{ij} \vect(x) = \vect(C^{ij}(x))$. This is guaranteed since both the padding and convolution operations of the layer are linear {\color{red}explain}. We also consider $V^{ij} \in \R^{c_{i} \times c_{j}k_1k_2}$ which is a matrix whose $c$'th row is the vectorized filter $V^{ij}_c := \vect(Z^{ij}_{c,~:~,~:~})$.
As a special case, the networks may also include fully-connected layers: $C:\R^{c_1 \times 1 \times 1} \times \R^{c_2 \times 1 \times 1}$ that take an input $x \times \R^{c_1 \times 1\times 1}$ and apply a linear transformation $\vect(C(x)) = W \vect(x)$, where $W \in \R^{c_2 \times c_1}$ is a trainable weight matrix. We note that fully-connected layers can be represented as $1\times 1$ convolutional layers, i.e., convolutional layers over $\R^{c_1 \times 1 \times 1} \to \R^{c_2 \times 1 \times 1}$ with $k_1=k_2$, $p=0$ and $s=1$.
{\bf Rearrangement layers.\enspace} A rearrangement layer $C:\R^{c_1 \times h_1 \times w_1} \to \R^{c_2 \times h_2 \times w_2}$ takes an input vector $x \in \R^{c_1 \times h_1 \times w_1}$ and `rearranges' its coordinates in a different shape and according to a certain permutation. Formally, it returns a vector $(x_{\pi(i)})_{i \in [c_2]\times [h_2] \times [w_2]}$, where $\pi:[c_2]\times [h_2] \times [w_2] \to [c_1]\times [h_1] \times [w_1]$ is an invertible function.
{\bf Pooling Layers.\enspace} A pooling layer~\citep{Zhou1988ComputationOO} (see~\citep{GoodBengCour16}) $C$ with kernel dimensions $(k_1,k_2)$ stride $s$ and padding $p$ takes an input $c_1 \times h_1 \times w_1$ and computes an output with $c_1$ channels, and dimensions $h_2 = \left(\lfloor(h_1 - k_1 +2p)\rfloor/s+1\right)$ and $w_2 = \left(\lfloor(w_1 - k_2 +2p)\rfloor/s+1\right)$. In this paper we consider average pooling and max-pooling layers. Each pooling layer computes the following tensor as output
\begin{equation}\label{eq:pool}
y_{c,i,j} = \max(\textnormal{Pad}_p(x)_{c,is~:~(i+1)s+k_1 , js~:~(j+1)s+k_2}),
\end{equation}
\section{Theoretical Results}\label{sec:theory}
In this section, we describe our main theoretical results. In our analysis we consider a regime of training a neural network with mini-batch SGD that takes place at the near-convergence.
To formally study SGD convergence, we employ the notion of convergence in mean of random variables. Namely, for a given sequence of random variables $\{W_t\}^{\infty}_{t=1} \subset \R^{N}$ starting from a certain vector $W_1 \in \R^{N}$ (constant), we say that it converges in mean if we have
\begin{equation}
\exists~ W^* \in \R^{N}:~\lim_{t \to \infty} \E[\|W_{t} - W^*\|] ~=~ 0.
\end{equation}
As a consequence of this definition, we have
\begin{equation}
\lim_{t\to \infty}\E[\|W_{t+1} - W_t\|] ~\leq~ \lim_{t\to \infty}\E[\|W_{t+1} - W^*\|] + \lim_{t\to \infty}\E[\|W_{t} - W^*\|] ~=~ 0.
\end{equation}
Therefore, convergence is possible only when the expected size of each step tends to zero. In particular, when training the network using mini-batch SGD along with weight decay, we have
\begin{equation}
\lim_{t \to \infty} \E_{\tiS} \| \nabla L^{\lambda}_{\tiS}(f_{W^t})\| = 0,
\end{equation}
since $\E[\|W^{t+1} - W^{t}\|] = \eta \cdot \E_{\tiS}[\|\nabla L^{\lambda}_{\tiS}(f_{W^t})\|]$.
To analyze the step taken in each step, we write it explicitly:
\begin{equation}
\begin{aligned}
\nabla_{V^{pq}} L^{\lambda}_{\tiS}(f_W) ~=~ \frac{1}{B} \sum_{(x,y)\in \tiS} \frac{\partial \ell(f_W(x),y)}{\partial f_W(x)} \cdot \frac{\partial f_W(x)}{\partial V^{pq}} + 2\lambda V^{pq}.
\end{aligned}
\end{equation}
In Sec.~\ref{sec:lowrank} we show that near-convergence, mini-batch SGD implicitly learns neural networks with low-rank matrices. In Sec.~\ref{sec:convergence}, we prove that SGD noise is inevitable when applying weight decay.
\subsection{Low-Rank Bias in Neural Networks}\label{sec:lowrank}
We begin our theoretical analysis with the simple observation that the rank of the gradients of the network's output with respect to the parameters matrix of a given convolutional layer $C^{ij}$ is upper bounded by the number of input patches $N^{ij}$.
\begin{restatable}{lemma}{ranklemma}\label{lem:rank}
Let $f_W$ be a convolutional neural network and let $C^{ij}$ be a given convolutional layer within $f_W$ with parameters matrix $V^{ij}$. Then, $\rank\left(\nabla_{V^{ij}} f_W(x)\right)\leq N^{ij}$.
\end{restatable}
We note that a fully-connected layer $M:\R^{c} \to \R^{c'}$ associated with a matrix $W \in \R^{c' \times c}$ can be represented as a convolutional layer $C:\R^{c \times 1 \times 1}\to \R^{c' \times 1 \times 1}$ whose parameters tensor $Z \in \R^{c' \times c \times 1 \times 1}$ satisfies $Z_{a,b,1,1}=W_{a,b}$ for all $(a,b)\in [c']\times [c]$. Since the number of patches in this case is $1$, the rank of the gradient associated with $V$ is at most $1$.
{\color{red}Some text}
\begin{restatable}{theorem}{batches}\label{thm:batches}
Let $\|\cdot \|$ be any matrix norm and $\ell$ any differentiable loss function. Let $f_W(x)$ be a ReLU neural network and $C^{pq}$ be a given matrix within $f_W$ and let $B^* \in [m]$. Then,
\begin{equation*}
\min_{\substack{V \in R^{d_p \times d_q}:\\ \rank(V) \leq kN^{pq}B^*}} \|V^{pq} - V\| ~\leq~ \frac{1}{2\lambda}\min_{\substack{\tiS \subset S: \\ \vert \tiS\vert = B^*}} \|\nabla_{V^{pq}} L^{\lambda}_{\tiS}(f_W)\|
\end{equation*}
\end{restatable}
\begin{proof}
Let $\tiS \subset S$ be a batch of size $B^*$. By the chain rule, we can write the gradient of the loss function as follows
\begin{equation*}
\begin{aligned}
\frac{\partial L^{\lambda}_{\tiS}(f_W)}{\partial V^{ij}} ~=~ & \frac{1}{B^*}\sum^{B^*}_{(x,y) \in \tiS}\frac{\partial \ell(f_W(x),y)}{\partial f_W(x)} \cdot \frac{\partial f_W(x)}{\partial V^{ij}} + 2\lambda V^{ij}.
\end{aligned}
\end{equation*}
According to Lems.~\ref{lem:rank} and~\ref{lem:sec_rank}, the matrix $E_{\tiS} := -\frac{1}{2\lambda B^*}\sum^{B^*}_{(x,y) \in \tiS}\frac{\partial \ell(f_W(x),y)}{\partial f_W(x)} \cdot \frac{\partial f_W(x)}{\partial W^{ij}}$ is of rank $\leq B^*N^{ij}$. Therefore, we obtain the following
\begin{equation*}
\begin{aligned}
\min_{V:~ \rank(V) \leq B^*N^{ij}}\|W^{ij} - V\| ~\leq~ \min_{\tiS \subset S:~\vert \tiS\vert = B^*}\|W^{ij} - E_{\tiS}\|
~=~ \frac{1}{2\lambda}\min_{\tiS \subset S:~\vert \tiS\vert = B^*}\left\|\frac{\partial L^{\lambda}_{\tiS}(f_W)}{\partial W^{ij}} \right\|,
\end{aligned}
\end{equation*}
which completes the proof of this theorem.
\end{proof}
The theorem above provides an upper bound on the minimal distance between the various weight matrices in the network and matrices of rank $\leq B^*$. The upper bound is proportional to the minimal norm (w.r.t. the selection of the batch) of the gradient of a regularized risk evaluated on a batch of size $B^*$ and essentially depends on the training hyperparameters (e.g., batch size, learning rate, weight decay, etc'). Therefore, it can be used to draw predictions about the relationships between certain hyperparameters, the rank and the performance. Several interesting predictions follow the theorem above.
\subsection{Degeneracy and the Origin of SGD Noise}\label{sec:convergence}
We use the term `SGD noise' to refer to the inherent inability of SGD to converge to a stationary solution for the weight matrices that is the same across different mini-batches. In this section, we characterize the convergence of mini-batch SGD. Our results are essentially impossibility results that show that convergence of SGD takes place only when the network is the zero function. This implies that asymptotic noise is generically unavoidable (when training is successful). For simplicity, we focus on normalization-free univariate neural networks $f_W:\R^{n}\to \R$ and assume that $\forall~i\in [m]:~x_i\neq 0$.
{\color{red}Why is it the correct definition?}
Let $W$ be a convergence point of SGD. By the definition of convergence, SGD will not make any further steps with probability $1$. Assume by contradiction that there exists a pair $(v_p,v_q) \in E_T$ and a subset $\tiS\subset S$ of size $B$ for which $\nabla_{W^{pq}} L_{\tiS}(f_W)\neq 0$.
We note that at convergence, we have $\nabla_{W^{pq}} L_{\tiS}(f_W)=0$ across all pairs $(v_p,v_q) \in E_T$ and mini-batches $\tiS \subset S$ of size $B < m$. Specifically, we we can write,
\begin{equation}\label{eq:1}
\begin{aligned}
0 ~&=~ \nabla_{W^{pq}} L^{\lambda}_{\tiS}(f_W) ~=~ \frac{1}{B} \sum_{(x,y)\in \tiS} \frac{\partial \ell(f_W(x),y)}{\partial f_W(x)} \cdot \frac{\partial f_W(x)}{\partial W^{pq}} + 2\lambda W^{pq}.
\end{aligned}
\end{equation}
Suppose we have two batches $\tiS_1,\tiS_2 \subset S$ of size $B$ that differ by only one sample. We denote the unique sample of each batch by $(x_i,y_i)$ and $(x_j,y_j)$ respectively. We notice that,
\begin{equation}\label{eq:substract}
\begin{aligned}
0 ~&=~ \nabla_{W^{pq}} L^{\lambda}_{\tiS_1}(f_W) - \nabla_{W^{pq}} L^{\lambda}_{\tiS_2}(f_W) \\
~&=~ \frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \cdot \frac{\partial f_W(x_{i})}{\partial W^{pq}} - \frac{\partial \ell(f_W(x_{j}),y_{j})}{\partial f_W(x_{j})} \cdot \frac{\partial f_W(x_{j})}{\partial W^{pq}}.
\end{aligned}
\end{equation}
Therefore, we conclude that for all $i,j\in [m]$,
\begin{equation}\label{eq:ij}
\begin{aligned}
M^{pq} ~=~ \frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \cdot \frac{\partial f_W(x_{i})}{\partial W^{pq}} ~=~ \frac{\partial \ell(f_W(x_{j}),y_{j})}{\partial f_W(x_{j})} \cdot \frac{\partial f_W(x_{j})}{\partial W^{pq}}.
\end{aligned}
\end{equation}
Hence, for all $(v_p,v_q) \in E_T$ and $i \in [m]$,
\begin{equation}\label{eq:batchone}
\frac{\partial \ell(f_W(x_i),y_i)}{\partial f_W(x_i)} \cdot \frac{\partial f_W(x_i)}{\partial W^{pq}} + 2\lambda W^{pq} ~=~ M^{pq} + 2\lambda W^{pq} ~=~ 0.
\end{equation}
Therefore, unless $\lambda=0$ or $\forall~(v_p,v_q)\in E_T:~W^{pq}=0$, we conclude that $\frac{\partial \ell(f_W(x_i),y_i)}{\partial f_W(x_i)}\neq 0$ for all $i\in [m]$. In this case, we also obtain that $\{\frac{\partial f_W(x_i)}{\partial \textnormal{vec}(W^{pq})}\}^{m}_{i=1}$ are collinear vectors by Eq.~\ref{eq:ij}.
Therefore, any convergence point of training a neural network using mini-batch SGD along with weight decay is highly degenerate and does not perfectly fit any one of the training labels. To better understand the essence of this degeneracy, we provide the following proposition, which is specialized for ReLU networks.
\begin{restatable}[$\lambda > 0$]{proposition}{collinear}\label{prop:propA}
Let $\ell(a,b)$ be a differentiable loss function, $\lambda > 0$, and let $f_W(x)$ be a ReLU neural network, where $\{(v_p,v_1)\in E \mid v_p \in V\} \subset E_T$. Let $W$ be a convergence point of mini-batch SGD for minimizing $L^{\lambda}_S(f_W)$ (see Eq.~\ref{eq:1}). Then, either $f_W \equiv 0$ or $\{x_i\}^{m}_{i=1}$ are collinear vectors.
\end{restatable}
The above proposition shows that unless the training samples lie on a one-dimensional linear space, any convergence point of SGD corresponds to the zero function. Since both conditions are unrealistic, we conclude that convergence is impossible in any practical scenario.
In the following proposition, we look at the convergence of SGD when training without weight decay.
\begin{restatable}[$\lambda = 0$]{proposition}{zerolambda}\label{prop:propB}
Let $\ell(a,b)$ be a differentiable loss function, $\lambda = 0$, and let $f_W(x)$ be a ReLU neural network, where $\{(v_p,v_1)\in E \mid v_p \in V\} \subset E_T$. Let $W$ be a convergence point of mini-batch SGD for minimizing $L^{\lambda}_S(f_W)$. Assume that $\{x_i\}^{m}_{i=1}$ are not collinear vectors. Then, for any $i\in [m]$, for which $\frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \neq 0$, we have, $f_W(x_i)=0$. In particular, if $\ell$ is convex and $\forall~b\in \R~\nexists~a^* \in \R:~ \ell(a^*,b)=\inf_a \ell(a,b)$, then, $\forall~i \in [m]:~f_W(x_i)=0$.
\end{restatable}
The above proposition shows that at convergence, the neural network outputs zero for any training sample that it does not perfectly fit (i.e., $\frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \neq 0$). In particular, if $\ell$ is an exponential-type loss function (e.g., binary cross-entropy, logistic loss and exponential loss), $\ell(a,b)$ has no minima $a$ for any $b \in \R$, and therefore, the only possible convergence points are ones for which $f_W(x_i)=0$ across all $i \in [m]$. While theoretically convergence to a non-zero function is not guaranteed, in practice training without weight decay can still fall into the regime of `almost convergence', in which $\max_{i \in [m]} \big\vert \frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \big\vert$ is tiny and the training steps $- \eta \cdot \frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \cdot \frac{\partial f_W(x_i)}{\partial W^{pq}}$ are very small as a result. As a general tendency, when increasing the size of the network (i.e., the width or the depth), we should expect the term $\max_{i \in [m]} \big\vert \frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \big\vert$ to tend to zero (if the network is trained successfully). Furthermore, with the MSE loss we may actually have convergence that happens when the network perfectly fits the training labels, i.e., $\forall~i\in [m]:~f_W(x_i)=y_i$.% (in this case $\forall~i \in [m]:~\frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} = 0$).
% To summarize, when training a network for cross-entropy minimization with mini-batch SGD, convergence is not guaranteed (with or without regularization). That is because the convergence points are highly degenerate. On the other hand, when training a network to minimize the squared loss without regularization, any point $W$ for which $\forall~i \in [m]:~f_W(x_i)=y_i$ is a convergence point, which is achievable with overparameterized networks.
It is worth noting that our analysis critically depends on using SGD instead of GD during optimization. While many papers analyze the training dynamics and critical points of GD, the solutions of SGD are highly degenerate and will, in general, look different from the solutions of GD. For instance, it is impossible to derive Eq.~\ref{eq:ij} without relying on optimization with SGD, and without weight decay, we could not argue that $W^{pq}$ is proportional to the low-rank matrix $M^{pq}$. Finally, the analysis above is independent of the batch size as long as it is strictly smaller than the full dataset's size. %As follows from Eq.~\ref{eq:batchone} the convergence points of mini-batch SGD are equivalent to the convergence points of SGD with batch size $1$.
\section{Experiments}\label{sec:experiments}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/rank_batch_size_mlp_100_10_lr_0001.pdf}&
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/rank_batch_size_mlp_100_10_lr_001.pdf}&
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/rank_batch_size_mlp_100_10_lr_01.pdf}&
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/rank_batch_size_mlp_100_10_lr_05.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/acc_batch_size_mlp_100_10_lr_0001.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/acc_batch_size_mlp_100_10_lr_001.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/acc_batch_size_mlp_100_10_lr_01.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/acc_batch_size_mlp_100_10_lr_05.pdf}
\\
{\small $\eta=0.001$} & {\small $\eta=0.01$} & {\small $\eta=0.1$} & {\small $\eta=0.5$}
\end{tabular}
\caption{{\bf Average rank of MLP-BN-10-100 trained on CIFAR10 with varying batch sizes {\color{blue}($\epsilon=0.001$)}.} Each line stands for the results of a different batch size. In the (top) row we plot the average rank across layers and in the (bottom) row we plot the train and test accuracy rates for each setting.}
\label{fig:bs}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/rank_wd_mlp_100_10_bs_16.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/rank_wd_mlp_100_10_bs_32.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/rank_wd_mlp_100_10_bs_64.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/rank_wd_mlp_100_10_bs_128.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/acc_wd_mlp_100_10_bs_16.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/acc_wd_mlp_100_10_bs_32.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/acc_wd_mlp_100_10_bs_64.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/acc_wd_mlp_100_10_bs_128.pdf}
\\
{\small $B=16$} & {\small $B=32$} & {\small $B=64$} & {\small $B=128$}
\end{tabular}
\caption{{\bf Average rank of MLP-BN-10-100 trained on CIFAR10 with varying weight decay {\color{blue}($\epsilon=0.001$)}.} See Fig.~\ref{fig:bs} for details.}
\label{fig:wd}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/rank_batch_size_resnet_100_5_lr_00125.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/rank_batch_size_resnet_100_5_lr_0025.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/rank_batch_size_resnet_100_5_lr_005.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/rank_batch_size_resnet_100_5_lr_01.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/acc_batch_size_resnet_100_5_lr_00125.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/acc_batch_size_resnet_100_5_lr_0025.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/acc_batch_size_resnet_100_5_lr_005.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/acc_batch_size_resnet_100_5_lr_01.pdf}
\\
{\small $\eta=0.0125$} & {\small $\eta=0.025$} & {\small $\eta=0.05$} & {\small $\eta=0.1$}
\end{tabular}
\caption{{\bf Average rank of RES-BN-5-500 trained on CIFAR10 with varying batch sizes {\color{blue}($\epsilon=0.001$)}.} See Fig.~\ref{fig:bs} for details.}
\label{fig:bs_resnet}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/rank_batch_size_mlp_nobn_500_5_lr_00125.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/rank_batch_size_mlp_nobn_500_5_lr_0025.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/rank_batch_size_mlp_nobn_500_5_lr_005.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/rank_batch_size_mlp_nobn_500_5_lr_01.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/acc_batch_size_mlp_nobn_500_5_lr_00125.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/acc_batch_size_mlp_nobn_500_5_lr_0025.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/acc_batch_size_mlp_nobn_500_5_lr_005.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/acc_batch_size_mlp_nobn_500_5_lr_01.pdf}
\\
{\small $\eta=0.0125$} & {\small $\eta=0.025$} & {\small $\eta=0.05$} & {\small $\eta=0.1$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of MLP-5-500 trained on CIFAR10 with varying batch sizes {\color{blue}($\epsilon=0.001$)}.} See Fig.~\ref{fig:bs} for details.}
\label{fig:bs_mlp_nobn}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_01_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_05_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_3_batch_size.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_01_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_05_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_3_batch_size.pdf}
\\
{\small $\eta=0.01$} & {\small $\eta=0.05$} & {\small $\eta=0.1$} & {\small $\eta=0.3$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of {\color{blue}MLP-BN-10-100 (no-momentum)} trained on CIFAR10 with varying batch sizes {\color{blue}($\epsilon=0.001$)}.} See Fig.~\ref{fig:bs} for details.}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/rank_mlp_bn_bs_4_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/rank_mlp_bn_bs_8_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/rank_mlp_bn_bs_16_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/rank_mlp_bn_bs_32_lr_1_weight_decay.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/acc_mlp_bn_bs_4_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/acc_mlp_bn_bs_8_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/acc_mlp_bn_bs_16_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/acc_mlp_bn_bs_32_lr_1_weight_decay.pdf}
\\
{\small $B=4$} & {\small $B=8$} & {\small $B=16$} & {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of {\color{blue}MLP-BN-10-100 (no-momentum)} trained on CIFAR10 with varying weight decay {\color{blue}($\epsilon=0.001$)}.} See Fig.~\ref{fig:bs} for details.}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/rank_mlp_no_bn_lr_05_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/rank_mlp_no_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/rank_mlp_no_bn_lr_2_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/rank_mlp_no_bn_lr_3_batch_size.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/acc_mlp_no_bn_lr_05_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/acc_mlp_no_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/acc_mlp_no_bn_lr_2_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/acc_mlp_no_bn_lr_3_batch_size.pdf}
\\
{\small $\eta=0.05$} & {\small $\eta=0.1$} & {\small $\eta=0.2$} & {\small $\eta=0.3$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of {\color{blue}MLP-5-500 (no-momentum)} trained on CIFAR10 with varying batch sizes {\color{blue}($\epsilon=0.001$)}.} See Fig.~\ref{fig:bs} for details.}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/rank_mlp_no_bn_bs_4_lr_025_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/rank_mlp_no_bn_bs_8_lr_025_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/rank_mlp_no_bn_bs_16_lr_025_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/rank_mlp_no_bn_bs_32_lr_025_weight_decay.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/acc_mlp_no_bn_bs_4_lr_025_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/acc_mlp_no_bn_bs_8_lr_025_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/acc_mlp_no_bn_bs_16_lr_025_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/acc_mlp_no_bn_bs_32_lr_025_weight_decay.pdf}
\\
{\small $B=4$} & {\small $B=8$} & {\small $B=16$} & {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of {\color{blue}MLP-5-500 (no-momentum)} trained on CIFAR10 with varying weight decay {\color{blue}($\epsilon=0.001$)}.} See Fig.~\ref{fig:bs} for details.}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/rank_resnet_bn_lr_5_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/rank_resnet_bn_lr_8_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/rank_resnet_bn_lr_10_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/rank_resnet_bn_lr_20_batch_size.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/acc_resnet_bn_lr_5_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/acc_resnet_bn_lr_8_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/acc_resnet_bn_lr_10_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/acc_resnet_bn_lr_20_batch_size.pdf}
\\
{\small $\eta=0.5$} & {\small $\eta=0.8$} & {\small $\eta=1.0$} & {\small $\eta=2.0$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of {\color{blue}ResNet-18} trained on CIFAR10 with varying batch sizes {\color{blue}($\epsilon=0.001$)}.} See Fig.~\ref{fig:bs} for details.}
\label{fig:bs_resnet18}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/rank_resnet_bn_bs_4_lr_15_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/rank_resnet_bn_bs_8_lr_15_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/rank_resnet_bn_bs_16_lr_15_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/rank_resnet_bn_bs_32_lr_15_weight_decay.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/acc_resnet_bn_bs_4_lr_15_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/acc_resnet_bn_bs_8_lr_15_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/acc_resnet_bn_bs_16_lr_15_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/acc_resnet_bn_bs_32_lr_15_weight_decay.pdf}
\\
{\small $B=4$} & {\small $B=8$} & {\small $B=16$} & {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of {\color{blue}ResNet-18} trained on CIFAR10 with varying weight decay {\color{blue}($\epsilon=0.001$)}.} See Fig.~\ref{fig:bs} for details.}
\label{fig:wd_resnet18}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/rank_resnet_bn_batch_size_4_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/rank_resnet_bn_batch_size_8_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/rank_resnet_bn_batch_size_16_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/rank_resnet_bn_batch_size_32_learning_rate.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/acc_resnet_bn_batch_size_4_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/acc_resnet_bn_batch_size_8_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/acc_resnet_bn_batch_size_16_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/acc_resnet_bn_batch_size_32_learning_rate.pdf}
\\
{\small $B=4$} & {\small $B=8$} & {\small $B=16$} & {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of {\color{blue}ResNet-18} trained on CIFAR10 with varying learning rate {\color{blue}($\epsilon=0.001$)}.} See Fig.~\ref{fig:bs} for details.}
\label{fig:lr_resnet18}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/rank_vgg_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/rank_vgg_bn_lr_2_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/rank_vgg_bn_lr_3_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/rank_vgg_bn_lr_5_batch_size.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/acc_vgg_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/acc_vgg_bn_lr_2_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/acc_vgg_bn_lr_3_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/acc_vgg_bn_lr_5_batch_size.pdf}
\\
{\small $\eta=0.1$} & {\small $\eta=0.2$} & {\small $\eta=0.3$} & {\small $\eta=0.5$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of {\color{blue}VGG-16} trained on CIFAR10 with varying batch sizes {\color{blue}($\epsilon=0.001$)}.} See Fig.~\ref{fig:bs} for details.}
\label{fig:bs_vgg}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/rank_vgg_bn_batch_size_4_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/rank_vgg_bn_batch_size_8_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/rank_vgg_bn_batch_size_16_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/rank_vgg_bn_batch_size_32_lr_1_weight_decay.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/acc_vgg_bn_batch_size_4_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/acc_vgg_bn_batch_size_8_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/acc_vgg_bn_batch_size_16_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/acc_vgg_bn_batch_size_32_lr_1_weight_decay.pdf}
\\
{\small $B=4$} & {\small $B=8$} & {\small $B=16$} & {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of {\color{blue}VGG-16} trained on CIFAR10 with varying weight decay {\color{blue}($\epsilon=0.01$)}.} See Fig.~\ref{fig:bs} for details.}
\label{fig:wd_vgg}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/rank_vgg_bn_batch_size_4_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/rank_vgg_bn_batch_size_8_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/rank_vgg_bn_batch_size_16_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/rank_vgg_bn_batch_size_32_learning_rate.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/acc_vgg_bn_batch_size_4_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/acc_vgg_bn_batch_size_8_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/acc_vgg_bn_batch_size_16_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/acc_vgg_bn_batch_size_32_learning_rate.pdf}
\\
{\small $B=4$} & {\small $B=8$} & {\small $B=16$} & {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of {\color{blue}VGG-16} trained on CIFAR10 with varying learning rate {\color{blue}($\epsilon=0.001$)}.} See Fig.~\ref{fig:bs} for details.}
\label{fig:lr_vgg}
\end{figure*}
\begin{figure*}[t] % vit batch size
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.20\linewidth]{plots/vit/batch_size/varyB_lr4e-3_avgrank.pdf} &
\includegraphics[width=0.20\linewidth]{plots/vit/batch_size/varyB_lr8e-3_avgrank.pdf} &
\includegraphics[width=0.20\linewidth]{plots/vit/batch_size/varyB_lr1e-2_avgrank.pdf} &
\includegraphics[width=0.20\linewidth]{plots/vit/batch_size/varyB_lr2e-2_avgrank.pdf} &
\includegraphics[width=0.20\linewidth]{plots/vit/batch_size/varyB_lr4e-2_avgrank.pdf}
\\
\includegraphics[width=0.20\linewidth]{plots/vit/batch_size/varyB_lr4e-3_accuracy.pdf} &
\includegraphics[width=0.20\linewidth]{plots/vit/batch_size/varyB_lr8e-3_accuracy.pdf} &
\includegraphics[width=0.20\linewidth]{plots/vit/batch_size/varyB_lr1e-2_accuracy.pdf} &
\includegraphics[width=0.20\linewidth]{plots/vit/batch_size/varyB_lr2e-2_accuracy.pdf} &
\includegraphics[width=0.20\linewidth]{plots/vit/batch_size/varyB_lr4e-2_accuracy.pdf}
\\
{\small $\eta=0.004$} & {\small $\eta=0.008$} & {\small $\eta=0.01$} & {\small $\eta=0.02$} & {\small $\eta=0.04$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of {\color{blue}ViT (no-momentum)} trained on CIFAR10 with varying batch sizes {\color{blue}($\epsilon=0.01$)}.} See Fig.~\ref{fig:bs} for details.}
\end{figure*}
\begin{figure*}[t] % vit weight decay
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.20\linewidth]{plots/vit/weight_decay/varywd_B8_avgrank.pdf} &
\includegraphics[width=0.20\linewidth]{plots/vit/weight_decay/varywd_B16_avgrank.pdf} &
\includegraphics[width=0.20\linewidth]{plots/vit/weight_decay/varywd_B32_avgrank.pdf} &
\includegraphics[width=0.20\linewidth]{plots/vit/weight_decay/varywd_B64_avgrank.pdf} &
\includegraphics[width=0.20\linewidth]{plots/vit/weight_decay/varywd_B128_avgrank.pdf}
\\
\includegraphics[width=0.20\linewidth]{plots/vit/weight_decay/varywd_B8_accuracy.pdf} &
\includegraphics[width=0.20\linewidth]{plots/vit/weight_decay/varywd_B16_accuracy.pdf} &
\includegraphics[width=0.20\linewidth]{plots/vit/weight_decay/varywd_B32_accuracy.pdf} &
\includegraphics[width=0.20\linewidth]{plots/vit/weight_decay/varywd_B64_accuracy.pdf} &
\includegraphics[width=0.20\linewidth]{plots/vit/weight_decay/varywd_B128_accuracy.pdf}
\\
{\small $B=8$} & {\small $B=16$} & {\small $B=32$} & {\small $B=64$} & {\small $B=128$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of {\color{blue}ViT (no-momentum)} trained on CIFAR10 with varying weight decay {\color{blue}($\epsilon=0.01$)}.} See Fig.~\ref{fig:bs} for details.}
\end{figure*}
In this section we experimentally analyze the convergence (and non-convergence) of training neural networks with SGD. In addition, we empirically study the implicit bias towards rank minimization in deep ReLU networks. Throughout the experiments we extensively vary different hyperparameters (e.g., the learning rate, weight decay, and the batch size) and study their effect on the rank of the various matrices in the network\footnote{The plots are best viewed when zooming into the pictures.}. For each run we used one Tesla-k80 GPU for 5-20 hours.
\subsection{Setup}
\paragraph{Evaluation process.} Following the problem setup, we consider $k$-class classification problems and train a multilayered neural network $f_W = e \circ h = e \circ g^L \circ \dots \circ g^1 :\R^n \to \R^k$ on some balanced training dataset $S$. The model is trained using CE/~MSE loss minimization between its logits and the one-hot encodings of the labels. Here, $g^1,\dots,g^L$ are the various hidden layers of the network (e.g., fully-connected layers, residual blocks, {\color{blue}or convolutional layers}), where $e$ is its top linear layer. After each epoch, we compute the averaged rank across the network's weight matrices and its train and test accuracy rates. To estimate the rank of a given matrix $M$, we count how many of the singular values of $M/\|M\|_2$ are {\color{blue}$\notin [-\epsilon, \epsilon]$, where $\epsilon$ is a small tolerance value}. Training was done on the MNIST and CIFAR10 datasets. %During training, we used random cropping, random horizontal flips and random rotations (by $15k$ degrees for $k$ uniformly sampled from $[24]$) and standardize the data. %For all datasets we standardized the data.
\paragraph{Architectures.} We consider several network architectures. %The first architecture is a convolutional network (ConvNet), beginning with a stack of a $2\times 2$ convolutional layer with stride $2$, batch normalization, a convolution of the same structure, batch normalization and ReLU. Following that we have a set of $L$ stacks of $3\times 3$ convolutional layer with stride $1$ and padding $1$, batch normalization and ReLU. The output tensors of these layers shares the same shape as their input's shape. The number of channels in each convolutional layer is $H$. The last layer is linear. The network is denoted by Conv-$L$-$H$.
{\bf (i)} The first architecture is an MLP, denoted by MLP-BN-$L$-$H$, which consists of $L$ hidden layers, where each layer contains a fully-connected layer of width $H$, followed by batch normalization and ReLU activations. On top of that, we compose a fully-connected output layer. {\bf (ii)} The second architecture, denoted by RES-BN-$L$-$H$, consists of a linear layer of width $H$, followed by $L$ residual blocks, ending with a fully-connected layer. Each block computes a function of the form $z+\sigma(n_2(W_2 \sigma(n_1(W_1z))))$, where $W_1,W_2 \in \R^{H\times H}$, $n_1,n_2$ are batch normalization layers and $\sigma$ is the ReLU function. We denote by MLP-$L$-$H$ and RES-$L$-$H$ the same architectures without applying batch normalization. {\color{blue}{\bf(iii)} The third, denoted VGG-16, is the convolutional network proposed by \cite{DBLP:journals/corr/SimonyanZ14a}, but with dropout replaced by batch normalization layers and with only one fully-connected layer at the end. {\bf(iv)} The fourth architecture is the residual network proposed in \cite{7780459}, denoted ResNet-18.}
%Unless mentioned otherwise,
% {\bf Datasets.\enspace} Throughout the experiments, we consider three datasets: (i) MNIST; (ii) CIFAR10; (iii) STL-10. For CIFAR10 we used random cropping, random horizontal flips and random rotations (by $15k$ degrees for $k$ uniformly sampled from $[24]$). For all datasets we standardized the data. {\color{red}Update}
\subsection{Experiments on SGD Noise}\label{sec:sgd_noise_exp}
In Sec.~\ref{sec:convergence} we showed that training a neural network with SGD along with weight decay, cannot converge to a non-zero function. To verify this prediction, we trained MLP-5-2000 networks with varying degrees of weight decay $\lambda \in \{0, \expnumber{1}{-6}, \expnumber{1}{-5}, \expnumber{1}{-4}, \expnumber{1}{-3}\}$ and investigated their convergence. Each model was trained using SGD with batch size 128 and learning rate 0.1 for 2000 epochs (without learning rate scheduling or momentum). We trained the models once with the cross-entropy loss and once with the squared loss.
To evaluate the convergence of the networks, we use the following quantity: $d(W_{t+1},W_{t}) := \frac{1}{\vert E_T\vert} \sum_{(p,q)\in E_T} \|W^{pq}_{t+1} - W^{pq}_{t}\|$, where $\{W^{pq}_t\}_{(p,q)\in E_T}$ are the various trainable matrices in the network at epoch $t$. This quantity measures the averaged distance between the network's matrices at epochs $t$ and $t+1$. Convergence is possible only if $\lim_{t\to \infty} d(W_{t+1},W_{t}) = 0$.
In Fig.~\ref{fig:convergence} we monitor (a) $d(W_{t+1},W_{t})$, (b) the train accuracy rates, (c) the train losses and (d) the averaged rank of the trainable matrices. As predicted in Prop.~\ref{prop:propA}, when training with $\lambda>0$, $W_t$ either converges to zero and $f_{W_t}$ to the zero function (e.g., see the results of training on CIFAR10 with the cross-entropy loss and $\lambda=\expnumber{1}{-4}, \expnumber{1}{-3}$) or $W_t$ does not converge (i.e., $d(W_{t+1},W_{t})$ does not tend to zero). On the other hand, the term $d(W_{t+1},W_{t})$ is smaller by orders of magnitude when using $\lambda=0$ in comparison to using $\lambda>0$ (except for cases when the selection of $\lambda > 0$ leads to $W^{pq}_t \to 0$). Interestingly, even though for certain values of $\lambda > 0$, the training loss and accuracy converged, the network's parameters do not converge.
Finally, when training for cross-entropy loss minimization without weight decay we encounter the `almost convergence' regime discussed in Sec.~\ref{sec:convergence}. Namely, even though perfect convergence is impossible (see Prop.~\ref{prop:propB}) for a fixed sized neural network, the term $d(W_{t+1},W_{t})$ may become as small as we wish by increasing the size of the neural network. Therefore, since the MLP-5-2000 is relatively large (compared to the dataset's size), we may inaccurately get the impression that $d(W_{t+1},W_{t})$ tends to zero.
\subsection{Experiments on Rank Minimization}\label{sec:rank_exp}
\paragraph{Rank, batch size and the learning rate.} {\color{blue} {\bf (i)}} We train instances of MLP-BN-10-100, RES-BN-5-500 and MLP-5-500 using SGD with different batch sizes $B$ and initial learning rates $\eta$. The learning rate is decayed by a factor of $0.1$ three times at epochs 60, 100, and 200 and training is stopped after 500 epochs. We trained the models with weight decay $\lambda=\expnumber{5}{-4}$ and momentum $0.9$. {\color{blue} {\bf (ii)} We train instances of VGG-16 and ResNet-18 with the previous hyperparameters except with momentum $0$.}
As can be seen in Figs.~\ref{fig:bs},~\ref{fig:bs_resnet}, \ref{fig:bs_mlp_nobn}, {\color{blue}\ref{fig:bs_resnet18}, and~\ref{fig:bs_vgg}} by decreasing the batch size, we essentially strengthen the low-rank constraint over the network's matrices, which eventually leads to matrices of lower ranks. This is consistent with the prediction made in Sec.~\ref{sec:lowrank} that we expect to learn matrices of smaller ranks when training the network with smaller batch sizes. In addition, we also notice that when increasing the learning rate, the rank minimization constraint strengthens as well, which leads to lower ranks and generally worse training performances. Therefore, to properly train a neural network with smaller batch sizes, it is necessary to decrease the learning rate accordingly.
\paragraph{Rank minimization and weight decay.} {\color{blue} {\bf (i)}} To study the effect of the weight decay on rank minimization, we train instances of MLP-BN-10-100 using SGD with different batch sizes $B \in \{16,32,64,128\}$ and values of $\lambda \in \{0, \expnumber{1}{-4}, \expnumber{2}{-4}, \expnumber{4}{-4}, \expnumber{8}{-4}, \expnumber{6}{-3}\}$. The learning rate is initially $\eta=0.1$ and is decayed by a factor of $0.1$ three times at epochs 60, 100, and 200 and training is stopped after 500 epoch and the weight decay is $\lambda=\expnumber{5}{-4}$. {\color{blue} {\bf (ii)} We additionally train instances of VGG-16 and ResNet-18 with the previous hyperparamaters except with momentum $0$.}
The results are reported in Fig. \ref{fig:wd}{\color{blue}, \ref{fig:wd_resnet18} and~\ref{fig:wd_vgg}}. As can be seen, by increasing $\lambda$ we typically impose stronger rank minimization constraints, as predicted in Sec.~\ref{sec:lowrank}. In addition, regardless of the batch size, we lose the rank minimization constraint when training with $\lambda=0$ as predicted in Sec.~\ref{sec:lowrank}.
% We repeated the same experiment with instances of MLP-10-100 without the batch normalization as well. In this case, we used $\lambda \in \{0, \expnumber{1}{-3}, \expnumber{2}{-3}, \expnumber{4}{-3}, \expnumber{8}{-3}, \expnumber{6}{-2}, \expnumber{2}{-1}\}$. For each run we trained the network for 500 epochs using SGD with momentum $0.9$, initial learning rate $0.01$ and batch size 64.
\section{Conclusions}\label{sec:conclusions}
{\color{red}needs to be rewritten} In this work we theoretically studied the source of SGD noise and the implicit bias towards low-rank weight matrices in deep neural networks. We make several theoretical and empirical observations:
\begin{itemize}[leftmargin=0.5cm]
\item SGD noise is inevitable when applying weight decay, regardless of the batch size;
\item We show that the batch size, weight decay and learning rate contribute to a regularizing effect on the rank of the trainable matrices of neural networks;
\item Weight decay and training with mini-batches are necessary to obtain rank minimization.
\end{itemize}
It has not escaped our attention that these results have a rich set of implications with respect to other mysteries of deep learning that are yet to be solved. They include questions such as why SGD training of networks generalizes well, the relationship between rank minimization and architectural choices (residual networks~\citep{He_2017_ICCV}, transformers~\citep{NIPS2017_3f5ee243} and convolutional networks~\citep{3a6e6cc3ac6e4326b35a6315d7cbddf7}) and whether rank minimization extends beyond the standard settings of supervised learning.
\bibliography{refs}
\bibliographystyle{iclr2023_conference}
\appendix
\onecolumn
\section{Additional Experiments}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c@{}c}
\rotatebox[origin=c]{90}{{\tiny MNIST, CE}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/mnist_dist_ce_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/mnist_acc_ce_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/mnist_loss_ce_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/mnist_rank_ce_mlp_2000_lr_01.pdf}}
\\
\rotatebox[origin=c]{90}{{\tiny MNIST, MSE}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/mnist_dist_mse_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/mnist_acc_mse_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/mnist_loss_mse_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/mnist_rank_mse_mlp_2000_lr_01.pdf}}
\\
\rotatebox[origin=c]{90}{{\tiny CIFAR10, CE}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/cifar10_dist_ce_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/cifar10_acc_ce_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/cifar10_loss_ce_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/cifar10_rank_ce_mlp_2000_lr_01.pdf}}
\\
\rotatebox[origin=c]{90}{{\tiny CIFAR10, MSE}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/cifar10_dist_mse_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/cifar10_acc_mse_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/cifar10_loss_mse_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/cifar10_rank_mse_mlp_2000_lr_01.pdf}}
\\
& {\small (a)} & {\small (b)} & {\small (c)} & {\small (d)}
\end{tabular}
\caption{{\bf Convergence of MLP-5-2000 trained on MNIST and CIFAR10 with CE/~MSE loss.} In {\bf (a)} we plot the averaged distance between the weight matrices at epoch $t$ and epoch $t+1$, captured by, $\frac{1}{\vert E_T\vert}\sum_{(p,q) \in E_{T}} \|W^{pq}_{t+1} - W^{pq}_{t}\|$. In {\bf (b)} we plot the train accuracy rates, in {\bf (c)} we plot the averaged train loss and in {\bf (d)} we plot the average rank across the trainable matrices.}
\label{fig:convergence}
\end{figure*}
\section{Additional Architectures}
\paragraph{Optimization.} {\color{red}edit the detaching} In this work, we consider optimizing $f_W$ using mini-batch stochastic gradient descent (SGD), possibly with detaching certain layers. Namely, we aim at minimizing the regularized empirical risk $L^{\lambda}_{S}(f_{W})$ by applying SGD for a certain number of iterations. We initialize $W_0$ using a standard initialization procedure and iteratively update $W_t$. At each iteration, we sample a subset $\tiS = \{(x_{ij},y_{ij})\}^{B}_{j=1} \subset S$ uniformly at random and update $W_{t+1} \leftarrow W_t - \eta_t \cdot \bar{\nabla}_W L^{\lambda}_{\tiS}(f_{W_t}(~\cdot~;\tiX))$, where $\tiX = \{x_{ij}\}^{B}_{j=1}$ and $\eta_t > 0$ is our learning rate at iteration $t$. Here, $f_{W}(~\cdot~;\tiX) : \R^{n} \to \R$ is the function given by $f_W(x;\tiX)$ for a fixed value of $\tiX$. Furthermore, for a given function $g_{\theta} = g^L_{\theta} \circ \dots g^1_{\theta}:\R^{n} \to \R^k$ where $\{g^{i_j}_{\theta}\}^{r}_{j=1}$ are detached, $\frac{\bar{\partial} g_{\theta}(z)}{\partial \theta} = \bar{\nabla}_{\theta} g_{\theta}(z)$ denotes the {\color{red}pseudo gradient} of $g_{\theta}$, which is computed using the chain rule, while treating $\{g^{i_j}_{\theta}\}^{r}_{j=1}$ as independent of $\theta$.
\paragraph{Network architecture.} Formally, $f_W$ is as a directed acyclic graph (DAG) $G=(V,E)$, where $V= \{v_1,\dots,v_L\}$ consists of the various layers of the network and each edge $e_{ij}=(v_i,v_j) \in E$ specifies a connection between two layers. Each layer is a function $v_{i}:\R^{c_{1} \times h_1 \times w_1} \to \R^{c_{i} \times h_i \times w_i}$ and each connection holds a linear transformation $C^{ij} : \R^{c_{j} \times h_j \times w_j} \to \R^{c_i \times h_i \times w_i}$. The layers are divided into three categories: the input layer $v_1$, the output layer $v_L$ and intermediate layers. In this setting, we do not have connections directed towards the input layer nor connections coming out of the output layer (i.e., $\forall~i \in [L]:~(v_i,v_L), (v_0,v_i) \notin E$). Given an input $x$ and batch $\tiX$, the output of a given layer $v_i$ is evaluated as follows
\begin{equation*}
v_i(x;\tiX) ~:=~ \sigma\left(n_i\left[\sum_{j \in \textnormal{pred}(i)} C^{ij}(v_{j}(x;\tiX));\tiX\right]\right),
\end{equation*}
except for the output layer $v_L$ that computes
\begin{equation*}
f_W(x;\tiX) ~:=~ v_L(x;\tiX) ~:=~ \sum_{j \in \textnormal{pred}(L)}C^{Lj} (v_{j}(x;\tiX)),
\end{equation*}
where $\textnormal{pred}(i)$ is the set of indices $j$, such that $(v_i,v_j) \in E$. Here, $\sigma$ is an element-wise activation function (e.g., ReLU, Leaky ReLU, sigmoid).
In this setting, the transformations $C^{pq}$ could be trainable, $(v_p,v_q) \in E_T$ (e.g., convolutional layers), or constant, $(v_p,v_q) \notin E_T$ (e.g., residual connections or a pooling layer). Throughout the analysis, we consider paths within the graph $G$, denoted by $\pi = (\pi_0,\dots,\pi_T)$, where $\pi_0=1$, $\pi_T = L$ and for all $i=0,\dots,T-1:~ (v_{\pi_{i}},v_{\pi_{i+1}}) \in E$.
\paragraph{Convolutional layers.} {\color{red}add a formal def maybe?} Throughout the paper we consider convolutional layers~\citep{6795724,10.5555/646469.691875} (see also~\citep{GoodBengCour16}) $C:\R^{c_1 \times h_1 \times w_1} \to \R^{c_2 \times h_2 \times w_2}$ with kernel sizes $(k_1,k_2)$, padding $p$ and stride $s$. The formula for the output dimensions are $h_2 = \left(\lfloor(h_1 - k_1 +2p)\rfloor/s+1\right)$ and $w_2 = \left(\lfloor(w_1 - k_2 +2p)\rfloor/s+1\right)$ and the number of patches on which the layer acts upon is $ h_2w_2$. For a given convolutional layer $C^{pq}$ in the network $f_W$, we denote $N^{pq}:=h_pw_p$. For a given convolutional layer $C^{pq}$ we define a matrix $W^{pq}$ that computes $\forall~ x \in \R^{c_q \times h_q \times w_q}:~ W^{pq} \vect(x) = \vect(C^{pq}(x))$. This is guaranteed since both the padding and convolution operations of the layer are linear. We also consider $V^{pq} \in \R^{c_{p} \times c_{q} k_1k_2}$ which is a matrix whose rows are $\vect(V^{pq}_{c,:,:})$.
\paragraph{Pooling Layers.} A pooling layer (see~\citep{Goodfellow-et-al-2016}) $C$ with kernel dimensions $(k_1,k_2)$ stride $s$ and padding $p$ takes an input $c_{in}\times h_1 \times w_1$ and computes an output with dimensions $h_2 = \left(\lfloor(h_1 - k_1 +2p)\rfloor/s+1\right)$ and $w_2 = \left(\lfloor(w_1 - k_2 +2p)\rfloor/s+1\right)$. In this paper we consider average pooling and max-pooling layers.
\paragraph{Batch normalization layers.} A batch normalization layer~\citep{pmlr-v37-ioffe15} is defined in the following manner,
\begin{small}
\begin{equation*}
\begin{aligned}
n_l(v_l(x;\tiX);\tiX) ~&=~ \gamma_l \cdot A_l(\tiX) \odot (v_l(x;\tiX) - b_l(\tiX)) \\
&\quad + \one_{c\times h_l \times w_l}\cdot \beta_l \\
A_l(\tiX) ~&=~ \left(\left(\sqrt{\Var_{\hat{x} \sim U[\tiX]}(v_l(\hat{x};\tiX)_{j})}\right)^{-1}\right)_{j} \\
b_l(\tiX) ~&=~ \E_{\hat{x} \sim U[\tiX]}[v_l(\hat{x};\tiX)],
\end{aligned}
\end{equation*}
\end{small}
where $v_l(\hat{x};\tiX)_{j}$ is the $j$'th coordinate of $v_l(\hat{x};\tiX)$ for $j \in [c_l]\times [h_l]\times [w_l]$. In addition, $\gamma_l,\beta_l \in \R$ are trainable scale and shift parameters and $\one_{a_1 \times \dots \times a_k} \in \R^{a_1\times \dots \times a_k}$ is a tensor whose all of its cells are equal $1$. To simplify the analysis, we detach (also known as stop-grad~\citep{NEURIPS2020_f3ada80d}) the matrices $A_l(\tiX)$ during backpropagation at any step during training. Namely, we treat $\frac{\bar{\partial} A_l(\tiX)}{\partial W}=0$ in order to compute the training updates (without changing the forward passes). We note that detaching is a commonly used approach for training a network with normalization layers~\citep{Wiesler2014MeannormalizedSG,renorm,pmlr-v22-raiko12,https://doi.org/10.48550/arxiv.1410.7455,NEURIPS2019_2f4fe03d,10.1007/s10994-020-05929-w,https://doi.org/10.48550/arxiv.2202.07304}. In fact,~\cite{NEURIPS2019_2f4fe03d} showed that detaching the variance term of layer normalization is competitive and even improves the performance on various tasks.
\section{Proofs}
\ranklemma*
\begin{proof} Let $x \in \R^{c_1\times h_1\times w_1}$ be an input tensor and $C^{ij}$ be a certain convolutional layer with kernel size $(k_1,k_2)$, stride $s$ and padding $p$. We would like to show that $\rank\left(\frac{\partial f_W(x)}{\partial V^{ij}}\right) \leq N^{ij}$. We begin by writing the output of $f_W$ as a sum over paths that pass through $C^{ij}$ and paths that do not. We note that the output can be written as follows,
\begin{equation*}
f_W(x) ~=~ \sum_{l_1 \in \textnormal{pred}(l_0)} C^{l_0l_1} \circ v_{l_1}(x),
\end{equation*}
where $l_0 = L$ and $C^{l_0l_1} \circ z := C^{l_0l_1}(z)$. In addition, each layer $v_{l}$ can also be written as
\begin{equation*}
\vect(v_{l_1}(x)) ~=~ D_{l_1} \odot \sum_{l_2\in \textnormal{pred}(l_1)} C^{l_1l_2} \circ v_{l_2}(x),
\end{equation*}
where $D_l := D_l(x) := \sigma'(v_l(x))) \in \R^{c_l \times h_l \times w_l}$. Therefore, for any input $x$, we can write $f_W(x)$ as the sum of matrix multiplications along paths $\pi$ from $v_1$ to $v_{l_0}$. Specifically, we can write $f_W(x)$ as a follows
%\begin{small}
\begin{equation*}
\begin{aligned}
f_W(x) ~&=~ \sum_{\pi \textnormal{ from $i$ to $l_0$}} C^{\pi_{T}\pi_{T-1}} \circ D_{\pi_{T-1}} \cdots D_{\pi_2} \odot C^{\pi_2\pi_1} \circ D_{\pi_1} \odot C^{ij} \circ v_j(x) \\
&\quad\quad + \sum_{\substack{\pi \textnormal{ from $1$ to $l_0$} \\ (i,j) \notin \pi }} C^{\pi_{T}\pi_{T-1}} \circ D_{\pi_{T-1}} \odot C^{\pi_{T-1}\pi_{T-2}} \cdots D_{\pi_2} \odot C^{\pi_2 \pi_1} \circ x, \\
~&=:~ A_W(x)+B_W(x)
\end{aligned}
\end{equation*}
%\end{small}
where $T(\pi)$ denotes the length of the path $\pi$ (denoted by $T$ for short). Since $\sigma$ is a piece-wise linear function, for any $x\in \R^{c_1\times h_1\times w_1}$, with measure $1$ over $W$, the matrices $\{D_{l}(x)\}^{L-1}_{l=1}$ are constant in the neighborhood of $W$. Furthermore, $V^{ij}$ does not appear in the multiplications along the paths $\pi$ from $1$ to $l_0$ that exclude $(i,j)$. Therefore, we conclude that $\frac{\partial B_W(x)}{\partial V^{ij}}=0$.
As a next step we would like to analyze the rank of $\frac{\partial A_W(x)}{\partial V^{ij}}$. For this purpose, we rewrite the convolutional layers and the multiplications by the matrices $D_{l}$ as matrix multiplications.
\paragraph{Representing $C^{ij}$.} We begin by representing the layer $C^{ij}$ as a linear transformation of its input with $N^{ij}$ blocks of $V^{ij}$. For this purpose, we define a representation of a given $3$-dimensional tensor input $x \in \R^{c_j \times h_j \times w_j}$ as a vector $\vect^{ij}(x) \in \R^{N^{ij} c_{j} k_1 k_2}$. First, we pad $x$ with $p$ rows and columns of zeros and obtain $\textnormal{Pad}_p(x)$. We then vectorize each one of its patches (of dimensions $c_{j}\times k_1\times k_2$) that the convolutional layer is acting upon (potentially overlapping) and concatenate them. Furthermore, we define a matrix $V^{ij} \in \R^{c_{i} \times (c_{j}k_1k_2)}$, whose $c$'th row is $\vect(Z^{ij}_{c,~:,~:~})$. We can write the vectorized output of the convolutional layer as $U^{ij}\cdot \vect^{ij}(x)$, where
\begin{equation}\label{eq:Uij}
U^{ij} := \renewcommand{\arraystretch}{1.2}
\left(
\begin{array}{ c c c c c c }
\multicolumn{1}{c}{} & & 0 & 0 & 0 & 0 \\
\multicolumn{2}{c}{\raisebox{.6\normalbaselineskip}[0pt][0pt]{$V^{ij}$}} & 0 & 0 & 0 & 0 \\
0 & 0 & & & 0 & 0 \\
0 & 0 & \multicolumn{2}{c}{\raisebox{.6\normalbaselineskip}[0pt][0pt]{$\ddots$}} & 0 & 0 \\
0 & 0 & 0 & 0 & & \multicolumn{1}{c}{} \\
0 & 0 & 0 & 0 & \multicolumn{2}{c}{\raisebox{.6\normalbaselineskip}[0pt][0pt]{$V^{ij}$}}
\end{array}
\right)
\end{equation}
is a $(N^{ij} c_{i}) \times (N^{ij} c_{j} k_1k_2)$ matrix with $N^{ij}$ copies of $V^{ij}$. We note that this is a non-standard representation of the convolutional layer's operation as a linear transformation. Typically, we write the convolutional layer as a linear transformation $W^{ij}$ acting on the vectorized version $\vect(x) \in \R^{c_{j}k_1k_2}$ of its input $x$. Since $\vect^{ij}(x)$ consists of the same variables as in $\vect(x)$ with potentially duplicate items, there is a linear transformation that translates $\vect(x)$ into $\vect^{ij}(x)$. Therefore, we can simply write $W^{ij} \cdot \vect(x) = U^{ij} \cdot \vect^{ij}(x)$.
{\bf Representing convolutional layers.\enspace} Except of $C^{ij}$, we represent each one of the network's convolutional layers $C^{ld}$ in $f_W$ as linear transformations. As mentioned earlier, we can write $\vect(C^{ld}(x)) = W^{ld}\vect(x)$, for any input $x\in \R^{c_d \times h_d \times w_d}$.
{\bf Representing pooling and rearrangement layers.\enspace} An average pooling layer or a rearrangement layer $C^{ld}$ can be represented as a (non-trainable) linear transformation of its input. Namely, we can write $\vect(C^{ld}(x)) = W^{ld}\vect(x)$ for some matrix $W^{ld}$. A max-pooling layer can be written as a composition of ReLU activations and multiple (non-trainable) linear transformations, since $\max(x,y) = \sigma(x-y) + y$. Therefore, without loss of generality we assume that the network is free of max-pooling layers.
{\bf Computing the rank.\enspace} Finally, we note that $\vect(C^{ij} \circ z) = U^{ij} \cdot \vect^{ij}(z) = W^{ij} \cdot \vect(z)$, $\vect(D_l \odot z) = P_l \cdot \vect(z)$ for $P_{l} := \textnormal{diag}(\vect(D_{l}))$. Therefore, we can write
\begin{equation*}
\begin{aligned}
A_W(x) ~&=~ \sum_{\pi \textnormal{ from $i$ to $l_0$}} W^{\pi_{T}\pi_{T-1}} \cdot P_{\pi_{T-1}} \cdots P_{\pi_2} \cdot W^{\pi_2\pi_1} \cdot P_{\pi_1} \cdot U^{ij} \cdot \vect^{ij}(v_j(x)) \\
~&=:~ a(x)^{\top} \cdot U^{ij} \cdot b(x),
\end{aligned}
\end{equation*}
where $a(x)^{\top} ~:=~ \sum_{\pi \textnormal{ from $i$ to $l_0$}} W^{\pi_{T}\pi_{T-1}} \cdot P_{\pi_{T-1}} \cdots P_{\pi_2} \cdot W^{\pi_2\pi_1} \cdot P_{\pi_1}$ and $b(x) := \vect^{ij}(v_j(x))$.
We note that with measure $1$, the matrices $\{P_l\}^{L-1}_{l=1}$ are constant in the neighborhood of $W$. In addition, $a(x)$ and $b(x)$ are computed as multiplications of matrices $W^{ld}$ and $P_l$ excluding $(i,j)=(p,q)$. Therefore, with measure $1$ over the selection of $W$, the Jacobians of $a(x)$ and $b(x)$ with respect to $V^{ij}$ are $0$. Furthermore, due to \eqref{eq:Uij} and the definition of $U^{ij}$, we can write
\begin{equation*}
\begin{aligned}
a(x)^{\top} \cdot U^{ij} \cdot b(x) ~=~ \sum^{N^{ij}}_{t=1} a_t(x)^{\top} \cdot V^{ij} \cdot b_t(x),
\end{aligned}
\end{equation*}
where $a_t(x)$ and $b_t(x)$ are the slices of $a(x)$ and $b(x)$ that are multiplied by the $t$'th $V^{ij}$ block in $U^{ij}$. Since the Jacobians of $a_i(x)$ and $b_i(x)$ with respect to $V^{ij}$ are $0$ with measure $1$ over the selection of $W$, we have,
\begin{equation}
\frac{\partial a(x)^{\top} \cdot U^{ij} \cdot b(x)}{\partial V^{ij}} ~=~ \sum^{N^{ij}}_{t=1} a_t(x) \cdot b_t(x)^{\top}.
\end{equation}
Therefore, we conclude that, with measure $1$ over the selection of $W$, we have $\frac{\partial f_W(x)}{\partial V^{ij}} = \sum^{N^{ij}}_{t=1} a_t(x) \cdot b_t(x)^{\top}$ which is a matrix of rank $\leq N^{ij}$.
\end{proof}
\begin{restatable}{lemma}{lowrankSec}\label{lem:sec_rank}
Let $\tf_W$ be a convolutional neural network with batch-normalization. Let $C^{pq}$ be a given convolutional layer within $\tf_W$ with kernel size $(k_1,k_2)$, stride $s$, padding $p$, $c_{q}$ input channels and $c_{p}$ output channels. Let $V^{pq} \in \R^{c_{q} \times (c_{p} k_1k_2)}$ be the parameters matrix of $C^{pq}$. Let $\tiX = \{x_i\}^{B}_{i=1}$ be a batch of samples. Then, the ranks of $\frac{\partial f_W(x;\tiX)}{\partial V^{pq}}$ and $\sum^{B}_{i=1}\alpha_i \frac{\partial f_W(x_i;\tiX)}{\partial V^{pq}}$ are at most $BN^{pq}$.
\end{restatable}
\begin{proof}
We would like to show that $\frac{\bar{\partial} f_W(x_i;\tiX)}{\partial V^{pq}}$ and $\sum^{B}_{i=1}\beta_i \cdot \frac{\bar{\partial} f_W(x_i;\tiX)}{\partial V^{pq}}$ are matrices of rank $\leq N^{pq}B$. We follow the proof of Lem.~2 in~\citep{https://doi.org/10.48550/arxiv.2206.05794} and extend it to the case where the network includes convolutional layers. For this purpose, we represent the top layers of the network as a sequence of linear transformations and the layer $C^{pq}$ as a block matrix with blocks $V^{pq}$ by the same technique as in the proof of Lem.~\ref{lem:rank}. The layers below $C^{pq}$ are composed together with the input into a vector that is independent of $V^{pq}$.
We note that for any fixed vector $x \in \R^{c_1 \times h_1 \times w_1}$, the network can be written as follows
\begin{equation*}
f(x;\tiX) ~=~ \sum_{l_1 \in \textnormal{pred}(L)} C^{Ll_1} (v_{l_1}(x;\tiX)).
\end{equation*}
In addition, each layer $v_{l_1}$ can be written as follows
\begin{equation*}
\begin{aligned}
v_{l_1}(x;\tiX) ~&=~ D_{l_1}(x;\tiX) \odot n_{l_1}(u_{l_1}(x;\tiX);\tiX),
\end{aligned}
\end{equation*}
where $D_{l_1}(x;\tiX) := \sigma'(n_{l_1}(u_{l_1}(x;\tiX);\tiX))$ and $u_{l_1}(x;\tiX) := \sum_{l_2\in \textnormal{pred}(l_1)} C^{l_1l_2} (v_{l_2}(x;\tiX))$. Therefore, we can write
\begin{equation*}
n_{l_1}(u_{l_1}(x;\tiX);\tiX) ~=~ \gamma_{l_1} \cdot A_{l_1}(\tiX) \odot \left(u_{l_1}(x;\tiX) - b_{l_1}(\tiX) \right) + \one_{c_{l_1} \times h_{l_1} \times w_{l_1}} \cdot \beta_{l_1},
\end{equation*}
where $b_{l}(\tiX) := \E_{\hat{x}\sim U[\tiX]}[u_{l}(\hat{x};\tiX)] = \sum_{t\in \textnormal{pred}(l)} C^{lt} (\mu_t(\tiX))$ and $\mu_t(\tiX) := \E_{x \sim U[\tiX]}[v_{t}(x;\tiX)]$. Hence, we can write
\begin{equation*}
\begin{aligned}
v_{l_1}(x;\tiX) ~&=~ \sum_{l_2\in \textnormal{pred}(l_1)} Q_{l_1}(x;\tiX) \odot \left[ C^{l_1l_2}(v_{l_2}(x;\tiX)) - C^{l_1l_2}(\mu_{l_2}(\tiX)) \right] + D_{l_1}(x;\tiX) \cdot \beta_{l_1} \\
~&=~ \sum_{l_2\in \textnormal{pred}(l_1)} Q_{l_1}(x;\tiX) \odot C^{l_1l_2}\left( v_{l_2}(x;\tiX) - \mu_{l_2}(\tiX) \right) + D_{l_1}(x;\tiX) \cdot \beta_{l_1}, \\
\end{aligned}
\end{equation*}
where $Q_{l}(x;\tiX) := D_{l}(x;\tiX) \odot A_{l}(\tiX)$ and the last equation follows from the linearity of convolutional layers.
% In addition, the flattened version of each layer $v_{l_1}$ can be written as
% \begin{equation*}
% \begin{aligned}
% \vect(v_{l_1}(x;\tiX)) ~&=~ \bar{D}_{l_1}(x;\tiX) \cdot \vect(n_l(u_l(x;\tiX);\tiX)) %\left[\gamma_{l_1} \cdot \bar{A}_{l_1}(\tiX) \cdot \left(\sum_{l_2\in \textnormal{pred}(l_1)} W^{l_1l_2} \cdot \vect(v_{l_2}(x;\tiX)) - \vect(b_{l_2}(\tiX)) \right) + \one_{c_{l_1} h_{l_1} w_{l_1}} \cdot \beta_{l_1}\right]\\
% %~&=~ \bar{D}_{l_1}(x;\tiX) \cdot \left[ \gamma_{l_1}\cdot \bar{A}_{l_1}(\tiX) \cdot \left(\sum_{l_2\in \textnormal{pred}(l_1)} W^{l_1l_2} \cdot \left[\vect(v_{l_2}(x;\tiX)) - \vect(\mu_{l_2}(\tiX)) \right] \right) +\one_{c_{l_1} h_{l_1} w_{l_1}} \cdot \beta_{l_1} \right]\\
% %~&=~ \sum_{l_2\in \textnormal{pred}(l_1)} Q_{l_1}(x;\tiX) \cdot W^{l_1l_2}\cdot \left( \vect(v_{l_2}(x;\tiX)) - \vect(\mu_{l_2}(\tiX)) \right) + \bar{D}_{l_1}(x;\tiX) \cdot \one_{c_{l_1} h_{l_1} w_{l_1}} \cdot \beta_{l_1},
% \end{aligned}
% \end{equation*}
% where $\bar{D}_l(x;\tiX) := \textnormal{diag}[D_l(x;\tiX)]$, $D_l(x;\tiX) := \sigma'(n_l(u_l(x;\tiX);\tiX))$ and
% \begin{equation*}
% \vect(n_l(u_l(x;\tiX);\tiX)) ~=~ \left[\gamma_{l} \cdot A_{l}(\tiX) \cdot \left(\sum_{l'\in \textnormal{pred}(l)} W^{ll'} \cdot \vect(v_{l'}(x;\tiX)) - \vect(b_{l'}(\tiX)) \right) + \one_{c_{l_1} h_{l_1} w_{l_1}} \cdot \beta_{l}\right].
% \end{equation*}
% \begin{equation*}
% \begin{aligned}
% \vect(v_{l_1}(x;\tiX)) ~&=~ \bar{D}_{l_1}(x;\tiX) \cdot \left[\gamma_{l_1} \cdot \bar{A}_{l_1}(\tiX) \cdot \left(\sum_{l_2\in \textnormal{pred}(l_1)} W^{l_1l_2} \cdot \vect(v_{l_2}(x;\tiX)) - \vect(b_{l_2}(\tiX)) \right) + \one_{c_{l_1} h_{l_1} w_{l_1}} \cdot \beta_{l_1}\right]\\
% ~&=~ \bar{D}_{l_1}(x;\tiX) \cdot \left[ \gamma_{l_1}\cdot \bar{A}_{l_1}(\tiX) \cdot \left(\sum_{l_2\in \textnormal{pred}(l_1)} W^{l_1l_2} \cdot \left[\vect(v_{l_2}(x;\tiX)) - \vect(\mu_{l_2}(\tiX)) \right] \right) +\one_{c_{l_1} h_{l_1} w_{l_1}} \cdot \beta_{l_1} \right]\\
% ~&=~ \sum_{l_2\in \textnormal{pred}(l_1)} Q_{l_1}(x;\tiX) \cdot W^{l_1l_2}\cdot \left( \vect(v_{l_2}(x;\tiX)) - \vect(\mu_{l_2}(\tiX)) \right) + \bar{D}_{l_1}(x;\tiX) \cdot \one_{c_{l_1} h_{l_1} w_{l_1}} \cdot \beta_{l_1},
% \end{aligned}
% \end{equation*}
% Here,
% \begin{equation*}
% b_l(\tiX) ~:=~ \sum_{t\in \textnormal{pred}(l)} W^{lt} \cdot \mu_t(\tiX) ~:=~ \sum_{t\in \textnormal{pred}(l)} W^{lt} \cdot \E_{x \sim U[\tiX]}\left[ v_{t}(x;\tiX) \right]
% \end{equation*}
% and $Q_l(x;\tiX) := \gamma_l \cdot D_{l}(x;\tiX) \cdot A_l(x;\tiX)$. For simplicity, we write $l_0 := L$, $\mu_l := \mu_l(x;\tiX)$, $v^j_{l} := v_{l}(x_j;\tiX)$ and $Q^{j}_l := Q_l(x_j;\tiX)$.
We denote $\mu_l := \mu_l(\tiX)$, $v^j_l := v_l(x_j;\tiX)$, $D^{j}_l := D_l(x_j;\tiX)$, $A_l := A_l(\tiX)$, $Q^{j}_l := Q_l(x_j;\tiX)$ and define
\begin{equation*}
\Delta^{j}_l := \twopartdef
{x_j} {l=1}
{v^{j}_l - \mu_{l}} {l\neq 1}.
\end{equation*}
Hence, we can write $f(x_{j_0};\tiX)$ as a follows
\begin{equation*}
\begin{aligned}
f(x_{j_0};\tiX) ~&=~ \sum_{l_1 \in \textnormal{pred}(l_0)} C^{l_0l_1} \cdot v^{j_0}_{l_1} \\
~&=~ \sum_{l_1 \in \textnormal{pred}(l_0)} C^{l_0l_1} \left(Q^{j_0}_{l_1} \odot \left( \sum_{l_2 \in \textnormal{pred}(l_1)} C^{l_1l_2} (v^{j_0}_{l_2} - \mu_{l_2}) \right) +D^{j_0}_{l_1}\cdot \beta_{l_1}\right) \\
~&=~ \sum_{l_1 \in \textnormal{pred}(l_0)} C^{l_0l_1} (D^{j_0}_{l_1}\cdot \beta_{l_1}) \\
&\quad + \sum_{l_1 \in \textnormal{pred}(l_0)} \sum_{l_2 \in \textnormal{pred}(l_1)} C^{l_0l_1} (Q^{j_0}_{l_1} \odot C^{l_1l_2} (v^{j_0}_{l_2} - \mu_{l_2})) \\
~&=~ \sum_{l_1 \in \textnormal{pred}(l_0)} C^{l_0l_1} (D^{j_0}_{l_1}\cdot \beta_{l_1}) \\
&\quad + \sum_{l_1 \in \textnormal{pred}(l_0)} \sum_{l_2 \in \textnormal{pred}(l_1)} C^{l_0l_1} (Q^{j_0}_{l_1} \odot C^{l_1l_2} (\Delta^{j_0}_{l_2})) \\
\end{aligned}
\end{equation*}
% Similarly,
% \begin{equation}
% \begin{aligned}
% \Delta^{j_0}_{l_2} ~&=~ v^{j_0}_{l_2} - \mu_{l_2} \\
% ~&=~ D^{j_0}_{l_2} \cdot \one \cdot \beta_{l_2} \\
% &\quad+ Q^{j_0}_{l_2} \cdot \left(\sum_{l_3\in \textnormal{pred}(l_2)} W^{l_2l_3} \cdot (v^{j_0}_{l_3} - \mu_{l_3}) \right) \\
% &\quad- \frac{1}{B}\sum^{B}_{j_1=1} Q^{j_1}_{l_2} \cdot \left(\sum_{l_3\in \textnormal{pred}(l_2)} W^{l_2l_3} \cdot (v^{j_1}_{l_3} - \mu_{l_3}) \right) \\
% ~&=~ D^{j_0}_{l_2} \cdot \one \cdot \beta_{l_2} + \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{l_2} \cdot \left(\sum_{l_3\in \textnormal{pred}(l_2)} W^{l_2l_3} \cdot (v^{j_1}_{l_3} - \mu_{l_3}) \right) \\
% ~&=~ D^{j_0}_{l_2} \cdot \one \cdot \beta_{l_2} + \sum_{l_3\in \textnormal{pred}(l_2)}\sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{l_2} \cdot W^{l_2l_3} \cdot \Delta^{j_1}_{l_3}, \\
% \end{aligned}
% \end{equation}
For any $l_{t+2}\neq 1$, we can write
\begin{equation}\label{eq:delta}
\begin{aligned}
\Delta^{j_t}_{l_{t+2}} ~&=~ v^{j_t}_{l_{t+2}} - \mu_{l_{t+2}} \\
~&=~ \left(D^{j_{t}}_{l_{t+2}} - \frac{1}{B}\sum^{B}_{j_{t+1}=1} D^{j_{t+1}}_{l_{t+2}}\right) \cdot \beta_{l_{t+2}} \\
&\quad + Q^{j_t}_{l_{t+2}} \odot \left(\sum_{l_{t+3}\in \textnormal{pred}(l_{t+2})} C^{l_{t+2}l_{t+3}} (v^{j_t}_{l_{t+3}} - \mu_{l_{t+3}}) \right) \\
&\quad- \frac{1}{B}\sum^{B}_{j_{t+1}=1} Q^{j_{t+1}}_{l_{t+2}} \odot \left(\sum_{l_{t+3}\in \textnormal{pred}(l_{t+2})} C^{l_{t+2}l_{t+3}} (v^{j_{t+1}}_{l_{t+3}} - \mu_{l_{t+3}}) \right) \\
~&=~ \sum^{B}_{j_{t+1}=1}\alpha_{j_t,j_{t+1}} \cdot D^{j_{t+1}}_{l_{t+2}} \cdot \beta_{l_{t+2}} + \sum_{l_{t+3}\in \textnormal{pred}(l_{t+2})}\sum^{B}_{j_{t+1}=1} \alpha_{j_t,j_{t+1}} Q^{j_{t+1}}_{l_{t+2}} \odot C^{l_{t+2}l_{t+3}} (\Delta^{j_{t+1}}_{l_{t+3}}), \\
\end{aligned}
\end{equation}
where $\alpha_{i,j}=-1/B$ for all pairs $i\neq j$ and $\alpha_{i,i}=\frac{B-1}{B}$ otherwise. For simplicity, from now on we omit writing the clauses `$(~~)$' when applying a convolutional layer $C^{ij}$ and write $\circ$ instead, i.e., $C^{ij}\circ x := C^{ij}(x)$. Hence, by recursion we can write the $T$'th expansion of $f_W$,
\begin{equation*}
\begin{aligned}
f_W(x_{j_0};\tiX) ~&=~ \sum_{l_1 \in \textnormal{pred}(l_0)} C^{l_0l_1} \circ D^{j_0}_{l_1}\cdot \beta_{l_1} \\
% &\quad + \sum_{l_1 \in \textnormal{pred}(l_0)} \sum_{l_2 \in \textnormal{pred}(l_1)} C^{l_0l_1} \circ Q^{j_0}_{l_1} \odot C^{l_1l_2} \circ D^{j_0}_{l_2} \cdot \beta_{l_2}\\
&\quad+\sum^{T}_{t=2} \sum_{l_1 \in \textnormal{pred}(l_0)} \dots \sum_{l_{t+1}\in \textnormal{pred}(l_{t})} C^{l_0l_1} \circ Q^{j_0}_{l_1} \odot C^{l_1l_2}\circ \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{l_2} \odot C^{l_2l_3} \\
&\quad\quad\quad \cdots C^{l_{t-1} l_{t}}\circ \sum^{B}_{j_{t-1}=1} \alpha_{j_{t-2},j_{t-1}} Q^{j_{t-1}}_{l_{t}} \odot C^{l_{t}l_{t+1}} \circ \sum^{B}_{j_t = 1} \alpha_{j_{t-1},j_t} \cdot D^{j_{t}}_{l_{t+1}} \cdot \beta_{l_{t+1}} \\
&\quad+\sum_{l_1 \in \textnormal{pred}(l_0)} \sum_{l_2 \in \textnormal{pred}(l_1)} \sum_{l_3 \in \textnormal{pred}(l_2)} \dots \sum_{l_{T+2}\in \textnormal{pred}(l_{T+1})} C^{l_0l_1} \circ Q^{j_0}_{l_1} \odot C^{l_1l_2}\circ \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{l_2} \odot C^{l_2l_3} \\
&\quad\quad\quad \cdots C^{l_T l_{T+1}}\circ \sum^{B}_{j_{T}=1} \alpha_{j_{T-1},j_T} Q^{j_T}_{l_{T+1}} \odot C^{l_{T+1}l_{T+2}}\circ \Delta^{j_{T}}_{l_{T+2}}. \\
\end{aligned}
\end{equation*}
Instead of recursively expanding $f_W$'s expression $T$ times by rewriting $\Delta^{j_{t}}_{l_{t+2}}$ using \eqref{eq:delta}, we can expand through each term $\Delta^{j_{t}}_{l_{t+2}}$ based on different kinds of criteria. For instance, we can keep $\Delta^{j_{t}}_{l_{t+2}}$ unchanged or expand it based on a whether $(l_{t+2},l_{t+3})=(p,q)$ or not. Therefore, we can write $f_W(x_{j_0};\tiX)$ as the sum of multiplications along paths $\pi$ that start with $(p,q)$ and end by $L$ and paths that start with $1$ and end with $L$ that do not pass through $(p,q)$. Formally, let $j_0 \in [m]$ and $\pi$ be a path from some node to $\pi_T=L$ of length $T$,
\begin{equation*}
E^{j_0}_{1,\pi} ~:=~ \twopartdef
{C^{\pi_2\pi_1} \circ D^{j_0}_{\pi_1} \cdot \beta_{\pi_1}} {T=2}
%{C^{\pi_3\pi_2} \circ Q^{j_0}_{\pi_2} \odot C^{\pi_2\pi_1} \circ D^{j_0}_{\pi_1} \cdot \beta_{\pi_1}} {T=3}
{
\begin{aligned}
&C^{\pi_T\pi_{T-1}} \circ Q^{j_0}_{\pi_{T-1}} \odot C^{\pi_{T-1}\pi_{T-2}} \circ \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{\pi_{T-2}} \odot C^{\pi_{T-2}\pi_{T-3}}\\
&\quad \cdots C^{\pi_{3} \pi_{2}} \circ \sum^{B}_{j_{T-2}=1} \alpha_{j_{T-3},j_{T-2}} Q^{j_{T-2}}_{\pi_{1}} \odot C^{\pi_2 \pi_1} \circ \sum^{B}_{j_{T-1}=1}\alpha_{j_{T-2},j_{T-1}}\cdot D^{j_{T-1}}_{\pi_{1}} \cdot \beta_{\pi_1}
\end{aligned}
}
{T > 2}
\end{equation*}
and also
\begin{equation*}
E^{j_0}_{2,\pi} ~:=~ \twopartdef
{
\begin{aligned}
&C^{\pi_T \pi_{T-1}} \circ Q^{j_0}_{\pi_{T-1}} \odot C^{\pi_{T-1}\pi_{T-2}} \circ \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{\pi_{T-2}} \\
&\quad \cdots C^{\pi_3 \pi_2} \circ \sum^{B}_{j_{T-3}=1} \alpha_{j_{T-4},j_{T-3}} Q^{j_{T-3}}_{\pi_{2}} \odot C^{\pi_2 \pi_1} \circ \Delta^{j_{T-3}}_{\pi_1}
\end{aligned}
} {\pi_{1}\neq 1}
{
\begin{aligned}
&C^{\pi_T \pi_{T-1}} \circ Q^{j_0}_{\pi_{T-1}} \cdot C^{\pi_{T-1}\pi_{T-2}} \circ \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{\pi_{T-2}} \\
&\quad \cdots C^{\pi_3 \pi_2} \circ \sum^{B}_{j_{T-3}=1} \alpha_{j_{T-4},j_{T-3}} Q^{j_{T-3}}_{\pi_2} \odot C^{\pi_2 \pi_1} \circ x
\end{aligned}
}
{\pi_1=1}
\end{equation*}
With this terminology we can write
\begin{equation*}
\begin{aligned}
f_W(x_{j_0};\tiX) ~&=~
\sum_{\substack{ \textnormal{$\pi$ from some node $\neq 1$ to $L$} \\ \textnormal{$(p,q) \notin \pi$}}} E^{j_0}_{1,\pi}
+ \sum_{\substack{ \textnormal{$\pi$ from 1 to $L$} \\ \textnormal{$(p,q) \notin \pi$}}} E^{j_0}_{2,\pi}
+\sum_{\substack{\pi \textnormal{ from $q$ to $L$} \\ \pi_2=p,\pi_1=q}} E^{j_0}_{2,\pi} ~=:~ Z^{j_0}_1+Z^{j_0}_2+Z^{j_0}_3.
\end{aligned}
\end{equation*}
We note that for any $x$ and $\tiX$, with measure $1$ over $W$, the matrices $D_{l}(x;\tiX)$ are constant in the neighborhood of $W$. Together with the assumption that $A_{l}(x;\tiX)$ are detached, the pseudo-gradients of the matrices $Q^j_l$ with respect to $V^{pq}$ are zero. Therefore, since the terms $Z^{j_0}_1$ and $Z^{j_0}_2$ do not involve $V^{pq}$, we have $\frac{\bar{\partial} f_W(x_{j_0};\tiX)}{\partial V^{pq}} = \frac{\bar{\partial} Z^{j_0}_3}{\partial V^{pq}}$. As a next step, we represent $\frac{\bar{\partial} Z^{j_0}_2}{\partial V^{pq}}$. For simplicity, we assume that each $\pi$ from $q$ to $L$ that crosses through $(p,q)$ is of length $T(\pi)\geq 3$. With a slight abuse of notation, we denote $T:=T(\pi)$. We can write
% \begin{equation*}
% \begin{aligned}
% Z^{j_0}_3 ~&=~ \sum_{\substack{\pi \textnormal{ from some node to $L$} \\ \pi_2=p,\pi_1=q}} C^{\pi_T\pi_{T-1}} \circ Q^{j_0}_{\pi_{T-1}} \odot C^{\pi_{T-1}\pi_{T-2}} \circ \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{\pi_{T-2}} \odot C^{\pi_{T-2}\pi_{T-3}} \\
% &\quad\quad\quad\quad\quad \cdots \sum^{B}_{j_{T-4}=1} \alpha_{j_{T-5},j_{T-4}} Q^{j_{T-4}}_{\pi_{3}} \odot C^{\pi_{3} \pi_{2}} \circ \sum^{B}_{j_{T-3}=1} \alpha_{j_{T-4},j_{T-3}} Q^{j_{T-3}}_{\pi_{2}} \odot C^{\pi_2 \pi_1} \circ D^{j_{T-2}}_{\pi_{1}} \cdot \beta_{\pi_1} \\
% ~&=~ \sum_{\substack{\pi \textnormal{ from $q$ to $L$} \\ \pi_2=p,\pi_1=q}} W^{\pi_T\pi_{T-1}} \cdot Q^{j_0}_{\pi_{T-1}} \cdot W^{\pi_{T-1}\pi_{T-2}} \cdot \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{\pi_{T-2}} \cdot W^{\pi_{T-2}\pi_{T-3}} \\
% &\quad\quad\quad\quad\quad \cdots \sum^{B}_{j_{T-3}=1} \alpha_{j_{T-4},j_{T-3}} Q^{j_{T-3}}_{\pi_{2}} \cdot W^{\pi_2\pi_1} \cdot \sum_{l \in \textnormal{pred}(\pi_1)}\sum^{B}_{j_{T-2}=1} \alpha_{j_{T-3},j_{T-2}} Q^{j_{T-2}}_{l} \cdot W^{\pi_1 l} \cdot D^{j_{T-2}}_{l} \cdot \beta_{l} \\
% ~&=~ \sum^{B}_{j_{T-3}=1} \left[\sum_{\substack{\pi \textnormal{ from $q$ to $L$} \\ \pi_2=p,\pi_1=q}} (u^{\pi,1}_{j_0,j_{T-3}})^{\top} \right] \cdot W^{pq} \cdot \sum_{l \in \textnormal{pred}(q)}\sum^{B}_{j_{T-2}=1} \alpha_{j_{T-3},j_{T-2}} Q^{j_{T-2}}_{q} \cdot W^{ql} \cdot D^{j_{T-2}}_{q} \cdot \beta_{q}\\
% ~&=~ \sum^{B}_{j_{T-3}=1} (u^{1}_{j_0,j_{T-3}})^{\top} \cdot W^{pq} \cdot u^{2}_{j_{T-3}},
% \end{aligned}
% \end{equation*}
\begin{equation*}
\begin{aligned}
Z^{j_0}_3 ~&=~ \sum_{\substack{\pi \textnormal{ from $p$ to $L$} \\ \pi_2=p,\pi_1=q}} C^{\pi_T \pi_{T-1}} \circ Q^{j_0}_{\pi_{T-1}} \odot C^{\pi_{T-1}\pi_{T-2}} \circ \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{\pi_{T-2}} \cdots C^{\pi_3 \pi_2} \circ \sum^{B}_{j_{T-3}=1} \alpha_{j_{T-4},j_{T-3}} Q^{j_{T-3}}_{\pi_{2}} \odot C^{\pi_2 \pi_1} \circ \Delta^{j_{T-3}}_{\pi_1} \\
~&=~ \sum_{\substack{\pi \textnormal{ from $p$ to $L$} \\ \pi_2=p,\pi_1=q}} C^{\pi_T \pi_{T-1}} \circ Q^{j_0}_{\pi_{T-1}} \odot C^{\pi_{T-1}\pi_{T-2}} \circ \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{\pi_{T-2}} \cdots C^{\pi_3 \pi_2} \circ \sum^{B}_{j_{T-3}=1} \alpha_{j_{T-4},j_{T-3}} Q^{j_{T-3}}_{\pi_2} \odot C^{pq} \circ \Delta^{j_{T-3}}_{q}. \\
%~&=~ \sum^{B}_{j_{T-3}=1} (z^{1}_{j_0,j_{T-3}})^{\top} \cdot W^{pq} \cdot z^{2}_{j_{T-3}},
\end{aligned}
\end{equation*}
As a next step we linearize the above expression. We notice that $\vect(C^{ij} \circ z) = U^{ij} \cdot \vect^{ij}(z) = W^{ij} \cdot \vect(z)$ and $\vect(Q^{j}_l \odot z) = \textnormal{diag}(\vect(Q^{j}_{l})) \cdot \vect(z)$. Therefore, we replace $C^{pq}$ by $U^{pq}$ and $\Delta^{j_{T-3}}_q$ with $\vect^{pq}(\Delta^{j_{T-3}}_q)$. In addition, each one of the convolutional layers $C^{ij}$ above $C^{pq}$ is replaced with $W^{ij}$ and each application of element-wise multiplication by $Q^{j}_l$ is replace with a (right side) matrix multiplication by $P^{j}_{l} := \textnormal{diag}(\vect(Q^{j}_{l}))$. We obtain that
\begin{equation*}
\begin{aligned}
Z^{j_0}_3 ~&=~ \sum_{\substack{\pi \textnormal{ from $p$ to $L$} \\ \pi_2=p,\pi_1=q}} W^{\pi_T \pi_{T-1}} \cdot P^{j_0}_{\pi_{T-1}} \cdot W^{\pi_{T-1}\pi_{T-2}} \cdot \sum^{B}_{j_1=1} \alpha_{j_0,j_1} P^{j_1}_{\pi_{T-2}} \cdots W^{\pi_3 \pi_2} \cdot \sum^{B}_{j_{T-3}=1} \alpha_{j_{T-4},j_{T-3}} P^{j_{T-3}}_{\pi_2} \cdot U^{pq} \cdot \vect^{pq}(\Delta^{j_{T-3}}_{q}) \\
~&=~ \sum^{B}_{j_{T-3}=1} (z^{1}_{j_0,j_{T-3}})^{\top} \cdot U^{pq} \cdot z^{2}_{j_{T-3}},
\end{aligned}
\end{equation*}
where
\begin{equation*}
\begin{aligned}
(z^{1}_{j_0,j_{T-3}})^{\top} ~:&=~ \sum_{\substack{\pi \textnormal{ from $p$ to $L$} \\ \pi_2=p,\pi_1=q}} W^{\pi_T \pi_{T-1}} \cdot P^{j_0}_{\pi_{T-1}} \cdot W^{\pi_{T-1}\pi_{T-2}} \cdot \sum^{B}_{j_1=1} \alpha_{j_0,j_1} P^{j_1}_{\pi_{T-2}} \cdots W^{\pi_3 \pi_2} \cdot \alpha_{j_{T-4},j_{T-3}} P^{j_{T-3}}_{\pi_2} \\
z^{2}_{j_{T-3}} ~:&=~ \vect^{pq}(\Delta^{j_{T-3}}_{q}).
\end{aligned}
\end{equation*}
We notice that $\frac{\bar{\partial} z^{1}_{j_0,j_{T-3}}}{\partial V^{pq}} = 0$ since it is decomposed of $W^{ij}$ for $(i,j)\neq (p,q)$ and $P^{j}_l$ which are reshapes of $Q^j_l = D^j_l \odot A_l$, where $\frac{\partial D^j_l}{\partial V^{pq}}=0$ since $D^j_l$ are locally constant as functions of $W$ and $\frac{\bar{\partial} A_l}{\partial V^{pq}}=0$ by assumption. Furthermore, $\frac{\partial z^{2}_{j_{T-3}}}{\partial V^{pq}}=0$ since $\Delta^{j_{T-3}}_q = v^{j_{T-3}}_q - \mu_q$, where $v_q$ is a layer below $C^{pq}$ (so its outputs are independent of $V^{pq}$). We note that $U^{pq}$ is a block matrix with $N^{pq}$ blocks of $V^{pq}$. Therefore, by we can write
\begin{equation*}
Z^{j_0}_{3} ~=~ \sum^{B}_{j_{T-3}=1} \sum^{N^{pq}}_{i=1} a^{\top}_{i,j_0,j_{T-3}} \cdot V^{pq} \cdot b_{i,j_{T-3}},
\end{equation*}
where $z^1_{j_0,j_{T-3}}= (a_{1,j_0,j_{T-3}} \| \dots \| a_{N^{pq},j_{T-3}})$ and $z^{2}_{j_{T-3}} = (b_{1,j_{T-3}} \| \dots \| b_{N^{pq},j_{T-3}})$ are partitions of the vectors $z^1_{j_0,j_{T-3}}$ and $z^2_{j_{T-3}}$ into $N^{pq}$ consecutive sub-vectors of the same size. In conclusion, we have
\begin{equation*}
\frac{\bar{\partial} f(x_{j_0};\tiX)}{\partial V^{pq}} ~=~ \frac{\partial Z^{j_0}_{3}}{\partial V^{pq}} ~=~ \sum^{B}_{j_{T-3}=1} \sum^{N^{pq}}_{i=1} a_{i,j_0,j_{T-3}} \cdot b^{\top}_{i,j_{T-3}}
\end{equation*}
which is a matrix of rank $\leq BN^{pq}$.
% Therefore, we can write $Z^{j_0}_2 = \sum^{B}_{j_{T-3}=1} (u^1_{j_0,j_{T-3}})^{\top} \cdot W^{pq} \cdot u^2_{j_{T-3}}$ and we obtain that $\frac{\bar{\partial} Z^{j_0}_2}{\partial W^{pq}} = \sum^{B}_{j_{T-3}=1} (u^1_{j_0,j_{T-3}})\cdot (u^2_{j_{T-3}})^{\top}$.
% Next, we would like to analyze the pseudo gradient $\frac{\bar{\partial} Z^{j_0}_4}{\partial W^{pq}}$. We can write
We also notice that $b_{i,j_{T-3}}$ are independent of $j_0$. Therefore, we obtain that
\begin{equation*}
\begin{aligned}
\sum^{B}_{j_0=1}\alpha_{j_0} \cdot \frac{\bar{\partial} f_W(x_{j_0};\tiX)}{\partial V^{pq}}
~&=~ \sum^{B}_{j_0=1} \alpha_{j_0}\cdot \frac{\bar{\partial} Z^{j_0}_3}{\partial
V^{pq}} \\
~&=~ \sum^{B}_{j_0=1} \alpha_{j_0} \sum^{B}_{j_{T-3}=1} \sum^{N^{pq}}_{i=1} a_{i,j_0,j_{T-3}} \cdot b^{\top}_{i,j_{T-3}}\\
~&=~ \sum^{B}_{j_{T-3}=1} \sum^{N^{pq}}_{i=1} \left(\sum^{B}_{j_0=1} \alpha_{j_0} \cdot a_{i,j_0,j_{T-3}}\right) \cdot b^{\top}_{i,j_{T-3}},
\end{aligned}
\end{equation*}
which is also a matrix of rank $\leq BN^{pq}$.
\end{proof}
\end{document}
%%%%%%%% ICML 2023 EXAMPLE LATEX SUBMISSION FILE %%%%%%%%%%%%%%%%%
\documentclass{article}
% Recommended, but optional, packages for figures and better typesetting:
\usepackage{microtype}
\usepackage{graphicx}
\usepackage{subfigure}
\usepackage{booktabs} % for professional tables
%\usepackage{subfig}
% hyperref makes hyperlinks in the resulting PDF.
% If your build breaks (sometimes temporarily if a hyperlink spans a page)
% please comment out the following usepackage line and replace
% \usepackage{icml2023} with \usepackage[nohyperref]{icml2023} above.
\usepackage{hyperref}
% Attempt to make hyperref and algorithmic work together better:
\newcommand{\theHalgorithm}{\arabic{algorithm}}
% Use the following line for the initial blind version submitted for review:
%\usepackage{icml2023}
% If accepted, instead use the following line for the camera-ready submission:
\usepackage[accepted]{icml2023_arxiv}
% For theorems and such
\usepackage{amsmath}
\usepackage{mathtools}
\usepackage{amsthm}
\usepackage{amsmath,amssymb,amsthm,bbm,thm-restate}
\input{math_commands.tex}
% if you use cleveref..
\usepackage[capitalize,noabbrev]{cleveref}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% THEOREMS
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\theoremstyle{plain}
\newtheorem{theorem}{Theorem}[section]
\newtheorem{proposition}[theorem]{Proposition}
\newtheorem{lemma}[theorem]{Lemma}
\newtheorem{corollary}[theorem]{Corollary}
\theoremstyle{definition}
\newtheorem{definition}[theorem]{Definition}
\newtheorem{assumption}[theorem]{Assumption}
\theoremstyle{remark}
\newtheorem{remark}[theorem]{Remark}
\newcommand{\expnumber}[2]{{#1}\mathrm{e}{#2}}
\newcommand{\cC}{\mathcal{C}}
\newcommand{\cD}{\mathcal{D}}
\newcommand{\cF}{\mathcal{F}}
\newcommand{\cG}{\mathcal{G}}
\newcommand{\cH}{\mathcal{H}}
\newcommand{\cN}{\mathcal{N}}
\newcommand{\cP}{\mathcal{P}}
\newcommand{\cU}{\mathcal{U}}
\newcommand{\cX}{\mathcal{X}}
\newcommand{\cY}{\mathcal{Y}}
\renewcommand{\P}{\mathbb{P}}
\newcommand{\B}{\mathbb{B}}
\newcommand{\bI}{\mathbb{I}}
\newcommand{\EE}[1]{\E\left[ #1\right]}
\newcommand{\PP}[1]{\P\left[ #1\right]}
\newcommand{\I}[1]{\bI\left[ #1\right]}
\newcommand{\tiX}{\tilde{X}}
\newcommand{\tiS}{\tilde{S}}
\newcommand{\ty}{\tilde{y}}
\newcommand{\tcG}{\tilde{\cG}}
\newcommand{\tcP}{\tilde{\cP}}
\newcommand{\tcU}{\tilde{\cU}}
\newcommand{\tf}{\tilde{f}}
\newcommand{\tig}{\tilde{g}}
\newcommand{\tih}{\tilde{h}}
\newcommand{\clD}{\cD_{\cC}}
\newcommand{\one}{\mathbbm{1}}
\newcommand{\rank}{\textnormal{rank}}
\newcommand{\vect}{\textnormal{vec}}
\newcommand{\threepartdef}[6]
{
\left\{
\begin{array}{lll}
#1 & \mbox{if } #2 \\
#3 & \mbox{if } #4 \\
#5 & \mbox{if } #6
\end{array}
\right.
}
\newcommand{\twopartdef}[4]
{
\left\{
\begin{array}{lll}
#1 & \mbox{if } #2 \\
#3 & \mbox{if } #4 \\
\end{array}
\right.
}
\def\fr#1#2{{\textstyle\frac{#1}{#2}}}
% Todonotes is useful during development; simply uncomment the next line
% and comment out the line below the next line to turn off comments
%\usepackage[disable,textsize=tiny]{todonotes}
\usepackage[textsize=tiny]{todonotes}
% The \icmltitle you define below is probably too long as a header.
% Therefore, a short form for the running title is supplied here:
\icmltitlerunning{SGD and Weight Decay Provably Induce a Low-Rank Bias in Deep Neural Networks}
\begin{document}
\twocolumn[
\icmltitle{SGD and Weight Decay Provably Induce \\ a Low-Rank Bias in Deep Neural Networks}
% It is OKAY to include author information, even for blind
% submissions: the style file will automatically remove it for you
% unless you've provided the [accepted] option to the icml2023
% package.
% List of affiliations: The first argument should be a (short)
% identifier you will use later to specify author affiliations
% Academic affiliations should list Department, University, City, Region, Country
% Industry affiliations should list Company, City, Region, Country
% You can specify symbols, otherwise they are numbered in order.
% Ideally, you should not use this facility. Affiliations will be numbered
% in order of appearance and this is the preferred way.
\icmlsetsymbol{equal}{*}
\begin{icmlauthorlist}
\icmlauthor{Tomer Galanti}{MIT}
\icmlauthor{Zachary S. Siegel}{P}
\icmlauthor{Aparna Gupte}{MIT}
\icmlauthor{Tomaso Poggio}{MIT}
\end{icmlauthorlist}
\icmlaffiliation{MIT}{Massachusetts Institute of Technology}
\icmlaffiliation{P}{Princeton University}
\icmlcorrespondingauthor{Tomer Galanti}{galanti@mit.edu}
% You may provide any keywords that you
% find helpful for describing your paper; these are used to populate
% the "keywords" metadata in the PDF but will not be shown in the document
\icmlkeywords{Machine Learning, ICML}
\vskip 0.3in
]
% this must go after the closing bracket ] following \twocolumn[ ...
% This command actually creates the footnote in the first column
% listing the affiliations and the copyright notice.
% The command takes one argument, which is text to display at the start of the footnote.
% The \icmlEqualContribution command is standard text for equal contribution.
% Remove it (just {}) if you do not need this facility.
%\printAffiliationsAndNotice{} % leave blank if no need to mention equal contribution
\printAffiliationsAndNotice{\icmlEqualContribution} % otherwise use the standard text.
\begin{abstract}
In this paper, we study the bias of Stochastic Gradient Descent (SGD) to learn low-rank weight matrices when training deep ReLU neural networks. Our results show that training neural networks with mini-batch SGD and weight decay causes a bias towards rank minimization over the weight matrices. Specifically, we show, both theoretically and empirically, that this bias is more pronounced when using smaller batch sizes, higher learning rates, or increased weight decay. Additionally, we predict and observe empirically that weight decay is necessary to achieve this bias. Finally, we empirically investigate the connection between this bias and generalization, finding that it has a marginal effect on generalization. Our analysis is based on a minimal set of assumptions and applies to neural networks of any width or depth, including those with residual connections and convolutional layers.
\end{abstract}
\section{Introduction}
Stochastic gradient descent (SGD) is a widely used optimization technique for training deep learning models~\citep{bottou-91c}. While it was initially developed to address the computational challenges of gradient descent, recent studies suggest that SGD also provides regularization that prevents overparameterized models from converging to minima that do not generalize well~\citep{DBLP:journals/corr/ZhangBHRV16,https://doi.org/10.48550/arxiv.1711.04623,https://doi.org/10.48550/arxiv.1609.04836,zhu}. For instance, empirical studies have shown that SGD outperforms gradient descent~\citep{zhu} and that smaller batch sizes result in better generalization~\citep{NIPS2017_a5e0ff62,https://doi.org/10.48550/arxiv.1609.04836}. However, the full range of regularization effects induced by SGD is not yet fully understood.
One area of recent research focuses on characterizing the implicit regularization of gradient-based optimization and its relationship to generalization in deep learning. Several papers have examined the potential bias of gradient descent or stochastic gradient descent toward rank minimization. Empirically, it was shown~\citep{NIPS2014_2afe4567,10.5555/3294771.3294853,s21165599,8099498,pmlr-v80-arora18b} that replacing weight matrices with low-rank approximations results in only a small drop in accuracy, suggesting that the weight matrices at convergence may be close to low-rank matrices. Following this line of work, various attempts were made to understand the origins of this low-rank bias and its potential relation with generalization.
Initially, it was believed that the implicit regularization in matrix factorization could be characterized in terms of the nuclear norm of the corresponding linear predictor~\citep{https://doi.org/10.48550/arxiv.1705.09280}. This conjecture was later refuted~\citep{DBLP:journals/corr/abs-2012-09839}. Subsequent conjecture posits that rank minimization may play a key role in explaining generalization in deep learning. For instance,~\cite{DBLP:journals/corr/abs-2005-06398} conjectured that the implicit regularization in matrix factorization can be explained by rank minimization, and also hypothesized that some notion of rank minimization may be key to explaining generalization in deep learning. Additionally,~\cite{DBLP:journals/corr/abs-2012-09839} established evidence that the implicit regularization in matrix factorization is a heuristic for rank minimization. Beyond factorization problems,~\cite{DBLP:journals/corr/abs-2006-06657} showed that gradient flow (GF) training of univariate linear networks with respect to exponentially-tailed classification losses learns weight matrices of rank 1. Intuitively, such networks generalize well due to their effectively limited capacities.
With nonlinear neural networks, the origin of this bias and its connection with generalization is less clear. Several papers~\citep{https://doi.org/10.48550/arxiv.2202.00856,le2022training} studied low-rank bias in linear layers at the top of a neural network. For instance,~\citep{le2022training} analyzes low-rank bias in neural networks trained with gradient flow (GF) without regularization. While this paper makes significant strides in extending the analysis in~\citep{DBLP:journals/corr/abs-2006-06657}, it makes several limiting assumptions. As a result, their analysis is only applicable under very specific conditions, such as when the data is linearly separable, and their low-rank analysis is limited to a set of linear layers aggregated at the top of the trained network. Later,~\cite{DBLP:journals/corr/abs-2201-12760} showed that for ReLU networks, GF generally does not minimize rank. They also argued that sufficiently deep ReLU networks exhibits low-rank solutions under $L_2$ norm minimization. This interesting result, however, applies only to the global minima and only to layers added to a pre-existing network that is capable of solving the problem.
Despite the recent progress in understanding the low-rank bias in deep networks, a complete understanding of its origins and its relationship with different hyperparameters is largely missing.
{\bf Contributions.\enspace} In this paper, we show that using mini-batch stochastic gradient descent (SGD) and weight decay effectively regularize the rank of the learned weight matrices during the training of neural networks. Our theoretical analysis predicts that smaller batch sizes, higher learning rates, or increased weight decay results in a decrease in the rank of the learned matrices, and that regularization is necessary to achieve this bias. The scope of the analysis is fairly general, covering deep ReLU networks trained with mini-batch SGD for minimizing a differentiable loss function with $L_2$ regularization (i.e., weight decay). The neural networks may include fully-connected layers, residual connections, and convolutional layers.
In addition to our theoretical analysis, we provide a comprehensive empirical study in which we examine the effects of different hyperparameters on the rank of weight matrices for various network architectures. Additionally, we carried out several experiments to examine the connection between low-rank bias and generalization. The results indicate that while low-rank bias is not a requirement for good generalization, it is correlated with a marginal improvement in performance.
\section{Problem Setup}\label{sec:setup}
In this paper, we study the influence of using mini-batch stochastic gradient descent (SGD) in conjunction with weight decay on the inductive biases of neural networks in standard supervised learning settings. The task at hand is defined by a distribution $P$ over samples $(x,y)\in \cX\times \cY$, where $\cX \subset \R^{c_{1} \times h_1 \times w_1}$ is the space of instances (e.g., images), and $\cY \subset \R^k$ is the label space.
We consider a parametric model $\cF \subset \{f': \cX \to \R^k\}$, where each function $f_W \in \cF$ is specified by a vector of parameters $W \in \R^{N}$. The function $f_W \in \cF$ assigns a prediction to any input point $x \in \cX$, and its performance is measured by the Expected Risk,
\begin{equation*}
L_P(f_W)~:=~\E_{(x,y)\sim P}[\ell(f_W(x),y)],
\end{equation*}
where $\ell:\R^k \times \cY \to [0,\infty)$ is a non-negative, differentiable, loss function (e.g., MSE or cross-entropy losses). For simplicity, in the analysis we assume that $k=1$.
Since we do not have direct access to the full population distribution $P$, the goal is to learn a predictor, $f_W$, from a training dataset $S = \{(x_i,y_i)\}_{i=1}^m$ of independent and identically distributed (i.i.d.) samples drawn from $P$. To avoid overfitting the training data, we typically use weight decay to control the complexity of the learned model. Specifically, we aim to minimize the Regularized Empirical Risk,
\begin{equation*}
L^{\lambda}_S(f_W)~:=~\frac{1}{m}\sum^{m}_{i=1}\ell(f_W(x_i),y_i) + \lambda \|W\|^2_2,
\end{equation*}
where $\lambda > 0$ is a predefined hyperparameter and $\|\cdot\|_2$ is the Frobenius norm. To accomplish this task, we typically use mini-batch SGD, as outlined in the following paragraph.
{\bf Optimization.\enspace} In this study, we employ stochastic gradient descent (SGD) to minimize the regularized empirical risk $L^{\lambda}_{S}(f_{W})$ over a specified number of iterations $T$. We begin by initializing $W_1$ using a standard initialization method, and then update $W_t$ for $T$ iterations, ultimately returning $W_T$. At each iteration $t$, we randomly select a batch $\tiS_t \subset S$ of $B$ samples, and update $W_{t+1} = W_t - \mu \nabla_W L^{\lambda}_{\tiS_t}(f{W_t})$, where $\mu > 0$ is the predefined learning rate.
{\bf Notation.\enspace} In this paper, we use the following notations. For an integer $k\geq 1$, we use the notation $[k]=\{1,\ldots,k\}$. The Euclidean norm of a vector $z \in \R^{n}$ is denoted by $\|z\|:=\sqrt{\sum^{n}_{i=1}z^2_i}$. For two vectors $x \in \R^{n}, y \in \R^{m}$ we define their concatenation as follows $(x\|y) := (x_1,\dots,x_n,y_1,\dots,y_m) \in \R^{n+m}$. For a given matrix $A \in \R^{n \times m}$, we denote $A_{i,:}$ its $i$th row and by $\vect(A) := (A_{1}\|\dots \| A_n)$ its vectorization. For a given tensor $A \in \R^{c \times h\times w}$, we denote by $\vect(A) := (\vect(A_1)\| \dots \| \vect(A_c))$ the vectorized form of $A$. Tensor slicing is defined as, $x_{a:b} := (x_a,\dots,x_b)$.
\subsection{Architectures}\label{sec:archs}
In this work, the function $f_W$ represents a neural network, consisting of a set of layers of weights interlaced with ReLU activation units. Our definition of a neural network is fairly generic, including convolutional layers, pooling layers, residual connections, and fully-connected layers.
{\bf Network architecture.\enspace} Formally, a neural network $f_W$ can be described as a directed acyclic graph (DAG) $G=(V,E)$, where $V= \{v_1,\dots,v_L\}$ consists of the various layers of the network, and each edge $e_{ij}=(v_i,v_j) \in E$ represents a connection between two layers. Each layer is a function $v_{i}:\R^{c_{1} \times h_1 \times w_1} \to \R^{c_{i} \times h_i \times w_i}$, and each connection $(v_i,v_j)$ holds a transformation $C^{ij} : \R^{c_{j} \times h_j \times w_j} \to \R^{c_i \times h_i \times w_i}$. The layers are divided into three categories: the input layer $v_1$, the output layer $v_L$, and intermediate layers. There are no connections directed towards the input layer or out of the output layer (i.e., $\forall i \in [L]:(v_i,v_L), (v_0,v_i) \notin E$). Given an input $x \in \R^{c_i \times h_i \times w_i}$, the output of a given layer $v_i$ is evaluated as follows $v_i(x) := \sigma(\sum_{j \in \textnormal{pred}(i)} C^{ij}(v_{j}(x)))$, except for the output layer $v_L$ that computes $f_W(x) := v_L(x) := \sum_{j \in \textnormal{pred}(L)}C^{Lj} (v_{j}(x))$. Here, $\textnormal{pred}(i):=\{j \in [L]\mid (v_i,v_j)\in E\}$, $\textnormal{succ}(i):=\{j \in [L]\mid (v_j,v_i)\in E\}$ denote the sets of predecessor and successor layers of the $i$th layer and $\sigma$ is the element-wise ReLU activation function. Each transformation $C^{ij}$ is either trainable (e.g., a convolutional layer) or a constant affine transformation (e.g., a residual connection). The set of trainable connections is denoted by $E_T$. In this paper, we consider the following types of layers.
{\bf Convolutional layers.\enspace} A convolutional layer~\citep{726791} (see also~\cite{Goodfellow-et-al-2016}), commonly used in image processing tasks, is defined by a kernel tensor $Z^{ij} \in \R^{c_{i}\times c_j \times k_1 \times k_2}$, where $c_j$, $c_i$, $k_1$, and $k_2$ represent the number of input and output channels and the kernel sizes respectively. The layer applies the kernel tensor to the input tensor by sliding it across the input tensor with a specified stride length, $s$, after zero-padding the input tensor with $p$ rows in each ``side'' of the tensor. The output tensor $y \in \R^{c_i \times h_i \times w_i}$ is computed by summing up the element-wise product of the kernel tensor and the corresponding section of the padded input tensor at each position of the sliding. Formally, for all indices $(c,t,l) \in [c_i]\times [h_i] \times [w_i]$,
\begin{small}
\begin{equation*}%\label{eq:conv}
\begin{aligned}
y_{c,t,l} ~&=~ \sum^{c_{j}}_{c'=1} \vect(Z_{c,c',~:~})^{\top} \\
&\quad\quad\quad\quad \cdot \vect(\textnormal{Pad}_p(x)_{c',ts~:~(t+1)s+k_1 , ls~:~(l+1)s+k_2}).
\end{aligned}
\end{equation*}
\end{small}
Here, $\textnormal{Pad}_p$ takes a tensor $x \in \R^{c_j \times h_j \times w_j}$ and returns a new tensor $x' \in \R^{c_j \times (h_j+2p) \times (w_j+2p)}$, where the first and last $p$ rows and columns of each channel $x'_{c,~:~,~:~}$ are zeros and the middle $1 \times h_j \times w_j$ tensor is equal to $x_{c,~:~,~:~}$. The output dimensions $h_i$ and $w_i$ are calculated using the formulas $h_i = \left(\lfloor(h_j - k_j +2p)\rfloor/s+1\right)$ and $w_i = \left(\lfloor(w_j - k_2 +2p)\rfloor/s+1\right)$.
We can also represent the convolutional layer as a linear operation by defining a matrix $V^{ij} \in \R^{c_ih_iw_i \times c_jh_jw_j}$, which computes the output of the layer for a given input vectorized as a column vector, and a matrix $W^{ij} \in \R^{c_{i} \times c_{j}k_1k_2}$, which has the vectorized filters as its rows, $W^{ij}_c := \vect(Z^{ij}_{c,~:~,~:~})$. This allows us to express the convolutional layer as a linear operator, making it possible to analyze its properties mathematically.
% A convolutional layer~\citep{726791} (see also~\cite{Goodfellow-et-al-2016}) $C^{ij}:\R^{c_j \times h_j \times w_j} \to \R^{c_i \times h_i \times w_i}$ with kernel sizes $(k_1,k_2)$, padding $p$ and stride $s$ is parameterized by a tensor $Z^{ij} \in \R^{c_{i}\times c_j \times k_1 \times k_2}$ and computes the following tensor as output
% \begin{small}
% \begin{equation}\label{eq:conv}
% \forall~ (c,t,l) \in [c_i]\times [h_i] \times [w_i]:~y_{c,t,l} ~=~ \sum^{c_{j}}_{c'=1} \vect(Z_{c,c',~:~})^{\top} \cdot \vect(\textnormal{Pad}_p(x)_{c',ts~:~(t+1)s+k_1 , ls~:~(l+1)s+k_2}).
% \end{equation}
% \end{small}
% Here, $\textnormal{Pad}_p$ takes a tensor $x \in \R^{c \times h \times w}$ and returns a new tensor $x' \in \R^{c \times (h+2p) \times (w+2p)}$, where the first and last $p$ rows and columns of each channel $x'_{c,~:~,~:~}$ are zeros and the middle $1 \times h \times w$ tensor is equal to $x_{c,~:~,~:~}$. The formulas for the output dimensions are $h_2 = \left(\lfloor(h_1 - k_1 +2p)\rfloor/s+1\right)$ and $w_2 = \left(\lfloor(w_1 - k_2 +2p)\rfloor/s+1\right)$.
% For a given convolutional layer $C^{ij}$ with weights $Z^{ij}$, we define a matrix $V^{ij} \in \R^{c_ih_iw_i \times c_jh_jw_j}$ that computes $\forall~ x \in \R^{c_j \times h_j \times w_j}:~ V^{ij} \vect(x) ~=~ \vect(C^{ij}(x))$.
% This matrix exists since both the padding and convolution operations can be represented as linear operations of the input. We also consider $W^{ij} \in \R^{c_{i} \times c_{j}k_1k_2}$ which is a matrix whose $c$'th row is the vectorized filter $W^{ij}_c := \vect(Z^{ij}_{c,~:~,~:~})$.
{\bf Fully-connected layers.\enspace} As a special case of convolutional layers, the network may also include fully-connected layers. A fully-connected layer $F:\R^{c_j} \to \R^{c_i}$, associated with a matrix $W \in \R^{c_i \times c_j}$, can be represented as a $1\times 1$ convolutional layer $C:\R^{c_j \times 1 \times 1}\to \R^{c_i \times 1 \times 1}$ with $k_1=k_2=1$, $p=0$ and $s=1$. The parameters tensor $Z \in \R^{c_i \times c_j \times 1 \times 1}$ satisfies $Z_{a,b,1,1}=W_{a,b}$ for all $(a,b)\in [c_i]\times [c_j]$, and the layer satisfies $\vect(C(x)) = W \vect(x)$.
{\bf Pooling layers.\enspace} A pooling layer~\citep{Zhou1988ComputationOO} (see also~\citep{Goodfellow-et-al-2016}) $C$ with kernel dimensions $(k_1,k_2)$ stride $s$ and padding $p$ takes an input $x \in \mathbb{R}^{c_j \times h_j \times w_j}$ and computes an output $y \in \mathbb{R}^{c_i \times h_i \times w_i}$ with $c_i=c_j$ channels, and dimensions $h_i = \left(\lfloor(h_j - k_1 +2p)\rfloor/s+1\right)$ and $w_i = \left(\lfloor(w_j - k_2 +2p)\rfloor/s+1\right)$. The output of each pooling layer is computed as follows:
\begin{equation*}%\label{eq:pool}
y_{c,t,l} ~=~ \textnormal{op}(\textnormal{Pad}_p(x)_{c,ts~:~(t+1)s+k_1,ls~:~(l+1)s+k_2}),
\end{equation*}
where $\textnormal{op}$ is either the maximum or average operator and $(c,t,l) \in [c_i]\times [h_i] \times [w_i]$.
{\bf Rearrangement layers.\enspace} To easily switch between convolutional and fully-connected layers, we should be able to represent tensor layers as vectors and vice versa. To change the shape of a specific layer, we use rearrangement layers. A rearrangement layer $C^{ij}:\R^{c_j \times h_j \times w_j} \to \R^{c_i \times h_i \times w_i}$ takes an input vector $x \in \R^{c_j \times h_j \times w_j}$ and rearranges its coordinates in a different shape and order. Formally, it returns a vector $(x_{\pi(k)})_{k \in [c_j]\times [h_j] \times [w_j]}$, where $\pi:[c_j]\times [h_j] \times [w_j] \to [c_i]\times [h_i] \times [w_i]$ is invertible (in particular $c_ih_iw_i=c_jh_jw_j$).
\section{Theoretical Results}\label{sec:theory}
In this section, we present our main theoretical result. We show that when training a ReLU neural network with SGD, the weight matrices tend to be close to matrices of a bounded rank. Specifically, with a simple observation (proved in Appendix~\ref{app:proofs}) that the number of input patches $N^{ij}$ of a certain convolutional layer $C^{ij}$ sets an upper bound on the rank of the gradient of the network with respect to the parameters matrix $W^{ij}$. By recursively unrolling the optimization, we express the weight matrix $W^{ij}_{t}$ as a sum of $(1-\mu\lambda)^k W^{ij}_{t-k}$ and $kB$ gradients of the loss function with respect to $W^{ij}$ for different samples at different iterations. Since each one of these terms is a matrix of rank $\leq N^{ij}$, we conclude that the distance between $W^{ij}_{t}$ and a matrix of rank $\leq N^{ij}Bk$ decays exponentially with increasing $k$.
\begin{restatable}{lemma}{ranklemma}\label{lem:rank}
Let $f_W$ be a neural network and let $C^{ij}$ be a convolutional layer within $f_W$ with parameters matrix $W^{ij}$. Then, $\rank\left(\nabla_{W^{ij}} f_W(x)\right)\leq N^{ij}$.
\end{restatable}
Interestingly, we observe particularly degenerate gradients for fully-connected layers. As discussed in Sec.~\ref{sec:archs}, for a fully-connected layer $C^{ij}:\R^{c_j \times 1 \times 1} \to \R^{c_i \times 1 \times 1}$ we have $N^{ij}=1$, and thus, $\rank\left(\nabla_{W^{ij}} f_W(x)\right) \leq 1$. To demonstrate this observation, we provide a simple proof for the case of fully-connected networks.
\begin{lemma}\label{lem:simple}
Let $f_W(x) = W^L \sigma(W^{L-1}\cdots \sigma(W_1 x)\cdots)$ be a neural network, where $W^l \in \mathbb{R}^{d_{l+1}\times d_l}$ for all $l \in [L]$ and $\sigma$ is the elementwise ReLU activation. Then, $\rank\left(\nabla_{W^{l}} f_W(x)\right)\leq 1$.
\end{lemma}
\begin{proof}
We would like to show that the matrix $\rank(\nabla_{W^{l}} f_W(x))\leq 1$. We note that for any given vector $z \in \mathbb{R}^{d}$, we have $\sigma(z) = \textnormal{diag}(\sigma'(z)) \cdot z$. Therefore, for any input vector $x \in \R^{d_1}$, the output of $f_W$ can be written as follows,
\begin{small}
\begin{equation*}
\begin{aligned}
f_W(x) &= W^L \sigma(W^{L-1}\cdots \sigma( W^1 x)\cdots)\\
&= W^L \cdot D_{L-1}(x;W) W^{L-1}\cdots D_{1}(x;W) \cdot W^1 \cdot x,
\end{aligned}
\end{equation*}
\end{small}
where $D_l(x;W) = \textnormal{diag}[\sigma'(u_l(x;W)))]$ and $u_l(x;W) = W^l \sigma(W^{l-1}\cdots \sigma(W^1 x)\dots)$. We denote by $u_{l,i}(x;W)$ the $i$'th coordinate of the vector $u_l(x;W)$. We note that $u_l(x;W)$ are continuous functions of $W$. Therefore, assuming that none of the coordinates $u_{l,i}(x;W)$ are zero, there exists a sufficiently small ball around $W$ for which $u_{l,i}(x;W)$ does not change its sign. Hence, within this ball, $\sigma'(u_{l,i}(x;W))$ are constant. We define a set $\mathcal{W}_{l,i}=\{W \mid u_{l,i}(x;W)=0\}$. We note that as long as $x\neq 0$, the set $\mathcal{W}_{l,i}$ is negligible within $\R^{N}$. Since there is a finite set of indices $l,i$, the set $\bigcup_{l,i} \mathcal{W}_{l,i}$ is also negligible.
Let $W$ be a set of parameters for which all of the coordinates $u_{l,i}(x;W)$ are non-zero. Then, the matrices $\{D_{l}(x;W)\}^{L-1}_{l=1}$ are constant in the neighborhood of $W$, and therefore, their derivative with respect to $W^l$ are zero. Let $a^{\top} = W^L \cdot D_{L-1}(x;W) W^{L-1} \cdots W^{l+1} D_{l}(x;W)$ and $b = D_{l-1}(x) \cdot W^{l-1} \cdots W^1 x$. We can write $f_W(x) = a(x;W)^{\top} \cdot W^l \cdot b(x;W)$. Since the derivatives of $a(x;W)$ and $b(x;W)$ with respect to $W^l$ are zero, by applying the formula $\frac{\partial a^{\top}Xb}{\partial X}=ab^{\top}$, we have $\nabla_{W^{l}} f_W(x) = a(x;W) \cdot b(x;W)^{\top}$ which is a matrix of rank at most $1$. Therefore, for any input $x\neq 0$, with measure $1$ (over the selection of $W$), $\nabla_{W^{l}} f_W(x)$ is a matrix of rank at most $1$.
\end{proof}
The following theorem provides an upper bound on the minimal distance between the network's weight matrices and low-rank matrices.
\begin{restatable}{theorem}{batches}\label{thm:batches}
Let $\|\cdot \|$ be any matrix norm and $\ell$ any differentiable loss function. Let $f_W(x)$ be a ReLU neural network and $C^{ij}$ be a convolutional layer within $f_W$ and let $B \in [m]$. Then, for all $k < t$,
\begin{equation*}
\min_{W:~ \rank(W) \leq N^{ij}Bk} \left\|\fr{W^{ij}_t}{\|W^{ij}_t\|} - W\right\| ~\leq~ (1-2\mu\lambda)^k \cdot \frac{\|W^{ij}_{t-k}\|}{\|W^{ij}_t\|}.
\end{equation*}
\end{restatable}
\begin{proof}
We denote by $\tiS_t \subset S$ the training batch that was used by SGD at iteration $t$. We have
\begin{equation*}
\begin{aligned}
W^{ij}_t ~&=~ W^{ij}_{t-1} - \mu \nabla_{W^{ij}}L_{\tilde{S}_{t-1}}(f_{W_{t-1}}) - 2\mu \lambda W^{ij}_{t-1}\\
~&=~ (1-2\mu\lambda) W^{ij}_{t-1} - \mu\nabla_{W^{ij}}L_{\tilde{S}_{t-1}}(f_{W_{t-1}}).
\end{aligned}
\end{equation*}
Similarly, we can write
\begin{equation*}
\begin{aligned}
W^{ij}_{t-1} ~&=~ (1-2\mu\lambda) W^{ij}_{t-2} - \mu\nabla_{W^{ij}}L_{\tilde{S}_{t-2}}(f_{W_{t-2}}).
\end{aligned}
\end{equation*}
This gives us
\begin{equation*}
\begin{aligned}
W^{ij}_{t} ~&=~ (1-2\mu\lambda)^2 W^{ij}_{t-2} \\
&\quad- \mu\nabla_{W^{ij}}L_{\tilde{S}_{t-1}}(f_{W_{t-1}}) \\
&\quad- \mu(1-2\mu\lambda)\nabla_{W^{ij}}L_{\tilde{S}_{t-2}}(f_{W_{t-2}}).
\end{aligned}
\end{equation*}
By recursively applying this process $k$ times, we have
\begin{equation*}
\begin{aligned}
W^{ij}_t ~&=~ (1-2\mu\lambda)^k W^{ij}_{t-k} \\
&\quad\quad \underbrace{- \mu\sum^{k}_{l=1}(1-2\mu\lambda)^{l-1}\nabla_{W^{ij}}L_{\tilde{S}_{t-l}}(f_{W_{t-l}})}_{=:U} %\\
%~&=:~ (1-2\mu\lambda)^k W^{ij}_{t-k} + U.
\end{aligned}
\end{equation*}
By the chain rule, we can write the gradient of the loss function as follows,
\begin{equation*}
\begin{aligned}
&\nabla_{W^{ij}} L_{\tiS_{t-l}}(f_{W_{t-l}}) \\
~&=~ \fr{1}{B}\sum_{(x,y) \in \tiS_{t-l}}\fr{\partial \ell(f_{W_{t-l}}(x),y)}{\partial f_{W_{t-l}}(x)} \cdot \nabla_{W^{ij}} f_{W_{t-l}}(x).
\end{aligned}
\end{equation*}
According to Lem.~\ref{lem:rank}, we have $\rank(\nabla_{W^{ij}} f_{W_{t-l}}(x)) \leq N^{ij}$. Since $f_W$ is a univariate function, each term $\fr{\partial \ell(f_{W_{t-l}}(x),y)}{\partial f_{W_{t-l}}(x)}$ is a scalar. Therefore, $\rank(\nabla_{W^{ij}} L_{\tiS_{t-l}}(f_{W_{t-l}})) \leq BN^{ij}$ since $\nabla_{W^{ij}} L_{\tiS_{t-l}}(f_{W_{t-l}})$ is an average of $B$ matrices of rank at most $N^{ij}$. In particular, $\rank(U)\leq N^{ij}Bk$ since $U$ is a sum of $k$ matrices of rank at most $N^{ij}B$. Therefore, we obtain that
\begin{equation*}
\begin{aligned}
&\min_{W:~ \rank(W) \leq N^{ij}Bk}\left\|W^{ij}_t - W\right\| \\
~&\leq~ \left\|W^{ij}_t - U\right\|
~=~ (1-2\mu\lambda)^{k} \|W^{ij}_{t-k}\|.
\end{aligned}
\end{equation*}
Finally, by dividing both sides by $\|W^{ij}_t\|$ we obtain the desired inequality.
\end{proof}
The theorem above provides an upper bound on the minimal distance between the parameters matrix $W^{ij}_t$ and a matrix of rank $\leq N^{ij}Bk$. The parameter $k$ controls the looseness of the bound and is independent of the optimization process. The bound is proportional to $(1-2\mu\lambda)^k\fr{\|W^{ij}_t\|}{\|W^{ij}_{t-k}\|}$, which decreases exponentially with $k$. Assuming the norm of $W^{ij}_t$ converges as $t$ approaches infinity and $k = o(t)$, we see that $\lim\limits_{t \to \infty} \fr{\|W^{ij}_t\|}{\|W^{ij}_{t-k}\|} =1$. Thus, SGD with weight decay provably induces a low-rank bias in each weight matrix $W^{ij}$. By selecting $k=\lceil\fr{\log(\epsilon)}{\log(1-2\mu\lambda)}\rceil$, we can ensure that $(1-2\mu\lambda)^k\leq \epsilon$. In this case, at the end of the training, the normalized matrix $\fr{W^{ij}_t}{\|W^{ij}_t\|}$ can be approximated by a second matrix $W$ of rank $\leq \fr{N^{ij}B\log(\epsilon)}{\log(1-2\mu\lambda)}\leq \fr{N^{ij}B\log(1/\epsilon)}{2\mu\lambda}$ (the second inequality assumes that $\mu\lambda < 0.5$) with an error of $\epsilon$. While the value of $\frac{N^{ij}B\log(1/\epsilon)}{2\mu\lambda}$ may be large in small or medium-scale learning settings, it still yields meaningful results for very wide neural networks. For example, since the bound is independent of the input and output channels $c_j$ and $c_i$ of $C^{ij}$, when $c_i$ and $c_j$ are very large, the dimensions of $W^{ij}$ are much larger than $\frac{N^{ij}B\log(1/\epsilon)}{2\mu\lambda}$, and thus, our bound implies that the use of SGD would provably reduce the rank of $W^{ij}$ during training.
While Thm.~\ref{thm:batches} provides an upper bound of $\frac{N^{ij}B\log(1/\epsilon)}{2\mu\lambda}$ on the rank of the learned matrix, it does not give a precise insight into how various parameters influence the rank. However, based on the bound, we can still make the prediction that {\bf training with smaller batch sizes, increasing weight decay or learning rate will lead to lower rank matrices learned by SGD}. These predictions are empirically validated in the next section.
% \subsection{Degeneracy and the Origin of ``SGD Noise''}\label{sec:convergence}
% As we mentioned, it is impossible for SGD to converge to a stationary point of the gradient dynamical system, resulting in inherent ``SGD noise''. In this section, we study the (non-)convergence of mini-batch SGD. Our results are essentially impossibility results: the assumption of SGD convergence to a critical point of the gradient implies that the network represents the zero function. As a result, asymptotic noise is inherently unavoidable (when training properly). For simplicity, we assume that $\forall~i\in [m]:~x_i\neq 0$.
% As shown in \eqref{eq:convergence}, any convergence point $W$ of SGD satisfies $\E_{\tiS} [\| \nabla L^{\lambda}_{\tiS}(f_{W})\|]=0$. Since the distribution over mini-batches $\tiS$ of size $B$ is discrete and $\| \nabla L^{\lambda}_{\tiS}(f_{W})\|\geq 0$, we obtain that at convergence we have $\| \nabla_{W^{ij}} L^{\lambda}_{\tiS}(f_{W})\| = 0$, for all mini-batches $\tiS$ of size $B$. In particular,
% \begin{equation}\label{eq:1}
% \begin{aligned}
% \forall \tiS:~ 0 ~&=~ \nabla_{W^{ij}} L^{\lambda}_{\tiS}(f_W) ~=~ \frac{1}{B} \sum_{(x,y)\in \tiS} \fr{\partial \ell(f_W(x),y)}{\partial f_W(x)} \cdot \nabla_{W^{ij}} f_W(x) + 2\lambda W^{ij}.
% \end{aligned}
% \end{equation}
% Suppose we have two batches $\tiS_1,\tiS_2 \subset S$ of size $B$ that differ by only one sample. We denote the unique sample of each batch by $(x_{j_1},y_{j_1})$ and $(x_{j_2},y_{j_2})$ respectively. We notice that,
% \begin{equation*}
% \begin{aligned}
% 0 ~&=~ \nabla_{W^{ij}} L^{\lambda}_{\tiS_1}(f_W) - \nabla_{W^{ij}} L^{\lambda}_{\tiS_2}(f_W) \\
% ~&=~ \fr{\partial \ell(f_W(x_{j_1}),y_{j_1})}{\partial f_W(x_{j_1})} \cdot \nabla_{W^{ij}} f_W(x_{j_1}) - \fr{\partial \ell(f_W(x_{j_2}),y_{j_2})}{\partial f_W(x_{j_2})} \cdot \nabla_{W^{ij}} f_W(x_{j_2}).
% \end{aligned}
% \end{equation*}
% Therefore, we conclude that for all $j_1,j_2\in [m]$,
% \begin{equation}\label{eq:j1j2}
% \begin{aligned}
% M^{ij} ~=~ \fr{\partial \ell(f_W(x_{j_1}),y_{j_1})}{\partial f_W(x_{j_1})} \cdot \nabla_{W^{ij}} f_W(x_{j_1}) ~=~ \fr{\partial \ell(f_W(x_{j_2}),y_{j_2})}{\partial f_W(x_{j_2})} \cdot \nabla_{W^{ij}} f_W(x_{j_2}).
% \end{aligned}
% \end{equation}
% Hence, for all $(v_i,v_j) \in E_T$ and $k \in [m]$,
% \begin{equation}\label{eq:batchone}
% \fr{\partial \ell(f_W(x_k),y_k)}{\partial f_W(x_k)} \cdot \nabla_{W^{ij}} f(x_k) + 2\lambda W^{ij} ~=~ M^{ij} + 2\lambda W^{ij} ~=~ 0.
% \end{equation}
% Therefore, unless $\lambda=0$ or $\forall~(v_i,v_j)\in E_T:~W^{ij}=0$, we conclude that $\fr{\partial \ell(f_W(x_k),y_k)}{\partial f_W(x_k)}\neq 0$ for all $k\in [m]$. In this case, we also obtain that $\{\fr{\partial f_W(x_k)}{\partial \textnormal{vec}(W^{ij})}\}^{m}_{k=1}$ are collinear vectors by \eqref{eq:j1j2}.
% Therefore, any convergence point of training a neural network using mini-batch SGD along with weight decay is highly degenerate and does not fit any one of the training labels. To better understand the essence of this degeneracy, we provide the following proposition (proved in Appendix~\ref{app:proofs}), which is specialized for ReLU networks.
% \begin{restatable}[$\lambda > 0$]{proposition}{collinear}\label{prop:propA}
% Let $\ell(a,b)$ be a differentiable loss function, $\lambda > 0$, and let $f_W(x)$ be a ReLU neural network, where $\textnormal{succ}(1)=\{p\}$ and $(v_p,v_1)\in E_T$. Let $\{x^k_{i}\}^{N^{p1}}_{k=1}$ be the $N^{p1}$ patches of $x_i$ used by the layer $C^{p1}$. Let $W$ be a convergence point of mini-batch SGD for minimizing $L^{\lambda}_S(f_W)$ (see \eqref{eq:1}). Then, either $f_W \equiv 0$ or $\forall~ i,j\in [m]:\{x^k_i,x^k_j\}^{N^{p1}}_{k=1}$ are linearly dependent tensors.
% \end{restatable}
% The preceding proposition shows that unless the patches of any two training samples are linearly dependent, any convergence point of SGD corresponds to the zero function. When $N^{p1}$ is small, the linear dependence criterion is unrealistic. For example, if $C^{p1}$ is a fully-connected layer, $N^{p1}=1$, and the condition asserts that any two training samples $x_i,x_j$ are collinear. Since this is unrealistic, we conclude that convergence is impossible unless $f_W \equiv 0$.
% As a next step, we consider convergence of SGD when training without weight decay.
% \begin{restatable}[$\lambda = 0$]{proposition}{zerolambda}\label{prop:propB}
% Let $\lambda=0$ and let $\ell$ be a differentiable loss function. Let $f_W(x)$ be a ReLU neural network, where $\textnormal{succ}(1)=\{p\}$ and $(v_p,v_1)\in E_T$ is fully-connected. Let $\{x^k_{i}\}^{N^{p1}}_{k=1}$ be the $N^{p1}$ patches of $x_i$ used by the layer $C^{p1}$. Let $W$ be a convergence point of mini-batch SGD for minimizing $L^{\lambda}_S(f_W)$. Then, $\forall~ i\in [m]:\fr{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})}=0$ or $f_W(x_j)=0$ or $\forall~ j\in [m]:\{x^k_i,x^k_j\}^{N^{p1}}_{k=1}$ are linearly dependent tensors. %In particular, if $\ell$ is convex and $\forall~b\in \R~\nexists~a^* \in \R:~ \ell(a^*,b)=\inf_a \ell(a,b)$, then, for all $i \in [m]$, $f_W(x_i)=0$ or $\{x^k_i,x^k_j\}^{N^{p1}}_{k=1}$ are linearly dependent tensors.
% \end{restatable}
% The preceding proposition (proved in Appendix~\ref{app:proofs}) provides conditions for SGD convergence of training a convolutional network without weight decay. It shows that at convergence, for every sample the network either perfectly fits the label (i.e., $\fr{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} = 0$) or outputs zero unless the patches of that sample are linearly dependent with the patches of any other sample. As mentioned, the linear dependence criteria is generally unrealistic when $(v_p,v_1)$ is fully-connected. Therefore, convergence with fully-connected networks is possible only when $f_W$ perfectly fits the training set labels.
% We note that if $\ell(a,b)$ is convex and has no minima $a$ for any $b \in \R$ (e.g., binary cross-entropy, logistic loss or exponential loss), then $\forall~ i \in [m]:\fr{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \neq 0$. Therefore, the only possible convergence points of fully-connected networks are ones for which $\forall~i \in [m]:f_W(x_i)=0$. Since this is in general absurd, we argue that perfect convergence of training a network with exponential-type loss functions is generally impossible. While convergence to a non-zero function is not guaranteed, in practice training without weight decay may still fall into the regime of `almost convergence', in which $\max_{i \in [m]} \big\vert \fr{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \big\vert$ is tiny and as a result the training steps $- \mu \cdot \fr{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \cdot \fr{\partial f_W(x_i)}{\partial W^{ij}}$ are very small. This is usually the case when an overparameterized network has been properly trained. For the squared loss, convergence may occur when the network perfectly fits the training labels, i.e., $\forall~i\in [m]:~f_W(x_i)=y_i$.
% It is worth noting that training with mini-batches during optimization and weight decay are critical to our analysis. While many papers examine the training dynamics and critical points of GD, as we show, SGD convergence points are highly degenerate and, in general, behave differently than GD solutions. %For example, it is impossible to derive \eqref{eq:j1j2} without using mini-batches. Without weight decay, we cannot claim that $W^{ij}$ is proportional to the low-rank matrix $M^{ij}$.
% Surprisingly, this analysis is unaffected by batch size, and therefore, the presence of SGD noise occurs regardless of batch size, as long as it is strictly smaller than the full dataset's size.
\section{Experiments}\label{sec:experiments}
\begin{figure*}[!ht]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.33\linewidth]{plots/resnet18/learning_rate/rank_resnet_bn_batch_size_4_learning_rate.pdf} &
\includegraphics[width=0.33\linewidth]{plots/resnet18/learning_rate/rank_resnet_bn_batch_size_8_learning_rate.pdf} &
\includegraphics[width=0.33\linewidth]{plots/resnet18/learning_rate/rank_resnet_bn_batch_size_16_learning_rate.pdf} %&
%\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/rank_resnet_bn_batch_size_32_learning_rate.pdf}
\\
\includegraphics[width=0.33\linewidth]{plots/resnet18/learning_rate/acc_resnet_bn_batch_size_4_learning_rate.pdf} &
\includegraphics[width=0.33\linewidth]{plots/resnet18/learning_rate/acc_resnet_bn_batch_size_8_learning_rate.pdf} &
\includegraphics[width=0.33\linewidth]{plots/resnet18/learning_rate/acc_resnet_bn_batch_size_16_learning_rate.pdf} %&
%\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/acc_resnet_bn_batch_size_32_learning_rate.pdf}
\\
{\small $B=4$} & {\small $B=8$} & {\small $B=16$} %& {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of ResNet-18 trained on CIFAR10 with varying $\mu$.} The top row shows the average rank across layers, while the bottom row shows the train and test accuracy rates for each setting. In this experiment, $\mu=\expnumber{5}{-4}$ and $\epsilon=\expnumber{1}{-3}$.}
\label{fig:lr_resnet18}
\begin{tabular}{c@{}c@{}c@{}c@{}c}
%\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/rank_resnet_bn_bs_4_lr_15_weight_decay.pdf} &
\includegraphics[width=0.33\linewidth]{plots/resnet18/weight_decay/rank_resnet_bn_bs_8_lr_15_weight_decay.pdf} &
\includegraphics[width=0.33\linewidth]{plots/resnet18/weight_decay/rank_resnet_bn_bs_16_lr_15_weight_decay.pdf} &
\includegraphics[width=0.33\linewidth]{plots/resnet18/weight_decay/rank_resnet_bn_bs_32_lr_15_weight_decay.pdf}
\\
%\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/acc_resnet_bn_bs_4_lr_15_weight_decay.pdf} &
\includegraphics[width=0.33\linewidth]{plots/resnet18/weight_decay/acc_resnet_bn_bs_8_lr_15_weight_decay.pdf} &
\includegraphics[width=0.33\linewidth]{plots/resnet18/weight_decay/acc_resnet_bn_bs_16_lr_15_weight_decay.pdf} &
\includegraphics[width=0.33\linewidth]{plots/resnet18/weight_decay/acc_resnet_bn_bs_32_lr_15_weight_decay.pdf}
\\
%{\small $B=4$} &
{\small $B=8$} & {\small $B=16$} & {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of ResNet-18 trained on CIFAR10 with different $\lambda$ values.} In this experiment, $\mu=1.5$ and $\epsilon=\expnumber{1}{-3}$.}
\label{fig:wd_resnet18}
\end{figure*}
% \begin{figure*}[t]
% \centering
% \end{figure*}
\begin{figure*}[!ht] % vit batch size
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.33\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_03_batch_size.pdf} &
%\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_05_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_3_batch_size.pdf}
\\
\includegraphics[width=0.33\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_03_batch_size.pdf} &
%\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_05_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_3_batch_size.pdf}
\\
{\small $\mu=0.03$} %& {\small $\mu=0.05$}
& {\small $\mu=0.1$} & {\small $\mu=0.3$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of MLP-BN-10-100 trained on CIFAR10 with various batch sizes.} In this experiment, $\lambda=\expnumber{5}{-4}$ and $\epsilon=\expnumber{1}{-3}$.}\label{fig:mlp_bn_10_100_lr/bs}
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.33\linewidth]{plots/vit/batch_size/varyB_lr4e-3_avgrank.pdf} &
\includegraphics[width=0.33\linewidth]{plots/vit/batch_size/varyB_lr8e-3_avgrank.pdf} &
%\includegraphics[width=0.20\linewidth]{plots/vit/batch_size/varyB_lr1e-2_avgrank.pdf} &
%\includegraphics[width=0.25\linewidth]{plots/vit/batch_size/varyB_lr2e-2_avgrank.pdf} &
\includegraphics[width=0.33\linewidth]{plots/vit/batch_size/varyB_lr4e-2_avgrank.pdf}
\\
\includegraphics[width=0.33\linewidth]{plots/vit/batch_size/varyB_lr4e-3_accuracy.pdf} &
\includegraphics[width=0.33\linewidth]{plots/vit/batch_size/varyB_lr8e-3_accuracy.pdf} &
%\includegraphics[width=0.20\linewidth]{plots/vit/batch_size/varyB_lr1e-2_accuracy.pdf} &
%\includegraphics[width=0.25\linewidth]{plots/vit/batch_size/varyB_lr2e-2_accuracy.pdf} &
\includegraphics[width=0.33\linewidth]{plots/vit/batch_size/varyB_lr4e-2_accuracy.pdf}
\\
{\small $\mu=0.004$} & {\small $\mu=0.008$} %& {\small $\mu=0.01$}
%& {\small $\mu=0.02$}
& {\small $\mu=0.04$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of ViT trained on CIFAR10 with various batch sizes.} In this experiment, $\lambda=\expnumber{5}{-4}$ and $\epsilon=\expnumber{1}{-2}$.}\label{fig:bs_vit}
%(no-momentum, patch size=4, depth=6, heads=8, dropout=0.1, $mlp_dim$=512, weight decay=5e-4)
\end{figure*}
\begin{figure*}[!ht] % resnet generalization
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.33\linewidth]{plots/generalization/resnet/rank_resnet_bn_lr_001_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/generalization/resnet/rank_resnet_bn_lr_005_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/generalization/resnet/rank_resnet_bn_lr_01_batch_size.pdf}
\\
\includegraphics[width=0.33\linewidth]{plots/generalization/resnet/acc_resnet_bn_lr_001_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/generalization/resnet/acc_resnet_bn_lr_005_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/generalization/resnet/acc_resnet_bn_lr_01_batch_size.pdf}
\\
{\bf {(a)}} & {\bf {(b)}} & {\bf {(c)}}
\end{tabular}
\label{fig:gen_1}
\caption{{\bf Average ranks and accuracy rates of ResNet-18 trained on CIFAR10 with different batch sizes.} {\bf(a)} was trained with $\mu=\expnumber{1}{-3}, \lambda=\expnumber{6}{-3}$, {\bf(b)} was trained with $\mu=\expnumber{5}{-3}, \lambda=\expnumber{6}{-3}$, and {\bf (c)} was trained with $\mu=\expnumber{1}{-2}, \lambda=\expnumber{4}{-4}$. We used a threshold of $\epsilon=\expnumber{1}{-2}$.}
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.33\linewidth]{plots/generalization/vgg/rank_vgg_bn_lr_001_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/generalization/vgg/rank_vgg_bn_lr_005_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/generalization/vgg/rank_vgg_bn_lr_01_batch_size.pdf}
\\
\includegraphics[width=0.33\linewidth]{plots/generalization/vgg/acc_vgg_bn_lr_001_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/generalization/vgg/acc_vgg_bn_lr_005_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/generalization/vgg/acc_vgg_bn_lr_01_batch_size.pdf}
\\
{\bf {(a)}} & {\bf {(b)}} & {\bf {(c)}}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of VGG-16 trained on CIFAR10 with different batch sizes.} {\bf(a)} was trained with $\mu=\expnumber{1}{-3}, \lambda=\expnumber{6}{-3}$, {\bf(b)} was trained with $\mu=\expnumber{5}{-3}, \lambda=\expnumber{5}{-4}$, and {\bf(c)} was trained with $\mu=\expnumber{1}{-2}, \lambda=\expnumber{4}{-4}$. We used a threshold of $\epsilon=\expnumber{4}{-2}$.} \label{fig:gen_2}
\end{figure*}
In this section, we empirically study how batch size, weight decay, and learning rate affect the rank of matrices in deep ReLU networks. We conduct separate experiments where we vary one hyperparameter while keeping the others constant to isolate its effect on the averaged rank\footnote{The plots are best viewed when zooming into the pictures.}. Additional experiments are provided in Appendix~\ref{app:additionalExp}.%{\color{blue} Each of the runs were trained using a single GPU on a computing cluster with several available GPU types. The length of the runs did not exceed 60 hours, and the available GPUs were NVIDIA's GeForce RTX 2080, GeForce RTX 2080 Ti, Quadro RTX 6000, Tesla V-100, GeForce RTX A6000, A100, and GeForce GTX 1080 Ti.}
\subsection{Setup}
{\bf Architectures.\enspace} We evaluate several network architectures in our study. (i) The first architecture is an MLP, denoted as MLP-BN-$L$-$H$, which comprises $L$ hidden layers, each containing a fully-connected layer with width $H$, followed by batch normalization and ReLU activations. This architecture ends with a fully-connected output layer. The same architecture without batch normalization is denoted by MLP-$L$-$H$. (ii) The second architecture, referred to as RES-BN-$L$-$H$, consists of a linear layer with width $H$, followed by $L$ residual blocks, and ending with a fully-connected layer. Each block performs a computation of the form $z+\sigma(n_2(W_2 \sigma(n_1(W_1z))))$, where $W_1,W_2 \in \R^{H\times H}$, $n_1,n_2$ are batch normalization layers, and $\sigma$ is the ReLU function. (iii) The third architecture is the convolutional network (VGG-16) proposed by~\cite{DBLP:journals/corr/SimonyanZ14a}, with dropout replaced by batch normalization layers, and a single fully-connected layer at the end. (iv) The fourth architecture is the residual network (ResNet-18) proposed in~\citep{7780459}. (v) The fifth architecture is a small visual transformer (ViT) \citep{dosovitskiy2020image}. Our implementation of ViT splits the input images into patches of size $4\times 4$, includes $8$ self-attention heads, each composed of $6$ self-attention layers. The self-attention layers are followed by two fully-connected layers with a dropout probability of 0.1, and a GELU activation in between them.
{\bf Training and evaluation.\enspace} We trained each model for CIFAR10 classification using Cross-Entropy loss minimization between its logits and the one-hot encodings of the labels. The training was carried out by SGD with batch size $B$, initial learning rate $\mu$, and weight decay $\lambda$. The MLP-BN-$L$-$H$, RES-BN-$L$-$H$, ResNet-18, and VGG-16 models were trained with a decreasing learning rate of 0.1 at epochs 60, 100, and 200, and the training was stopped after 500 epochs. The ViT models were trained using SGD with a learning rate that was decreased by a factor of 0.2 at epochs 60 and 100 and training was stopped after 200 epochs. During training, we applied random cropping, random horizontal flips, and random rotations (by $15k$ degrees for $k$ uniformly sampled from $[24]$) and standardized the data.
To study the influence of different hyperparameters on the rank of the weight matrices, in each experiment, we trained the models while varying one hyperparameter at a time, while keeping other hyperparameters constant. After each epoch, we compute the average rank across the network's weight matrices and its train and test accuracy rates. For a convolutional layer $C^{ij}$, we use $W^{ij}$ as its weight matrix. To estimate the rank of a given matrix $M$, we count how many of the singular values of $\fr{M}{\|M\|_2}$ are out of the range $[-\epsilon, \epsilon]$, where $\epsilon$ is a small tolerance value.
%Unless mentioned otherwise,
% {\bf Datasets.\enspace} Throughout the experiments, we consider three datasets: (i) MNIST; (ii) CIFAR10; (iii) STL-10. For CIFAR10 we used random cropping, random horizontal flips and random rotations (by $15k$ degrees for $k$ uniformly sampled from $[24]$). For all datasets we standardized the data. {\color{red}Update}
\subsection{Results}%\label{sec:rank_exp}
{\bf Low-rank bias and the batch size.\enspace} As shown in Figs.\ref{fig:mlp_bn_10_100_lr/bs}-\ref{fig:gen_2}, decreasing the batch size strengthens the low-rank constraint on the network's matrices, resulting in matrices of lower ranks. This aligns with the prediction made in Sec.~\ref{sec:theory} that training with smaller batch sizes leads to matrices of lower ranks. This observation highlights the impact of batch size on the rank of the weight matrices and how it can be used to control the complexity of the network.
{\bf Low-rank bias, weight decay and learning rate.\enspace} As shown in Fig.~\ref{fig:wd_resnet18}, increasing $\lambda$ imposes stronger rank minimization constraints on the weight matrices. Interestingly, the effect of batch size on the ranks of the weight matrices appears to be minimal when $\lambda=0$, which suggests that weight decay is essential for imposing a noticeable low-rank bias on the weight matrices. Furthermore, Figs.~\ref{fig:lr_resnet18},~\ref{fig:mlp_bn_10_100_lr/bs} and~\ref{fig:bs_vit} show that increasing the learning rate tends to lead to smaller ranks of weight matrices, which aligns with the prediction made in Sec.~\ref{sec:theory}.
{\bf Low-Rank Bias and Generalization.\enspace} We investigated the relationship between low-rank bias and generalization by training ResNet-18 and VGG-16 models on CIFAR10 with varying batch sizes, while keeping $\lambda$ and $\mu$ constant. To provide a fair comparison, we selected $\lambda$ and $\mu$ to ensure all models fit the training data equally. Our results, shown in Figs.~\ref{fig:gen_1}-\ref{fig:gen_2}, indicate that models trained with smaller batch sizes (i.e. lower rank in their weights) tend to generalize better as the test accuracy rate monotonically increases as the batch size decreases. Based on these findings, we hypothesize that when altering a certain hyperparameter, a neural network with a lower average rank will have better performance than a network with the same architecture but higher rank matrices, assuming both networks perfectly fit the training data.
%To study the effect of the weight decay on rank minimization, we train instances of MLP-BN-10-100 using SGD with various batch sizes $B \in \{16,32,64,128\}$ and values of $\lambda \in \{0, \expnumber{1}{-4}, \expnumber{2}{-4}, \expnumber{4}{-4}, \expnumber{8}{-4}, \expnumber{6}{-3}\}$. The learning rate is initially $\mu=0.1$ and is decayed by a factor of $0.1$ three times at epochs 60, 100, and 200 and training is stopped after 500 epoch and the weight decay is $\lambda=\expnumber{5}{-4}$. {\color{blue} {\bf (ii)} We additionally train instances of VGG-16 and ResNet-18 with the previous hyperparamaters except with momentum $0$.}
%The results are reported in Fig. .
% \subsection{Experiments on SGD Noise}\label{sec:sgd_noise_exp}
% In Sec.~\ref{sec:convergence} we showed that convergence to a non-zero function is impossible when training a fully-connected neural network with SGD and weight decay. To validate this prediction, we trained MLP-5-2000 instances examined their convergence as a function of $\lambda$. Each model was trained for CIFAR10 classification using SGD with batch size 128 and learning rate 0.1 for 2000 epochs.
% To investigate the convergence of the networks, we measure the averaged distance between the network's matrices between consecutive epochs: $d(W_{t+1},W_{t}) := \fr{1}{\vert E_T\vert} \sum_{(p,q)\in E_T} \|W^{ij}_{t+1} - W^{ij}_{t}\|$, where $\{W^{ij}_t\}_{(p,q)\in E_T}$ are the various trainable matrices in the network at epoch $t$. As mentioned in Sec.~\ref{sec:theory}, convergence is possible only when $\lim\limits_{t\to \infty} d(W_{t+1},W_{t}) = 0$.
% In Fig.~\ref{fig:convergence} we monitor {\bf (a)} $d(W_{t+1},W_{t})$, {\bf (b)} the train accuracy rates, {\bf (c)} the train losses and {\bf (d)} the averaged rank of the trainable matrices. As predicted in Prop.~\ref{prop:propA}, when training with $\lambda>0$, $W_t$ either converges to zero and $f_{W_t}$ to the zero function (e.g., see the results with $\lambda=\expnumber{1}{-4}, \expnumber{1}{-3}$) or $W_t$ does not converge (i.e., $d(W_{t+1},W_{t})$ does not tend to zero). Furthermore, we observe that for $\lambda=0$, $d(W_{t+1},W_{t})$ is smaller by orders of magnitude compared to using $\lambda>0$ (except for cases when the selection of $\lambda > 0$ leads to $W^{ij}_t \to 0$). Interestingly, even though for certain values of $\lambda > 0$, the training loss and accuracy converged, the network's parameters do not converge.
% Finally, when training for cross-entropy loss minimization without weight decay we encounter the `almost convergence' regime discussed in Sec.~\ref{sec:convergence}. Namely, even though perfect convergence is impossible, the term $d(W_{t+1},W_{t})$ may become as small as we wish by increasing the size of the neural network. Therefore, since the MLP-5-2000 is relatively large (compared to the dataset's size), we may inaccurately get the impression that $d(W_{t+1},W_{t})$ tends to zero.
% We repeated the same experiment with instances of MLP-10-100 without the batch normalization as well. In this case, we used $\lambda \in \{0, \expnumber{1}{-3}, \expnumber{2}{-3}, \expnumber{4}{-3}, \expnumber{8}{-3}, \expnumber{6}{-2}, \expnumber{2}{-1}\}$. For each run we trained the network for 500 epochs using SGD with momentum $0.9$, initial learning rate $0.01$ and batch size 64.
\section{Conclusions}\label{sec:conclusions}
A mathematical characterization of the biases associated with SGD-trained neural networks is regarded as a significant open problem in the theory of deep learning~\citep{NIPS2017_10ce03a1}. In addition to its independent interest, a low-rank bias -- though probably not necessary for generalization -- may be a key ingredient in an eventual characterization of the generalization properties of deep networks. In fact, recent results~\citep{huh2022the} and our preliminary experiments (see Figs.~\ref{fig:gen_1}-\ref{fig:gen_2} in the appendix) suggest that low-rank bias in neural networks improves generalization.
Our study of deep ReLU neural networks trained with mini-batch Stochastic Gradient Descent (SGD) and weight decay shows that the resulting weight matrices tend to be low-rank when training with small batch sizes, high learning rates, or high levels of weight decay. Our theoretical and empirical results provide a better understanding of how these hyperparameters can be used to control the complexity of the network and potentially improve generalization.
While this paper focused on a basic supervised learning setting using SGD and weight decay, but future studies could investigate the structure of weights and activations in neural networks trained with other optimization methods and regularization techniques. Additionally, it would be valuable to study these biases in unsupervised and self-supervised learning settings for deeper insights into network training. Another interesting direction would be to examine the effects of different architectures such as recurrent neural networks or transformer networks on the rank minimization bias.
\newpage
\section*{Acknowledgements}
We thank Mikhail Belkin, Eran Malach and Vardan Papyan for illuminating discussions during the preparation of this manuscript. This material is based upon work supported by the Center for Minds, Brains and Machines (CBMM), funded by NSF STC award CCF-1231216.
\bibliography{refs}
\bibliographystyle{icml2023}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% APPENDIX
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\newpage
\appendix
\onecolumn
\section{Proofs}\label{app:proofs}
\ranklemma*
\begin{proof} Let $x \in \R^{c_1\times h_1\times w_1}$ be an input tensor, and $C^{ij}$ be a certain convolutional layer with kernel size $(k_1,k_2)$, stride $s$, and padding $p$. We wish to show that $\rank\left(\nabla_{W^{ij}} f_W(x)\right) \leq N^{ij}$. We begin by expressing the output of $f_W$ as a sum of paths that pass through $C^{ij}$ and paths that do not. We can express the output as follows:
\begin{equation*}
f_W(x) ~=~ \sum_{l_1 \in \textnormal{pred}(l_0)} C^{l_0l_1} \circ v_{l_1}(x),
\end{equation*}
where $l_0 = L$ and $C^{l_0l_1} \circ z := C^{l_0l_1}(z)$. Each layer $v_{l}$ can also be expressed as:
\begin{equation*}
v_{l_1}(x) ~=~ D_{l_1} \odot \sum_{l_2\in \textnormal{pred}(l_1)} C^{l_1l_2} \circ v_{l_2}(x),
\end{equation*}
where $D_l := D_l(x) := \sigma'(v_l(x))) \in \R^{c_l \times h_l \times w_l}$.
A path $\pi$ within the network's graph $G$ is a sequence $\pi = (\pi_0,\dots,\pi_T)$, where $\pi_0=1$, $\pi_T = L$ and for all $i=0,\dots,T-1:~ (v_{\pi_{i}},v_{\pi_{i+1}}) \in E$. We can write $f_W(x)$ as the sum of matrix multiplications along paths $\pi$ from $v_1$ to $v_{l_0}$. Specifically, we can write $f_W(x)$ as a follows
%\begin{small}
\begin{equation*}
\begin{aligned}
f_W(x) ~&=~ \sum_{\pi \textnormal{ from $i$ to $l_0$}} C^{\pi_{T}\pi_{T-1}} \circ D_{\pi_{T-1}} \cdots D_{\pi_2} \odot C^{\pi_2\pi_1} \circ D_{\pi_1} \odot C^{ij} \circ v_j(x) \\
&\quad\quad + \sum_{\substack{\pi \textnormal{ from $1$ to $l_0$} \\ (i,j) \notin \pi }} C^{\pi_{T}\pi_{T-1}} \circ D_{\pi_{T-1}} \odot C^{\pi_{T-1}\pi_{T-2}} \cdots D_{\pi_2} \odot C^{\pi_2 \pi_1} \circ x, \\
~&=:~ A_W(x)+B_W(x)
\end{aligned}
\end{equation*}
%\end{small}
where $T=T(\pi)$ denotes the length of the path $\pi$. Since $\sigma$ is a piece-wise linear function with a finite number of pieces, for any $x\in \R^{c_1\times h_1\times w_1}$, with measure $1$ over $W$, the matrices $\{D_{l}(x)\}^{L-1}_{l=1}$ are constant in the neighborhood of $W$. Furthermore, $W^{ij}$ does not appear in the multiplications along the paths $\pi$ from $1$ to $l_0$ that exclude $(i,j)$. Therefore, we conclude that $\frac{\partial B_W(x)}{\partial W^{ij}}=0$.
As a next step we would like to analyze the rank of $\frac{\partial A_W(x)}{\partial W^{ij}}$. For this purpose, we rewrite the convolutional layers and the multiplications by the matrices $D_{l}(x)$ as matrix multiplications.
{\bf Representing $C^{ij}$.\enspace} We begin by representing the layer $C^{ij}$ as a linear transformation of its input with $N^{ij}$ blocks of $W^{ij}$. For this purpose, we define a representation of a given $3$-dimensional tensor input $z \in \R^{c_j \times h_j \times w_j}$ as a vector $\vect^{ij}(z) \in \R^{N^{ij} c_{j} k_1 k_2}$. First, we pad $z$ with $p$ rows and columns of zeros and obtain $\textnormal{Pad}_p(z)$. We then vectorize each one of its patches (of dimensions $c_{j}\times k_1\times k_2$) that the convolutional layer is acting upon (potentially overlapping) and concatenate them. We can write the vectorized output of the convolutional layer as $U^{ij} \vect^{ij}(z)$, where
\begin{equation}\label{eq:Uij}
U^{ij} := \renewcommand{\arraystretch}{1.2}
\left(
\begin{array}{ c c c c c c }
\multicolumn{1}{c}{} & & 0 & 0 & 0 & 0 \\
\multicolumn{2}{c}{\raisebox{.6\normalbaselineskip}[0pt][0pt]{$W^{ij}$}} & 0 & 0 & 0 & 0 \\
0 & 0 & & & 0 & 0 \\
0 & 0 & \multicolumn{2}{c}{\raisebox{.6\normalbaselineskip}[0pt][0pt]{$\ddots$}} & 0 & 0 \\
0 & 0 & 0 & 0 & & \multicolumn{1}{c}{} \\
0 & 0 & 0 & 0 & \multicolumn{2}{c}{\raisebox{.6\normalbaselineskip}[0pt][0pt]{$W^{ij}$}}
\end{array}
\right)
\end{equation}
is a $(N^{ij} c_{i}) \times (N^{ij} c_{j} k_1k_2)$ matrix with $N^{ij}$ copies of $W^{ij}$. We note that this is a non-standard representation of the convolutional layer's operation as a linear transformation. Typically, we write the convolutional layer as a linear transformation $V^{ij}$ acting on the vectorized version $\vect(z) \in \R^{c_{j}k_1k_2}$ of its input $z$. Since $\vect^{ij}(z)$ consists of the same variables as in $\vect(z)$ with potentially duplicate items, there is a linear transformation that translates $\vect(z)$ into $\vect^{ij}(z)$. Therefore, we can simply write $V^{ij} \vect(z) = U^{ij} \vect^{ij}(z)$.
{\bf Representing convolutional layers.\enspace} Except of $C^{ij}$, we represent each one of the network's convolutional layers $C^{ld}$ in $f_W$ as linear transformations. As mentioned earlier, we can write $\vect(C^{ld}(z)) = V^{ld}\vect(z)$, for any input $z\in \R^{c_d \times h_d \times w_d}$.
{\bf Representing pooling and rearrangement layers.\enspace} An average pooling layer or a rearrangement layer $C^{ld}$ can be easily represented as a (non-trainable) linear transformation of its input. Namely, we can write $\vect(C^{ld}(z)) = V^{ld}\vect(z)$ for some matrix $V^{ld}$. A max-pooling layer can be written as a composition of ReLU activations and multiple (non-trainable) linear transformations, since $\max(x,y) = \sigma(x-y) + y$. Therefore, without loss of generality we can replace the pooling layers with non-trainable linear transformations and ReLU activations.
{\bf Computing the rank.\enspace} Finally, we note that $\vect(C^{ij} \circ z) = U^{ij} \vect^{ij}(z) = V^{ij} \vect(z)$, $\vect(D_l \odot z) = P_l \cdot \vect(z)$ for $P_{l} := \textnormal{diag}(\vect(D_{l}))$. Therefore, we can write
\begin{equation*}
\begin{aligned}
A_W(x) ~&=~ \sum_{\pi \textnormal{ from $i$ to $l_0$}} W^{\pi_{T}\pi_{T-1}} \cdot P_{\pi_{T-1}} \cdots P_{\pi_2} \cdot W^{\pi_2\pi_1} \cdot P_{\pi_1} \cdot U^{ij} \cdot \vect^{ij}(v_j(x)) \\
~&=:~ a(x)^{\top} \cdot U^{ij} \cdot b(x),
\end{aligned}
\end{equation*}
where $a(x)^{\top} ~:=~ \sum_{\pi \textnormal{ from $i$ to $l_0$}} W^{\pi_{T}\pi_{T-1}} \cdot P_{\pi_{T-1}} \cdots P_{\pi_2} \cdot W^{\pi_2\pi_1} \cdot P_{\pi_1}$ and $b(x) := \vect^{ij}(v_j(x))$.
We note that with measure $1$, the matrices $\{P_l\}^{L-1}_{l=1}$ are constant in the neighborhood of $W$. In addition, $a(x)$ and $b(x)$ are computed as multiplications of matrices $W^{ld}$ and $P_l$ excluding $(i,j)=(p,q)$. Therefore, with measure $1$ over the selection of $W$, the Jacobians of $a(x)$ and $b(x)$ with respect to $V^{ij}$ are $0$. Furthermore, due to \eqref{eq:Uij} and the definition of $U^{ij}$, we can write
\begin{equation*}
\begin{aligned}
a(x)^{\top} \cdot U^{ij} \cdot b(x) ~=~ \sum^{N^{ij}}_{t=1} a_t(x)^{\top} \cdot W^{ij} \cdot b_t(x),
\end{aligned}
\end{equation*}
where $a_t(x)$ and $b_t(x)$ are the slices of $a(x)$ and $b(x)$ that are multiplied by the $t$'th $W^{ij}$ block in $U^{ij}$. Since the Jacobians of $a_i(x)$ and $b_i(x)$ with respect to $W^{ij}$ are $0$ with measure $1$ over the selection of $W$ (from the same argument as in the proof of Lem.~\ref{lem:simple}), we have,
\begin{equation}
\frac{\partial a(x)^{\top} \cdot U^{ij} \cdot b(x)}{\partial W^{ij}} ~=~ \sum^{N^{ij}}_{t=1} a_t(x) \cdot b_t(x)^{\top}.
\end{equation}
Therefore, we conclude that, with measure $1$ over the selection of $W$, we have $\frac{\partial f_W(x)}{\partial W^{ij}} = \sum^{N^{ij}}_{t=1} a_t(x) \cdot b_t(x)^{\top}$ which is a matrix of rank $\leq N^{ij}$.
\end{proof}
\section{Additional Experiments}\label{app:additionalExp}
To further demonstrate the bias towards rank minimization of SGD with weight decay, we conducted a series of experiments with different learning settings. We follow the same training and evaluation protocol described in Sec.~\ref{sec:setup}. The results are summarized in Figs.~\ref{fig:bs}-\ref{fig:vit_var_wd}.
\begin{figure*}[ht]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/rank_batch_size_mlp_100_10_lr_0001.pdf}&
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/rank_batch_size_mlp_100_10_lr_001.pdf}&
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/rank_batch_size_mlp_100_10_lr_01.pdf}&
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/rank_batch_size_mlp_100_10_lr_05.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/acc_batch_size_mlp_100_10_lr_0001.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/acc_batch_size_mlp_100_10_lr_001.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/acc_batch_size_mlp_100_10_lr_01.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/acc_batch_size_mlp_100_10_lr_05.pdf}
\\
{\small $\mu=0.001$} & {\small $\mu=0.01$} & {\small $\mu=0.1$} & {\small $\mu=0.5$}
\end{tabular}
\caption{{\bf Average rank of MLP-BN-10-100 trained on CIFAR10 with various batch sizes.} In this experiment, $\lambda=\expnumber{5}{-4}$ and $\epsilon=\expnumber{1}{-3}$.}
\label{fig:bs}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/rank_wd_mlp_100_10_bs_16.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/rank_wd_mlp_100_10_bs_32.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/rank_wd_mlp_100_10_bs_64.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/rank_wd_mlp_100_10_bs_128.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/acc_wd_mlp_100_10_bs_16.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/acc_wd_mlp_100_10_bs_32.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/acc_wd_mlp_100_10_bs_64.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/acc_wd_mlp_100_10_bs_128.pdf}
\\
{\small $B=16$} & {\small $B=32$} & {\small $B=64$} & {\small $B=128$}
\end{tabular}
\caption{{\bf Average rank of MLP-BN-10-100 trained on CIFAR10 with varying $\lambda$.} In this experiment, $\mu=0.1$, momentum $0.9$ and $\epsilon=\expnumber{1}{-3}$.}
\label{fig:momentum}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/rank_batch_size_resnet_100_5_lr_00125.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/rank_batch_size_resnet_100_5_lr_0025.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/rank_batch_size_resnet_100_5_lr_005.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/rank_batch_size_resnet_100_5_lr_01.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/acc_batch_size_resnet_100_5_lr_00125.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/acc_batch_size_resnet_100_5_lr_0025.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/acc_batch_size_resnet_100_5_lr_005.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/acc_batch_size_resnet_100_5_lr_01.pdf}
\\
{\small $\mu=0.0125$} & {\small $\mu=0.025$} & {\small $\mu=0.05$} & {\small $\mu=0.1$}
\end{tabular}
\caption{{\bf Average rank of RES-BN-5-500 trained on CIFAR10 with various batch sizes.} In this experiment, $\lambda=\expnumber{5}{-4}$ and $\epsilon=\expnumber{1}{-3}$.}
\label{fig:bs_resnet}
\end{figure*}
% \begin{figure*}[t]
% \centering
% \begin{tabular}{c@{}c@{}c@{}c@{}c}
% \includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/rank_batch_size_mlp_nobn_500_5_lr_00125.pdf} &
% \includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/rank_batch_size_mlp_nobn_500_5_lr_0025.pdf} &
% \includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/rank_batch_size_mlp_nobn_500_5_lr_005.pdf} &
% \includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/rank_batch_size_mlp_nobn_500_5_lr_01.pdf}
% \\
% \includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/acc_batch_size_mlp_nobn_500_5_lr_00125.pdf} &
% \includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/acc_batch_size_mlp_nobn_500_5_lr_0025.pdf} &
% \includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/acc_batch_size_mlp_nobn_500_5_lr_005.pdf} &
% \includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/acc_batch_size_mlp_nobn_500_5_lr_01.pdf}
% \\
% {\small $\mu=0.0125$} & {\small $\mu=0.025$} & {\small $\mu=0.05$} & {\small $\mu=0.1$}
% \end{tabular}
% \caption{{\bf Average ranks and accuracy rates of MLP-5-500 trained on CIFAR10 with various batch sizes.} Each model was trained with a $\lambda=\expnumber{5}{-4}$ weight decay. To estimate the rank, we used an $\epsilon=0.001$ threshold.}
% \label{fig:bs_mlp_nobn}
% \end{figure*}
% \begin{figure*}[t]
% \centering
% \begin{tabular}{c@{}c@{}c@{}c@{}c}
% \includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_01_batch_size.pdf} &
% \includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_05_batch_size.pdf} &
% \includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_1_batch_size.pdf} &
% \includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_3_batch_size.pdf}
% \\
% \includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_01_batch_size.pdf} &
% \includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_05_batch_size.pdf} &
% \includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_1_batch_size.pdf} &
% \includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_3_batch_size.pdf}
% \\
% {\small $\mu=0.01$} & {\small $\mu=0.05$} & {\small $\mu=0.1$} & {\small $\mu=0.3$}
% \end{tabular}
% \caption{{\bf Average ranks and accuracy rates of {\color{blue}MLP-BN-10-100 (no-momentum)} trained on CIFAR10 with various batch sizes {\color{blue}($\epsilon=0.001$)}.} See Fig.~\ref{fig:bs} for details.}
% \end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/rank_mlp_bn_bs_4_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/rank_mlp_bn_bs_8_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/rank_mlp_bn_bs_16_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/rank_mlp_bn_bs_32_lr_1_weight_decay.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/acc_mlp_bn_bs_4_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/acc_mlp_bn_bs_8_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/acc_mlp_bn_bs_16_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/acc_mlp_bn_bs_32_lr_1_weight_decay.pdf}
\\
{\small $B=4$} & {\small $B=8$} & {\small $B=16$} & {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of MLP-BN-10-100 trained on CIFAR10 with varying $\lambda$.} In this experiment, $\mu=0.1$ and $\epsilon=\expnumber{1}{-3}$.}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/rank_mlp_no_bn_lr_05_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/rank_mlp_no_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/rank_mlp_no_bn_lr_2_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/rank_mlp_no_bn_lr_3_batch_size.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/acc_mlp_no_bn_lr_05_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/acc_mlp_no_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/acc_mlp_no_bn_lr_2_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/acc_mlp_no_bn_lr_3_batch_size.pdf}
\\
{\small $\mu=0.05$} & {\small $\mu=0.1$} & {\small $\mu=0.2$} & {\small $\mu=0.3$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of MLP-5-500 trained on CIFAR10 with various batch sizes.} In this experiment, $\lambda=\expnumber{5}{-4}$ and $\epsilon=\expnumber{1}{-3}$.}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/rank_mlp_no_bn_bs_4_lr_025_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/rank_mlp_no_bn_bs_8_lr_025_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/rank_mlp_no_bn_bs_16_lr_025_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/rank_mlp_no_bn_bs_32_lr_025_weight_decay.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/acc_mlp_no_bn_bs_4_lr_025_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/acc_mlp_no_bn_bs_8_lr_025_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/acc_mlp_no_bn_bs_16_lr_025_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/acc_mlp_no_bn_bs_32_lr_025_weight_decay.pdf}
\\
{\small $B=4$} & {\small $B=8$} & {\small $B=16$} & {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of MLP-5-500 trained on CIFAR10 with varying $\lambda$.} The models were trained with a $\mu=0.025$ initial learning rate. To estimate the rank, we used an $\epsilon=0.001$ threshold.}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/rank_resnet_bn_lr_5_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/rank_resnet_bn_lr_8_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/rank_resnet_bn_lr_10_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/rank_resnet_bn_lr_20_batch_size.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/acc_resnet_bn_lr_5_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/acc_resnet_bn_lr_8_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/acc_resnet_bn_lr_10_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/acc_resnet_bn_lr_20_batch_size.pdf}
\\
{\small $\mu=0.5$} & {\small $\mu=0.8$} & {\small $\mu=1.0$} & {\small $\mu=2.0$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of ResNet-18 trained on CIFAR10 with various batch sizes.} In this experiment, $\lambda=\expnumber{5}{-4}$ and $\epsilon=\expnumber{1}{-3}$.}
\label{fig:bs_resnet18}
\end{figure*}
% \begin{figure*}[t]
% \centering
% \end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/rank_vgg_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/rank_vgg_bn_lr_2_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/rank_vgg_bn_lr_3_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/rank_vgg_bn_lr_5_batch_size.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/acc_vgg_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/acc_vgg_bn_lr_2_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/acc_vgg_bn_lr_3_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/acc_vgg_bn_lr_5_batch_size.pdf}
\\
{\small $\mu=0.1$} & {\small $\mu=0.2$} & {\small $\mu=0.3$} & {\small $\mu=0.5$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of VGG-16 trained on CIFAR10 with various batch sizes.} In this experiment, $\mu=\expnumber{5}{-4}$ and $\epsilon=\expnumber{1}{-3}$.}
\label{fig:bs_vgg}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/rank_vgg_bn_batch_size_4_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/rank_vgg_bn_batch_size_8_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/rank_vgg_bn_batch_size_16_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/rank_vgg_bn_batch_size_32_lr_1_weight_decay.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/acc_vgg_bn_batch_size_4_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/acc_vgg_bn_batch_size_8_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/acc_vgg_bn_batch_size_16_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/acc_vgg_bn_batch_size_32_lr_1_weight_decay.pdf}
\\
{\small $B=4$} & {\small $B=8$} & {\small $B=16$} & {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of VGG-16 trained on CIFAR10 with varying $\lambda$.} In this experiment, $\mu=0.1$ and $\epsilon=0.01$.}
\label{fig:wd_vgg}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/rank_vgg_bn_batch_size_4_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/rank_vgg_bn_batch_size_8_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/rank_vgg_bn_batch_size_16_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/rank_vgg_bn_batch_size_32_learning_rate.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/acc_vgg_bn_batch_size_4_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/acc_vgg_bn_batch_size_8_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/acc_vgg_bn_batch_size_16_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/acc_vgg_bn_batch_size_32_learning_rate.pdf}
\\
{\small $B=4$} & {\small $B=8$} & {\small $B=16$} & {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of VGG-16 trained on CIFAR10 with varying $\mu$.} In this experiment, $\mu=\expnumber{5}{-4}$ and $\epsilon=\expnumber{1}{-3}$.}
\label{fig:lr_vgg}
\end{figure*}
\begin{figure*}[t] % vit weight decay
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
%\includegraphics[width=0.20\linewidth]{plots/vit/weight_decay/varywd_B8_avgrank.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/weight_decay/varywd_B16_avgrank.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/weight_decay/varywd_B32_avgrank.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/weight_decay/varywd_B64_avgrank.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/weight_decay/varywd_B128_avgrank.pdf}
\\
%\includegraphics[width=0.20\linewidth]{plots/vit/weight_decay/varywd_B8_accuracy.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/weight_decay/varywd_B16_accuracy.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/weight_decay/varywd_B32_accuracy.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/weight_decay/varywd_B64_accuracy.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/weight_decay/varywd_B128_accuracy.pdf}
\\
%{\small $B=8$} &
{\small $B=16$} & {\small $B=32$} & {\small $B=64$} & {\small $B=128$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of ViT trained on CIFAR10 with varying $\lambda$.} In this experiment, $\mu=\expnumber{4}{-2}$ and $\epsilon=0.01$.}
\label{fig:vit_var_wd}
\end{figure*}
\end{document}
% This document was modified from the file originally made available by
% Pat Langley and Andrea Danyluk for ICML-2K. This version was created
% by Iain Murray in 2018, and modified by Alexandre Bouchard in
% 2019 and 2021 and by Csaba Szepesvari, Gang Niu and Sivan Sabato in 2022.
% Modified again in 2023 by Sivan Sabato and Jonathan Scarlett.
% Previous contributors include Dan Roy, Lise Getoor and Tobias
% Scheffer, which was slightly modified from the 2010 version by
% Thorsten Joachims & Johannes Fuernkranz, slightly modified from the
% 2009 version by Kiri Wagstaff and Sam Roweis's 2008 version, which is
% slightly modified from Prasad Tadepalli's 2007 version which is a
% lightly changed version of the previous year's version by Andrew
% Moore, which was in turn edited from those of Kristian Kersting and
% Codrina Lauth. Alex Smola contributed to the algorithmic style files.
%%%%% NEW MATH DEFINITIONS %%%%%
\usepackage{amsmath,amsfonts,bm}
% Mark sections of captions for referring to divisions of figures
\newcommand{\figleft}{{\em (Left)}}
\newcommand{\figcenter}{{\em (Center)}}
\newcommand{\figright}{{\em (Right)}}
\newcommand{\figtop}{{\em (Top)}}
\newcommand{\figbottom}{{\em (Bottom)}}
\newcommand{\captiona}{{\em (a)}}
\newcommand{\captionb}{{\em (b)}}
\newcommand{\captionc}{{\em (c)}}
\newcommand{\captiond}{{\em (d)}}
% Highlight a newly defined term
\newcommand{\newterm}[1]{{\bf #1}}
% Figure reference, lower-case.
\def\figref#1{figure~\ref{#1}}
% Figure reference, capital. For start of sentence
\def\Figref#1{Figure~\ref{#1}}
\def\twofigref#1#2{figures \ref{#1} and \ref{#2}}
\def\quadfigref#1#2#3#4{figures \ref{#1}, \ref{#2}, \ref{#3} and \ref{#4}}
% Section reference, lower-case.
\def\secref#1{section~\ref{#1}}
% Section reference, capital.
\def\Secref#1{Section~\ref{#1}}
% Reference to two sections.
\def\twosecrefs#1#2{sections \ref{#1} and \ref{#2}}
% Reference to three sections.
\def\secrefs#1#2#3{sections \ref{#1}, \ref{#2} and \ref{#3}}
% Reference to an equation, lower-case.
\def\eqref#1{equation~\ref{#1}}
% Reference to an equation, upper case
\def\Eqref#1{Equation~\ref{#1}}
% A raw reference to an equation---avoid using if possible
\def\plaineqref#1{\ref{#1}}
% Reference to a chapter, lower-case.
\def\chapref#1{chapter~\ref{#1}}
% Reference to an equation, upper case.
\def\Chapref#1{Chapter~\ref{#1}}
% Reference to a range of chapters
\def\rangechapref#1#2{chapters\ref{#1}--\ref{#2}}
% Reference to an algorithm, lower-case.
\def\algref#1{algorithm~\ref{#1}}
% Reference to an algorithm, upper case.
\def\Algref#1{Algorithm~\ref{#1}}
\def\twoalgref#1#2{algorithms \ref{#1} and \ref{#2}}
\def\Twoalgref#1#2{Algorithms \ref{#1} and \ref{#2}}
% Reference to a part, lower case
\def\partref#1{part~\ref{#1}}
% Reference to a part, upper case
\def\Partref#1{Part~\ref{#1}}
\def\twopartref#1#2{parts \ref{#1} and \ref{#2}}
\def\ceil#1{\lceil #1 \rceil}
\def\floor#1{\lfloor #1 \rfloor}
\def\1{\bm{1}}
\newcommand{\train}{\mathcal{D}}
\newcommand{\valid}{\mathcal{D_{\mathrm{valid}}}}
\newcommand{\test}{\mathcal{D_{\mathrm{test}}}}
\def\eps{{\epsilon}}
% Random variables
\def\reta{{\textnormal{$\eta$}}}
\def\ra{{\textnormal{a}}}
\def\rb{{\textnormal{b}}}
\def\rc{{\textnormal{c}}}
\def\rd{{\textnormal{d}}}
\def\re{{\textnormal{e}}}
\def\rf{{\textnormal{f}}}
\def\rg{{\textnormal{g}}}
\def\rh{{\textnormal{h}}}
\def\ri{{\textnormal{i}}}
\def\rj{{\textnormal{j}}}
\def\rk{{\textnormal{k}}}
\def\rl{{\textnormal{l}}}
% rm is already a command, just don't name any random variables m
\def\rn{{\textnormal{n}}}
\def\ro{{\textnormal{o}}}
\def\rp{{\textnormal{p}}}
\def\rq{{\textnormal{q}}}
\def\rr{{\textnormal{r}}}
\def\rs{{\textnormal{s}}}
\def\rt{{\textnormal{t}}}
\def\ru{{\textnormal{u}}}
\def\rv{{\textnormal{v}}}
\def\rw{{\textnormal{w}}}
\def\rx{{\textnormal{x}}}
\def\ry{{\textnormal{y}}}
\def\rz{{\textnormal{z}}}
% Random vectors
\def\rvepsilon{{\mathbf{\epsilon}}}
\def\rvtheta{{\mathbf{\theta}}}
\def\rva{{\mathbf{a}}}
\def\rvb{{\mathbf{b}}}
\def\rvc{{\mathbf{c}}}
\def\rvd{{\mathbf{d}}}
\def\rve{{\mathbf{e}}}
\def\rvf{{\mathbf{f}}}
\def\rvg{{\mathbf{g}}}
\def\rvh{{\mathbf{h}}}
\def\rvu{{\mathbf{i}}}
\def\rvj{{\mathbf{j}}}
\def\rvk{{\mathbf{k}}}
\def\rvl{{\mathbf{l}}}
\def\rvm{{\mathbf{m}}}
\def\rvn{{\mathbf{n}}}
\def\rvo{{\mathbf{o}}}
\def\rvp{{\mathbf{p}}}
\def\rvq{{\mathbf{q}}}
\def\rvr{{\mathbf{r}}}
\def\rvs{{\mathbf{s}}}
\def\rvt{{\mathbf{t}}}
\def\rvu{{\mathbf{u}}}
\def\rvv{{\mathbf{v}}}
\def\rvw{{\mathbf{w}}}
\def\rvx{{\mathbf{x}}}
\def\rvy{{\mathbf{y}}}
\def\rvz{{\mathbf{z}}}
% Elements of random vectors
\def\erva{{\textnormal{a}}}
\def\ervb{{\textnormal{b}}}
\def\ervc{{\textnormal{c}}}
\def\ervd{{\textnormal{d}}}
\def\erve{{\textnormal{e}}}
\def\ervf{{\textnormal{f}}}
\def\ervg{{\textnormal{g}}}
\def\ervh{{\textnormal{h}}}
\def\ervi{{\textnormal{i}}}
\def\ervj{{\textnormal{j}}}
\def\ervk{{\textnormal{k}}}
\def\ervl{{\textnormal{l}}}
\def\ervm{{\textnormal{m}}}
\def\ervn{{\textnormal{n}}}
\def\ervo{{\textnormal{o}}}
\def\ervp{{\textnormal{p}}}
\def\ervq{{\textnormal{q}}}
\def\ervr{{\textnormal{r}}}
\def\ervs{{\textnormal{s}}}
\def\ervt{{\textnormal{t}}}
\def\ervu{{\textnormal{u}}}
\def\ervv{{\textnormal{v}}}
\def\ervw{{\textnormal{w}}}
\def\ervx{{\textnormal{x}}}
\def\ervy{{\textnormal{y}}}
\def\ervz{{\textnormal{z}}}
% Random matrices
\def\rmA{{\mathbf{A}}}
\def\rmB{{\mathbf{B}}}
\def\rmC{{\mathbf{C}}}
\def\rmD{{\mathbf{D}}}
\def\rmE{{\mathbf{E}}}
\def\rmF{{\mathbf{F}}}
\def\rmG{{\mathbf{G}}}
\def\rmH{{\mathbf{H}}}
\def\rmI{{\mathbf{I}}}
\def\rmJ{{\mathbf{J}}}
\def\rmK{{\mathbf{K}}}
\def\rmL{{\mathbf{L}}}
\def\rmM{{\mathbf{M}}}
\def\rmN{{\mathbf{N}}}
\def\rmO{{\mathbf{O}}}
\def\rmP{{\mathbf{P}}}
\def\rmQ{{\mathbf{Q}}}
\def\rmR{{\mathbf{R}}}
\def\rmS{{\mathbf{S}}}
\def\rmT{{\mathbf{T}}}
\def\rmU{{\mathbf{U}}}
\def\rmV{{\mathbf{V}}}
\def\rmW{{\mathbf{W}}}
\def\rmX{{\mathbf{X}}}
\def\rmY{{\mathbf{Y}}}
\def\rmZ{{\mathbf{Z}}}
% Elements of random matrices
\def\ermA{{\textnormal{A}}}
\def\ermB{{\textnormal{B}}}
\def\ermC{{\textnormal{C}}}
\def\ermD{{\textnormal{D}}}
\def\ermE{{\textnormal{E}}}
\def\ermF{{\textnormal{F}}}
\def\ermG{{\textnormal{G}}}
\def\ermH{{\textnormal{H}}}
\def\ermI{{\textnormal{I}}}
\def\ermJ{{\textnormal{J}}}
\def\ermK{{\textnormal{K}}}
\def\ermL{{\textnormal{L}}}
\def\ermM{{\textnormal{M}}}
\def\ermN{{\textnormal{N}}}
\def\ermO{{\textnormal{O}}}
\def\ermP{{\textnormal{P}}}
\def\ermQ{{\textnormal{Q}}}
\def\ermR{{\textnormal{R}}}
\def\ermS{{\textnormal{S}}}
\def\ermT{{\textnormal{T}}}
\def\ermU{{\textnormal{U}}}
\def\ermV{{\textnormal{V}}}
\def\ermW{{\textnormal{W}}}
\def\ermX{{\textnormal{X}}}
\def\ermY{{\textnormal{Y}}}
\def\ermZ{{\textnormal{Z}}}
% Vectors
\def\vzero{{\bm{0}}}
\def\vone{{\bm{1}}}
\def\vmu{{\bm{\mu}}}
\def\vtheta{{\bm{\theta}}}
\def\va{{\bm{a}}}
\def\vb{{\bm{b}}}
\def\vc{{\bm{c}}}
\def\vd{{\bm{d}}}
\def\ve{{\bm{e}}}
\def\vf{{\bm{f}}}
\def\vg{{\bm{g}}}
\def\vh{{\bm{h}}}
\def\vi{{\bm{i}}}
\def\vj{{\bm{j}}}
\def\vk{{\bm{k}}}
\def\vl{{\bm{l}}}
\def\vm{{\bm{m}}}
\def\vn{{\bm{n}}}
\def\vo{{\bm{o}}}
\def\vp{{\bm{p}}}
\def\vq{{\bm{q}}}
\def\vr{{\bm{r}}}
\def\vs{{\bm{s}}}
\def\vt{{\bm{t}}}
\def\vu{{\bm{u}}}
\def\vv{{\bm{v}}}
\def\vw{{\bm{w}}}
\def\vx{{\bm{x}}}
\def\vy{{\bm{y}}}
\def\vz{{\bm{z}}}
% Elements of vectors
\def\evalpha{{\alpha}}
\def\evbeta{{\beta}}
\def\evepsilon{{\epsilon}}
\def\evlambda{{\lambda}}
\def\evomega{{\omega}}
\def\evmu{{\mu}}
\def\evpsi{{\psi}}
\def\evsigma{{\sigma}}
\def\evtheta{{\theta}}
\def\eva{{a}}
\def\evb{{b}}
\def\evc{{c}}
\def\evd{{d}}
\def\eve{{e}}
\def\evf{{f}}
\def\evg{{g}}
\def\evh{{h}}
\def\evi{{i}}
\def\evj{{j}}
\def\evk{{k}}
\def\evl{{l}}
\def\evm{{m}}
\def\evn{{n}}
\def\evo{{o}}
\def\evp{{p}}
\def\evq{{q}}
\def\evr{{r}}
\def\evs{{s}}
\def\evt{{t}}
\def\evu{{u}}
\def\evv{{v}}
\def\evw{{w}}
\def\evx{{x}}
\def\evy{{y}}
\def\evz{{z}}
% Matrix
\def\mA{{\bm{A}}}
\def\mB{{\bm{B}}}
\def\mC{{\bm{C}}}
\def\mD{{\bm{D}}}
\def\mE{{\bm{E}}}
\def\mF{{\bm{F}}}
\def\mG{{\bm{G}}}
\def\mH{{\bm{H}}}
\def\mI{{\bm{I}}}
\def\mJ{{\bm{J}}}
\def\mK{{\bm{K}}}
\def\mL{{\bm{L}}}
\def\mM{{\bm{M}}}
\def\mN{{\bm{N}}}
\def\mO{{\bm{O}}}
\def\mP{{\bm{P}}}
\def\mQ{{\bm{Q}}}
\def\mR{{\bm{R}}}
\def\mS{{\bm{S}}}
\def\mT{{\bm{T}}}
\def\mU{{\bm{U}}}
\def\mV{{\bm{V}}}
\def\mW{{\bm{W}}}
\def\mX{{\bm{X}}}
\def\mY{{\bm{Y}}}
\def\mZ{{\bm{Z}}}
\def\mBeta{{\bm{\beta}}}
\def\mPhi{{\bm{\Phi}}}
\def\mLambda{{\bm{\Lambda}}}
\def\mSigma{{\bm{\Sigma}}}
% Tensor
\DeclareMathAlphabet{\mathsfit}{\encodingdefault}{\sfdefault}{m}{sl}
\SetMathAlphabet{\mathsfit}{bold}{\encodingdefault}{\sfdefault}{bx}{n}
\newcommand{\tens}[1]{\bm{\mathsfit{#1}}}
\def\tA{{\tens{A}}}
\def\tB{{\tens{B}}}
\def\tC{{\tens{C}}}
\def\tD{{\tens{D}}}
\def\tE{{\tens{E}}}
\def\tF{{\tens{F}}}
\def\tG{{\tens{G}}}
\def\tH{{\tens{H}}}
\def\tI{{\tens{I}}}
\def\tJ{{\tens{J}}}
\def\tK{{\tens{K}}}
\def\tL{{\tens{L}}}
\def\tM{{\tens{M}}}
\def\tN{{\tens{N}}}
\def\tO{{\tens{O}}}
\def\tP{{\tens{P}}}
\def\tQ{{\tens{Q}}}
\def\tR{{\tens{R}}}
\def\tS{{\tens{S}}}
\def\tT{{\tens{T}}}
\def\tU{{\tens{U}}}
\def\tV{{\tens{V}}}
\def\tW{{\tens{W}}}
\def\tX{{\tens{X}}}
\def\tY{{\tens{Y}}}
\def\tZ{{\tens{Z}}}
% Graph
\def\gA{{\mathcal{A}}}
\def\gB{{\mathcal{B}}}
\def\gC{{\mathcal{C}}}
\def\gD{{\mathcal{D}}}
\def\gE{{\mathcal{E}}}
\def\gF{{\mathcal{F}}}
\def\gG{{\mathcal{G}}}
\def\gH{{\mathcal{H}}}
\def\gI{{\mathcal{I}}}
\def\gJ{{\mathcal{J}}}
\def\gK{{\mathcal{K}}}
\def\gL{{\mathcal{L}}}
\def\gM{{\mathcal{M}}}
\def\gN{{\mathcal{N}}}
\def\gO{{\mathcal{O}}}
\def\gP{{\mathcal{P}}}
\def\gQ{{\mathcal{Q}}}
\def\gR{{\mathcal{R}}}
\def\gS{{\mathcal{S}}}
\def\gT{{\mathcal{T}}}
\def\gU{{\mathcal{U}}}
\def\gV{{\mathcal{V}}}
\def\gW{{\mathcal{W}}}
\def\gX{{\mathcal{X}}}
\def\gY{{\mathcal{Y}}}
\def\gZ{{\mathcal{Z}}}
% Sets
\def\sA{{\mathbb{A}}}
\def\sB{{\mathbb{B}}}
\def\sC{{\mathbb{C}}}
\def\sD{{\mathbb{D}}}
% Don't use a set called E, because this would be the same as our symbol
% for expectation.
\def\sF{{\mathbb{F}}}
\def\sG{{\mathbb{G}}}
\def\sH{{\mathbb{H}}}
\def\sI{{\mathbb{I}}}
\def\sJ{{\mathbb{J}}}
\def\sK{{\mathbb{K}}}
\def\sL{{\mathbb{L}}}
\def\sM{{\mathbb{M}}}
\def\sN{{\mathbb{N}}}
\def\sO{{\mathbb{O}}}
\def\sP{{\mathbb{P}}}
\def\sQ{{\mathbb{Q}}}
\def\sR{{\mathbb{R}}}
\def\sS{{\mathbb{S}}}
\def\sT{{\mathbb{T}}}
\def\sU{{\mathbb{U}}}
\def\sV{{\mathbb{V}}}
\def\sW{{\mathbb{W}}}
\def\sX{{\mathbb{X}}}
\def\sY{{\mathbb{Y}}}
\def\sZ{{\mathbb{Z}}}
% Entries of a matrix
\def\emLambda{{\Lambda}}
\def\emA{{A}}
\def\emB{{B}}
\def\emC{{C}}
\def\emD{{D}}
\def\emE{{E}}
\def\emF{{F}}
\def\emG{{G}}
\def\emH{{H}}
\def\emI{{I}}
\def\emJ{{J}}
\def\emK{{K}}
\def\emL{{L}}
\def\emM{{M}}
\def\emN{{N}}
\def\emO{{O}}
\def\emP{{P}}
\def\emQ{{Q}}
\def\emR{{R}}
\def\emS{{S}}
\def\emT{{T}}
\def\emU{{U}}
\def\emV{{V}}
\def\emW{{W}}
\def\emX{{X}}
\def\emY{{Y}}
\def\emZ{{Z}}
\def\emSigma{{\Sigma}}
% entries of a tensor
% Same font as tensor, without \bm wrapper
\newcommand{\etens}[1]{\mathsfit{#1}}
\def\etLambda{{\etens{\Lambda}}}
\def\etA{{\etens{A}}}
\def\etB{{\etens{B}}}
\def\etC{{\etens{C}}}
\def\etD{{\etens{D}}}
\def\etE{{\etens{E}}}
\def\etF{{\etens{F}}}
\def\etG{{\etens{G}}}
\def\etH{{\etens{H}}}
\def\etI{{\etens{I}}}
\def\etJ{{\etens{J}}}
\def\etK{{\etens{K}}}
\def\etL{{\etens{L}}}
\def\etM{{\etens{M}}}
\def\etN{{\etens{N}}}
\def\etO{{\etens{O}}}
\def\etP{{\etens{P}}}
\def\etQ{{\etens{Q}}}
\def\etR{{\etens{R}}}
\def\etS{{\etens{S}}}
\def\etT{{\etens{T}}}
\def\etU{{\etens{U}}}
\def\etV{{\etens{V}}}
\def\etW{{\etens{W}}}
\def\etX{{\etens{X}}}
\def\etY{{\etens{Y}}}
\def\etZ{{\etens{Z}}}
% The true underlying data generating distribution
\newcommand{\pdata}{p_{\rm{data}}}
% The empirical distribution defined by the training set
\newcommand{\ptrain}{\hat{p}_{\rm{data}}}
\newcommand{\Ptrain}{\hat{P}_{\rm{data}}}
% The model distribution
\newcommand{\pmodel}{p_{\rm{model}}}
\newcommand{\Pmodel}{P_{\rm{model}}}
\newcommand{\ptildemodel}{\tilde{p}_{\rm{model}}}
% Stochastic autoencoder distributions
\newcommand{\pencode}{p_{\rm{encoder}}}
\newcommand{\pdecode}{p_{\rm{decoder}}}
\newcommand{\precons}{p_{\rm{reconstruct}}}
\newcommand{\laplace}{\mathrm{Laplace}} % Laplace distribution
\newcommand{\E}{\mathbb{E}}
\newcommand{\Ls}{\mathcal{L}}
\newcommand{\R}{\mathbb{R}}
\newcommand{\emp}{\tilde{p}}
\newcommand{\lr}{\alpha}
\newcommand{\reg}{\lambda}
\newcommand{\rect}{\mathrm{rectifier}}
\newcommand{\softmax}{\mathrm{softmax}}
\newcommand{\sigmoid}{\sigma}
\newcommand{\softplus}{\zeta}
\newcommand{\KL}{D_{\mathrm{KL}}}
\newcommand{\Var}{\mathrm{Var}}
\newcommand{\standarderror}{\mathrm{SE}}
\newcommand{\Cov}{\mathrm{Cov}}
% Wolfram Mathworld says $L^2$ is for function spaces and $\ell^2$ is for vectors
% But then they seem to use $L^2$ for vectors throughout the site, and so does
% wikipedia.
\newcommand{\normlzero}{L^0}
\newcommand{\normlone}{L^1}
\newcommand{\normltwo}{L^2}
\newcommand{\normlp}{L^p}
\newcommand{\normmax}{L^\infty}
\newcommand{\parents}{Pa} % See usage in notation.tex. Chosen to match Daphne's book.
\DeclareMathOperator*{\argmax}{arg\,max}
\DeclareMathOperator*{\argmin}{arg\,min}
\DeclareMathOperator{\sign}{sign}
\DeclareMathOperator{\Tr}{Tr}
\let\ab\allowbreak
\documentclass{article} % For LaTeX2e
\usepackage{iclr2023_conference,times}
% Optional math commands from https://github.com/goodfeli/dlbook_notation.
\input{math_commands.tex}
\usepackage{hyperref}
\usepackage{url}
\usepackage{amsmath,amssymb,amsthm,bbm}
\usepackage{thm-restate,enumitem}
\usepackage{tabularx}
\usepackage{rotating}
\newtheorem{theorem}{Theorem}
\newtheorem{lemma}{Lemma}
\newtheorem{proposition}{Proposition}
\newtheorem{corollary}{Corollary}
\newtheorem{definition}{Definition}
\newtheorem{remark}{Remark}
\newcommand{\expnumber}[2]{{#1}\mathrm{e}{#2}}
\newcommand{\cC}{\mathcal{C}}
\newcommand{\cD}{\mathcal{D}}
\newcommand{\cF}{\mathcal{F}}
\newcommand{\cG}{\mathcal{G}}
\newcommand{\cH}{\mathcal{H}}
\newcommand{\cN}{\mathcal{N}}
\newcommand{\cP}{\mathcal{P}}
\newcommand{\cU}{\mathcal{U}}
\newcommand{\cX}{\mathcal{X}}
\newcommand{\cY}{\mathcal{Y}}
\renewcommand{\P}{\mathbb{P}}
\newcommand{\B}{\mathbb{B}}
\newcommand{\bI}{\mathbb{I}}
\newcommand{\EE}[1]{\E\left[ #1\right]}
\newcommand{\PP}[1]{\P\left[ #1\right]}
\newcommand{\I}[1]{\bI\left[ #1\right]}
\newcommand{\tiX}{\tilde{X}}
\newcommand{\tiS}{\tilde{S}}
\newcommand{\ty}{\tilde{y}}
\newcommand{\tcG}{\tilde{\cG}}
\newcommand{\tcP}{\tilde{\cP}}
\newcommand{\tcU}{\tilde{\cU}}
\newcommand{\tf}{\tilde{f}}
\newcommand{\tig}{\tilde{g}}
\newcommand{\tih}{\tilde{h}}
\newcommand{\clD}{\cD_{\cC}}
\newcommand{\one}{\mathbbm{1}}
\newcommand{\rank}{\textnormal{rank}}
\newcommand{\vect}{\textnormal{vec}}
\newcommand{\threepartdef}[6]
{
\left\{
\begin{array}{lll}
#1 & \mbox{if } #2 \\
#3 & \mbox{if } #4 \\
#5 & \mbox{if } #6
\end{array}
\right.
}
\newcommand{\twopartdef}[4]
{
\left\{
\begin{array}{lll}
#1 & \mbox{if } #2 \\
#3 & \mbox{if } #4 \\
\end{array}
\right.
}
\def\fr#1#2{{\textstyle\frac{#1}{#2}}}
\def\draft{1}
\ifnum\draft=1
\def\ShowAuthNotes{1}
\else
\def\ShowAuthNotes{0}
\fi
\ifnum\ShowAuthNotes=1
\newcommand{\authnote}[3]{\textcolor{#3}{[{\footnotesize {\bf #1:} { {#2}}}]}}
\else
\newcommand{\authnote}[3]{}
\fi
\newcommand{\vnote}[1]{\authnote{Aparna}{#1}{BrickRed}}
\title{SGD and Weight Decay Provably Induce a Low-Rank Bias in Neural Networks}
% Authors must not appear in the submitted version. They should be hidden
% as long as the \iclrfinalcopy macro remains commented out below.
% Non-anonymous submissions will be rejected without review.
\author{Antiquus S.~Hippocampus, Natalia Cerebro \& Amelie P. Amygdale \thanks{ Use footnote for providing further information
about author (webpage, alternative address)---\emph{not} for acknowledging
funding agencies. Funding acknowledgements go at the end of the paper.} \\
Department of Computer Science\\
Cranberry-Lemon University\\
Pittsburgh, PA 15213, USA \\
\texttt{\{hippo,brain,jen\}@cs.cranberry-lemon.edu} \\
\And
Ji Q. Ren \& Yevgeny LeNet \\
Department of Computational Neuroscience \\
University of the Witwatersrand \\
Joburg, South Africa \\
\texttt{\{robot,net\}@wits.ac.za} \\
\AND
Coauthor \\
Affiliation \\
Address \\
\texttt{email}
}
% The \author macro works with any number of authors. There are two commands
% used to separate the names and addresses of multiple authors: \And and \AND.
%
% Using \And between authors leaves it to \LaTeX{} to determine where to break
% the lines. Using \AND forces a linebreak at that point. So, if \LaTeX{}
% puts 3 of 4 authors names on the first line, and the last on the second
% line, try using \AND instead of \And before the third author name.
\newcommand{\fix}{\marginpar{FIX}}
\newcommand{\new}{\marginpar{NEW}}
%\iclrfinalcopy % Uncomment for camera-ready version, but NOT for submission.
\begin{document}
\maketitle
\begin{abstract}
We analyze deep ReLU neural networks trained with mini-batch Stochastic Gradient Descent (SGD) and weight decay. We show, both theoretically and empirically, that when training a neural network using SGD with weight decay and small batch size, the resulting weight matrices tend to be of small rank. Our analysis relies on a minimal set of assumptions; the neural networks may be arbitrarily wide or deep and may include residual connections, as well as convolutional layers.
The same analysis implies the inherent presence of SGD ``noise'', defined as the inability of SGD to converge to a stationary point. In particular, we prove that SGD noise must always be present, even asymptotically, as long as we incorporate weight decay and the batch size is smaller than the total number of training samples.
%TP%Furthermore, we study the source of SGD noise and prove that when training with weight decay, the only solutions of SGD at convergence are zero functions.
\end{abstract}
\section{Introduction}\label{sec:intro}
% Over the past few years, neural networks have challenged machine learning theory with several mysteries. We consider two of them here that are, to our knowledge, so far unsolved. The first is about the mathematical reasons for a bias in Stochastic Gradient Descent (SGD) towards neural networks with low-rank weight matrices, a fact which has been empirically observed (e.g., \cite{DBLP:journals/corr/abs-2201-12760,DBLP:journals/corr/abs-2103-10427}). The second is regarding the role and origin of SGD noise (e.g., \cite{DBLP:journals/corr/abs-2110-06914}). The origin of the SGD noise represents a puzzling and rather neglected question: {\em why does SGD, under certain conditions, never stop making updates?}
Stochastic gradient descent (SGD) is one of the standard workhorses for optimizing deep models~\citep{bottou-91c}. Though initially proposed to remedy the computational bottleneck of gradient descent (GD), recent studies suggest that SGD also induces crucial regularization, which prevents overparameterized models from converging to minima that cannot generalize well~\citep{DBLP:journals/corr/ZhangBHRV16,https://doi.org/10.48550/arxiv.1711.04623,https://doi.org/10.48550/arxiv.1609.04836,zhu}. Empirical studies suggest that SGD outperforms GD~\cite{zhu} and SGD generalizes better when used with smaller batch sizes~\citep{NIPS2017_a5e0ff62,https://doi.org/10.48550/arxiv.1609.04836}, and (iii) gradient descent with additional noise cannot compete with SGD~\cite{zhu}.
The full range of regularization effects induced by SGD, however, is not yet fully understood.
%Furthermore, we are unaware of any study of the noise in the output of the network during SGD training, a noise which seems to be present even at very long training times, suggesting that true convergence to a single set of weight matrices is never reached.
%\subsection{Contributions}
%{\bf Contributions.\enspace}
In this paper we present a mathematical analysis of the bias of SGD towards rank-minimization. To investigate this bias, we propose the {\em SGD Near-Convergence Regime} as a novel approach for investigating inductive biases of SGD-trained neural networks. This setting considers the case where SGD reaches a point in training where the expected update is small in comparison to the actual weights' norm. Our analysis is fairly generic: we consider deep ReLU networks trained with mini-batch SGD for minimizing a differentiable loss function with $L_2$ regularization (i.e., weight decay). The neural networks may include fully-connected layers, residual connections and convolutions.
Our main contributions are:
\begin{itemize}[leftmargin=0.5cm]
\item In Thm.~\ref{thm:batches}, we demonstrate that training neural networks with mini-batch SGD and weight decay results in a low-rank bias in their weight matrices. We theoretically demonstrate that when training with smaller batch sizes, the rank of the learned matrices tends to decrease. This observation is validated as part of an extensive empirical study of the effect of certain hyperparameters on the rank of learned matrices with various architectures.
\item In Sec.~\ref{sec:convergence}, we study the inherent inability of SGD to converge to a stationary point, that we call `SGD noise'. In Props.~\ref{prop:propA}-\ref{prop:propB} we describe conditions in which `SGD noise' is inevitable when training convolutional neural networks. In particular, we demonstrate that when training a fully-connected neural network, SGD noise must always be present, even asymptotically, as long as we incorporate weight decay and the batch size is smaller than the total number of samples. These predictions are empirically validated in Sec.~\ref{sec:sgd_noise_exp}.
\end{itemize}
\subsection{Related Work}\label{sec:related_work}
A prominent thread in the recent literature revolves around characterizing the implicit regularization of gradient-based optimization in the belief that this is key to generalization in deep learning. Several papers have focused on a potential bias of gradient descent or stochastic gradient descent towards rank minimization. The initial interest was motivated by the matrix factorization problem, which corresponds to training a depth-2 linear neural network with multiple outputs w.r.t. the square loss. \cite{https://doi.org/10.48550/arxiv.1705.09280} initially conjectured that the implicit regularization in matrix factorization can be characterized in terms of the nuclear norm of the corresponding linear predictor. This conjecture, however, was formally refuted by~\cite{DBLP:journals/corr/abs-2012-09839}. Later, \cite{DBLP:journals/corr/abs-2005-06398} conjectured that the implicit regularization in matrix factorization can be explained by rank minimization, and also hypothesized that some notion of rank minimization may be key to explaining generalization in deep learning. \cite{DBLP:journals/corr/abs-2012-09839} established evidence that the implicit regularization in matrix factorization is a heuristic for rank minimization. Beyond factorization problems, \cite{DBLP:journals/corr/abs-2006-06657} showed that gradient flow (GF) training of univariate linear networks w.r.t. exponentially-tailed classification losses learns weight matrices of rank 1.
With nonlinear neural networks, however, things are less clear. Empirically, a series of papers~\citep{NIPS2014_2afe4567,10.5555/3294771.3294853,s21165599,8099498,pmlr-v80-arora18b} showed that replacing the weight matrices by low-rank approximations results in only a small drop in accuracy. This suggests that the weight matrices at convergence may be close to low-rank matrices. However, whether they provably behave this way remains unclear. \cite{DBLP:journals/corr/abs-2201-12760} showed that for ReLU networks, GF generally does not minimize rank. They also argued that sufficiently deep ReLU networks can have low-rank solutions under $L_2$ norm minimization. This interesting result, however, applies to layers added to a network that already solves the problem and may not have any low-rank bias. It is not directly related to the mechanism described in this paper, which applies to all layers in the network but only in the presence of regularization and SGD unlike~\citep{DBLP:journals/corr/abs-2201-12760}. A recent paper~\citep{le2022training} analyzes low-rank bias in neural networks trained with GF (without regularization). While this paper makes significant strides in extending the analysis in~\citep{DBLP:journals/corr/abs-2006-06657}, it makes several limiting assumptions. As a result, their analysis is only applicable under very specific conditions, such as when the data is linearly separable, and their low-rank analysis is limited to a set of linear layers aggregated at the top of the trained network.
\section{Problem Setup}\label{sec:setup}
In this work we consider a standard supervised learning setting (classification or regression), and study the inductive biases induced by training a neural network with mini-batch SGD along with weight decay. Formally, the task is defined by a distribution $P$ over samples $(x,y)\in \cX\times \cY$, where $\cX \subset \R^{c_{1} \times h_1 \times w_1}$ is the instance space (e.g., images), and $\cY \subset \R^k$ is a label space.
We consider a parametric model $\cF \subset \{f':\cX\to \R^k\}$, where each function $f_W \in \cF$ is specified by a vector of parameters $W \in \R^{N}$. A function $f_W \in \cF$ assigns a prediction to any input point $x \in \cX$, and its performance is measured by the {\em Expected Risk}, $L_P(f_W)~:=~\E_{(x,y)\sim P}[\ell(f_W(x),y)]$,
%\begin{equation}
%\label{eq:loss}
%\end{equation}
where $\ell:\R^k \times \cY \to [0,\infty)$ is a non-negative, differentiable, loss function (e.g., MSE or cross-entropy losses). For simplicity, in the analysis we focus on the case where $k=1$.
Since we do not have direct access to the full population distribution $P$, the goal is to learn a predictor, $f_W$, from some training dataset $S = \{(x_i,y_i)\}_{i=1}^m$ of independent and identically distributed (i.i.d.) samples drawn from $P$. Traditionally, in order to avoid overfitting the training data, we typically employ weight decay in order to control the complexity of the learned model. Namely, we intend to minimize the {\em Regularized Empirical Risk}, $L^{\lambda}_S(f_W)~:=~\fr{1}{m}\sum^{m}_{i=1}\ell(f_W(x_i),y_i) + \lambda \|W\|^2_2$,
% \begin{equation}
% ,
% \end{equation}
where $\lambda > 0$ is predefined hyperparameter. In order to minimize this objective, we typically use mini-batch SGD, as detailed below.
{\bf Optimization.\enspace} In this work, we minimize the regularized empirical risk $L^{\lambda}_{S}(f_{W})$ by applying stochastic gradient descent (SGD) for a certain number of iterations $T$. Formally, we initialize $W_0$ using a standard initialization procedure, iteratively update $W_t$ for $T$ iterations and return $W_T$. At each iteration, we sample a subset $\tiS = \{(x_{ij},y_{ij})\}^{B}_{j=1} \subset S$ uniformly at random and update $W_{t+1} \leftarrow W_t - \mu \nabla_W L^{\lambda}_{\tiS}(f_{W_t})$, where $\mu > 0$ is a predefined learning rate.
{\bf Notation.\enspace} Throughout the paper, we use the following notations. For an integer $k\geq 1$, $[k]=\{1,\ldots,k\}$. $\|z\|$ denotes the Euclidean norm. For two given vectors $x \in \R^{n}, y \in \R^{m}$ we define their concatenation as follows $(x\|y) := (x_1,\dots,x_n,y_1,\dots,y_m) \in \R^{n+m}$. For a given matrix $x \in \R^{n \times m}$, we denote $x_{i,:}$ its $i$'th row and by $\vect(x) := (x_{1}\|\dots \| x_n)$ its vectorization. For a given tensor $x \in \R^{c \times h\times w}$, we denote by $\vect(x) := (\vect(x_1)\| \dots \| \vect(x_c))$ the vectorized form of $x$. We define tensor slicing as follows, $x_{a:b} := (x_a,\dots,x_b)$.
\subsection{Architectures}\label{sec:archs}
In this work, the function $f_W$ represents a neural network, consisting of a set of layers of weights interlaced with ReLU activation units. We employ a fairly generic definition of a neural network, that includes convolutional layers, pooling layers, residual connections and fully-connected layers.
{\bf Network architecture.\enspace} Formally, $f_W$ is as a directed acyclic graph (DAG) $G=(V,E)$, where $V= \{v_1,\dots,v_L\}$ consists of the various layers of the network and each edge $e_{ij}=(v_i,v_j) \in E$ specifies a connection between two layers. Each layer is a function $v_{i}:\R^{c_{1} \times h_1 \times w_1} \to \R^{c_{i} \times h_i \times w_i}$ and each connection $(v_i,v_j)$ holds a transformation $C^{ij} : \R^{c_{j} \times h_j \times w_j} \to \R^{c_i \times h_i \times w_i}$. The layers are divided into three categories: {\bf (i)} the input layer $v_1$, {\bf (ii)} the output layer $v_L$ and {\bf (iii)} intermediate layers. In this setting, we do not have connections directed towards the input layer nor connections coming out of the output layer (i.e., $\forall~i \in [L]:~(v_L,v_i), (v_0,v_i) \notin E$). Given an input $x \in \R^{c_i \times h_i \times w_i}$, the output of a given layer $v_i$ is evaluated as follows $v_i(x) := \sigma(\sum_{j \in \textnormal{pred}(i)} C^{ij}(v_{j}(x)))$, except for the output layer $v_L$ that computes $f_W(x) := v_L(x) := \sum_{j \in \textnormal{pred}(L)}C^{Lj} (v_{j}(x))$. Here, $\textnormal{pred}(i):=\{j \in [L]\mid (v_i,v_j)\in E\}$, $\textnormal{succ}(i):=\{j \in [L]\mid (v_j,v_i)\in E\}$ and $\sigma$ is the ReLU activation function. Each transformation $C^{ij}$ is either trainable (e.g., a convolutional layer) or a constant affine transformation (e.g., a residual connection). We denote by $E_T$ the set of trainable connections. In this paper, we consider the following transformations.
{\bf Convolutional layers.\enspace} A convolutional layer~\citep{726791} (see also~\cite{Goodfellow-et-al-2016}) $C^{ij}:\R^{c_j \times h_j \times w_j} \to \R^{c_i \times h_i \times w_i}$ with kernel sizes $(k_1,k_2)$, padding $p$ and stride $s$ is parameterized by a tensor $Z^{ij} \in \R^{c_{i}\times c_j \times k_1 \times k_2}$ and computes the following tensor as output
\begin{small}
\begin{equation}\label{eq:conv}
\forall~ (c,t,l) \in [c_i]\times [h_i] \times [w_i]:~y_{c,t,l} ~=~ \sum^{c_{j}}_{c'=1} \vect(Z_{c,c',~:~})^{\top} \cdot \vect(\textnormal{Pad}_p(x)_{c',ts~:~(t+1)s+k_1 , ls~:~(l+1)s+k_2}).
\end{equation}
\end{small}
Here, $\textnormal{Pad}_p$ takes a tensor $x \in \R^{c \times h \times w}$ and returns a new tensor $x' \in \R^{c \times (h+2p) \times (w+2p)}$, where the first and last $p$ rows and columns of each channel $x'_{c,~:~,~:~}$ are zeros and the middle $1 \times h \times w$ tensor is equal to $x_{c,~:~,~:~}$. The formulas for the output dimensions are $h_2 = \left(\lfloor(h_1 - k_1 +2p)\rfloor/s+1\right)$ and $w_2 = \left(\lfloor(w_1 - k_2 +2p)\rfloor/s+1\right)$. For a given convolutional layer $C^{ij}$ with weights $Z^{ij}$, we define a matrix $V^{ij} \in \R^{c_ih_iw_i \times c_jh_jw_j}$ that computes $\forall~ x \in \R^{c_j \times h_j \times w_j}:~ V^{ij} \vect(x) ~=~ \vect(C^{ij}(x))$.
This matrix exists since both the padding and convolution operations can be represented as linear operations of the input. We also consider $W^{ij} \in \R^{c_{i} \times c_{j}k_1k_2}$ which is a matrix whose $c$'th row is the vectorized filter $W^{ij}_c := \vect(Z^{ij}_{c,~:~,~:~})$.
{\bf Fully-connected layers.\enspace} As a special case of convolutional layers, the network may also include fully-connected layers. Namely, a fully-connected layer $F:\R^{c_1} \to \R^{c_2}$ associated with a matrix $W \in \R^{c_2 \times c_1}$ can be represented as a $1\times 1$ convolutional layer $C:\R^{c_1 \times 1 \times 1}\to \R^{c_2 \times 1 \times 1}$ with $k_1=k_2=1$, $p=0$ and $s=1$. Namely, the parameters tensor $Z \in \R^{c_2 \times c_1 \times 1 \times 1}$ satisfies $Z_{a,b,1,1}=W_{a,b}$ for all $(a,b)\in [c_2]\times [c_1]$ and the layer satisfies $\vect(C(x)) = W \vect(x)$.
{\bf Pooling layers.\enspace} A pooling layer~\citep{Zhou1988ComputationOO} (see also~\citep{Goodfellow-et-al-2016}) $C$ with kernel dimensions $(k_1,k_2)$ stride $s$ and padding $p$ takes an input $x \in \mathbb{R}^{c_1 \times h_1 \times w_1}$ and computes an output $y \in \mathbb{R}^{c_2 \times h_2 \times w_2}$ with $c_2=c_1$ channels, and dimensions $h_2 = \left(\lfloor(h_1 - k_1 +2p)\rfloor/s+1\right)$ and $w_2 = \left(\lfloor(w_1 - k_2 +2p)\rfloor/s+1\right)$. Each pooling layer computes the following tensor as output
\begin{equation}\label{eq:pool}
\forall~ (c,t,l) \in [c_i]\times [h_i] \times [w_i]:~y_{c,t,l} ~=~ \textnormal{op}(\textnormal{Pad}_p(x)_{c,ts~:~(t+1)s+k_1,ls~:~(l+1)s+k_2}),
\end{equation}
where $\textnormal{op}$ is either the maximum or average operator.
{\bf Rearrangement layers.\enspace} To conveniently switch between convolutional and fully-connected layers, we should be able to represent tensor layers as vectors and vice versa. To reshape the representation of a certain layer, we allow the networks to include rearrangement layers. A rearrangement layer $C^{ij}:\R^{c_j \times h_j \times w_j} \to \R^{c_i \times h_i \times w_i}$ takes an input vector $x \in \R^{c_j \times h_j \times w_j}$ and `rearranges' its coordinates in a different shape and permutation. Formally, it returns a vector $(x_{\pi(k)})_{k \in [c_j]\times [h_j] \times [w_j]}$, where $\pi:[c_j]\times [h_j] \times [w_j] \to [c_i]\times [h_i] \times [w_i]$ is invertible (in particular, $c_ih_iw_i=c_jh_jw_j$). %{\color{red}relationship with fully-connected layers.}
\section{Theoretical Results}\label{sec:theory}
In this section we describe our main theoretical results. We investigate the inductive biases that emerge near convergence when training with SGD. For this purpose, we begin by introducing our definitions of SGD convergence (and near convergence) points.
To formally study convergence, we employ the notion of convergence in mean of random variables. Namely, we say that a sequence of random variables $\{W_{t}\}^{\infty}_{t=1} \subset \R^{N}$ starting from a certain vector $W_0 \in \R^{N}$ (constant) converges in mean if the following holds
\begin{equation*}
\exists~ W^* \in \R^{N}:~\lim_{t \to \infty} \E[\|W_{t} - W^*\|] ~=~ 0.
\end{equation*}
As a consequence of this definition, we have
\begin{equation*}
\lim_{t\to \infty}\E[\|W_{t+1} - W_t\|] ~\leq~ \lim_{t\to \infty}\E[\|W_{t+1} - W^*\|] + \lim_{t\to \infty}\E[\|W_{t} - W^*\|] ~=~ 0,
\end{equation*}
where the expectations are taken over the selections of the mini-batches. In particular, convergence is possible only when the expected size of each step tends to zero. In particular, when training the network using mini-batch SGD along with weight decay, we have
\begin{equation}\label{eq:convergence}
\lim_{t \to \infty} \E[\| \nabla L^{\lambda}_{\tiS}(f_{W_t})\|] ~=~ \lim_{t\to \infty} \fr{1}{\mu}\E[\|W_{t+1} - W_{t}\|] ~=~ 0.
\end{equation}
%where $\mu$ is the learning rate.
In this work we study the implicit biases of SGD by investigating the {\em ``SGD Near Convergence Regime''}, where SGD arrives at a point in training where each subsequent step is small compared to the actual weights, i.e., $\E_{\tiS} [\| \nabla_{W^{ij}} L^{\lambda}_{\tiS}(f_{W_{T}})\|/\|W^{ij}_T\|]$ is small. In Sec.~\ref{sec:lowrank} we show that near-convergence, mini-batch SGD learns neural networks with low-rank matrices and in Sec.~\ref{sec:convergence} we prove that perfect SGD convergence is impossible in the presence of weight decay.
We acknowledge that our definition differs from the traditional notions of convergence. Several papers study the convergence of SGD to a point where the performance is near-optimal when the learning rate $\mu_t$ decays with the number of iterations. In these cases, convergence is guaranteed since $\mu_t$ tends to zero. These papers, however, do not analyze whether the expected gradients $\E_{\tiS}[\|\nabla_{W^{ij}_t} L^{\lambda}_{\tiS}(f_{W_t})\|]$ also decays at each step. Other papers (e.g.,~\citep{https://doi.org/10.48550/arxiv.1605.08361,doi:10.1137/19M1308943}) study the critical points of the objective function to better understand the solutions of gradient-based optimization methods. While understanding the critical points of the objective function is necessary for characterizing the convergence points of GD and GF, SGD does not necessarily converge at these points. For instance, suppose $W_t$ is a stationary point of $L^{\lambda}_{S}(f_W)$ and there exists a batch $\tiS$ for which $\nabla_W L^{\lambda}_{\tiS}(f_{W_t})\neq 0$. With probability $\binom{m}{B}^{-1}>0$ SGD selects the batch $\tiS$ and updates $W_{t+1} = W_t - \mu \nabla L^{\lambda}_{\tiS}(f_{W_t}) \neq W_t$. Therefore, $p=\mathbb{P}[W_{t+1} \neq W_t \mid W_t] > 0$ and $\mathbb{P}[\exists~ l \in [T']:~W_{t+l} \neq W_t \mid W_t]\geq 1-(1-p)^{T'} \xrightarrow[T' \to \infty]{} 1$. As a result, the probability of the optimization becoming stuck at $W_t$ indefinitely is zero.
%Other papers (e.g.,~\cite{}) consider {\color{red}complete}, however, in practice we typically return $W_{T}$ as the output of SGD. However, the type of questions that we ask in this paper are different: {\color{red}complete}.
\subsection{Low-Rank Bias in Neural Networks}\label{sec:lowrank}
We begin our theoretical analysis with the simple observation (proved in Appendix~\ref{app:proofs}) that the number of input patches $N^{ij}$ of a certain convolutional layer $C^{ij}$ upper bounds the rank of the gradient of the network w.r.t. $W^{ij}$.
\begin{restatable}{lemma}{ranklemma}\label{lem:rank}
Let $f_W$ be a neural network and let $C^{ij}$ be a convolutional layer within $f_W$ with parameters matrix $W^{ij}$. Then, $\rank\left(\nabla_{W^{ij}} f_W(x)\right)\leq N^{ij}$.
\end{restatable}
Interestingly, we obtain particularly degenerate gradients for fully-connected layers. As discussed in Sec.~\ref{sec:archs} for a fully-connected layer $C^{ij}:\R^{c_j \times 1 \times 1} \to \R^{c_i \times 1 \times 1}$ we have $N^{ij}=1$, and therefore, $\rank\left(\nabla_{W^{ij}} f_W(x)\right) \leq 1$.
The following theorem provides an upper bound on the minimal distance between the network's weight matrices and low-rank matrices.
\begin{restatable}{theorem}{batches}\label{thm:batches}
Let $\|\cdot \|$ be any matrix norm and $\ell$ any differentiable loss function. Let $f_W(x)$ be a ReLU neural network and $C^{ij}$ be a convolutional layer within $f_W$ and let $B \in [m]$. Then,
\begin{equation*}
\min_{\substack{W \in R^{d_i \times d_j}:\\ \rank(W) \leq N^{ij}B}} \|\fr{W^{ij}}{\|W^{ij}\|} - W\| ~\leq~ \fr{1}{2\lambda}\min_{\substack{\tiS \subset S: \\ \vert \tiS\vert = B}} \|\nabla_{W^{ij}} L^{\lambda}_{\tiS}(f_W)\| / \|W^{ij}\|
\end{equation*}
\end{restatable}
\begin{proof}
Let $\tiS \subset S$ be a batch of size $B$. By the chain rule, we can write the gradient of the loss function as follows
%\begin{small}
\begin{equation*}
\begin{aligned}
\nabla_{W^{ij}} L^{\lambda}_{\tiS}(f_W) ~=~ & \fr{1}{B}\sum_{(x,y) \in \tiS}\fr{\partial \ell(f_W(x),y)}{\partial f_W(x)} \cdot \nabla_{W^{ij}} f_W(x) + 2\lambda W^{ij} =: - E_{\tiS} + 2\lambda W^{ij}.
\end{aligned}
\end{equation*}
%\end{small}
According to Lem.~\ref{lem:rank}, we have $\rank(\fr{1}{2\lambda}E_{\tiS}) \leq BN^{ij}$. Therefore, we obtain that
\begin{small}
\begin{equation*}
\begin{aligned}
\min_{W:~ \rank(W) \leq BN^{ij}}\|W^{ij} - W\| ~\leq~ \min_{\tiS \subset S:~\vert \tiS\vert = B}\|W^{ij} - \fr{1}{2\lambda}E_{\tiS}\|
~=~ \fr{1}{2\lambda}\min_{\tiS \subset S:~\vert \tiS\vert = B}\left\|\nabla_{W^{ij}} L^{\lambda}_{\tiS}(f_W) \right\|.
\end{aligned}
\end{equation*}
\end{small}
Finally, by dividing both sides by $\|W^{ij}\|$ we obtain the desired inequality.
\end{proof}
The theorem above provides an upper bound on the minimal distance between the parameters matrix $W^{ij}$ and a matrix of rank $\leq BN^{ij}$. The upper bound is proportional to, $\min_{\substack{\tiS}}\|\nabla_{W^{ij}} L^{\lambda}_{\tiS}(f_W)\|/\|W^{ij}\|$, which is the minimal norm of the gradient of the regularized empirical risk evaluated on batches of size $B$, normalized by the norm of the weight matrix. As shown in \eqref{eq:convergence}, near convergence we expect $\E_{\tiS}[\|\nabla_{W^{ij}} L^{\lambda}_{\tiS}(f_W)\| /\|W^{ij}\|]$ to be small, and therefore, $\min_{\substack{\tiS}} \|\nabla_{W^{ij}} L^{\lambda}_{\tiS}(f_W)\|/\|W^{ij}\|$ should also be small ($\tiS$ is distributed uniformly as a batch of samples size $B$). In particular, by the theorem above, we expect $\min_{W:~ \rank(W) \leq BN^{ij}}\|\fr{W^{ij}}{\|W^{ij}\|} - W\|$ to also be small. As a result, we predict that the rank of the learned parameter matrices $W^{ij}$ decreases with the batch size. In Sec.~\ref{sec:experiments} we validate this idea with an extensive set of experiments -- and also study the relationship between the rank and other hyperparameters.
% Therefore, it can be used to draw predictions about the relationships between certain hyperparameters, the rank and the performance. Several interesting predictions follow the theorem above.
\subsection{Degeneracy and the Origin of ``SGD Noise''}\label{sec:convergence}
As we mentioned, it is impossible for SGD to converge to a stationary point of the gradient dynamical system, resulting in inherent ``SGD noise''. In this section, we study the (non-)convergence of mini-batch SGD. Our results are essentially impossibility results: the assumption of SGD convergence to a critical point of the gradient implies that the network represents the zero function. As a result, asymptotic noise is inherently unavoidable (when training properly). For simplicity, we assume that $\forall~i\in [m]:~x_i\neq 0$.
As shown in \eqref{eq:convergence}, any convergence point $W$ of SGD satisfies $\E_{\tiS} [\| \nabla L^{\lambda}_{\tiS}(f_{W})\|]=0$. Since the distribution over mini-batches $\tiS$ of size $B$ is discrete and $\| \nabla L^{\lambda}_{\tiS}(f_{W})\|\geq 0$, we obtain that at convergence we have $\| \nabla_{W^{ij}} L^{\lambda}_{\tiS}(f_{W})\| = 0$, for all mini-batches $\tiS$ of size $B$. In particular,
\begin{equation}\label{eq:1}
\begin{aligned}
\forall \tiS:~ 0 ~&=~ \nabla_{W^{ij}} L^{\lambda}_{\tiS}(f_W) ~=~ \frac{1}{B} \sum_{(x,y)\in \tiS} \fr{\partial \ell(f_W(x),y)}{\partial f_W(x)} \cdot \nabla_{W^{ij}} f_W(x) + 2\lambda W^{ij}.
\end{aligned}
\end{equation}
Suppose we have two batches $\tiS_1,\tiS_2 \subset S$ of size $B$ that differ by only one sample. We denote the unique sample of each batch by $(x_{j_1},y_{j_1})$ and $(x_{j_2},y_{j_2})$ respectively. We notice that,
\begin{equation*}
\begin{aligned}
0 ~&=~ \nabla_{W^{ij}} L^{\lambda}_{\tiS_1}(f_W) - \nabla_{W^{ij}} L^{\lambda}_{\tiS_2}(f_W) \\
~&=~ \fr{\partial \ell(f_W(x_{j_1}),y_{j_1})}{\partial f_W(x_{j_1})} \cdot \nabla_{W^{ij}} f_W(x_{j_1}) - \fr{\partial \ell(f_W(x_{j_2}),y_{j_2})}{\partial f_W(x_{j_2})} \cdot \nabla_{W^{ij}} f_W(x_{j_2}).
\end{aligned}
\end{equation*}
Therefore, we conclude that for all $j_1,j_2\in [m]$,
\begin{equation}\label{eq:j1j2}
\begin{aligned}
M^{ij} ~=~ \fr{\partial \ell(f_W(x_{j_1}),y_{j_1})}{\partial f_W(x_{j_1})} \cdot \nabla_{W^{ij}} f_W(x_{j_1}) ~=~ \fr{\partial \ell(f_W(x_{j_2}),y_{j_2})}{\partial f_W(x_{j_2})} \cdot \nabla_{W^{ij}} f_W(x_{j_2}).
\end{aligned}
\end{equation}
Hence, for all $(v_i,v_j) \in E_T$ and $k \in [m]$,
\begin{equation}\label{eq:batchone}
\fr{\partial \ell(f_W(x_k),y_k)}{\partial f_W(x_k)} \cdot \nabla_{W^{ij}} f(x_k) + 2\lambda W^{ij} ~=~ M^{ij} + 2\lambda W^{ij} ~=~ 0.
\end{equation}
Therefore, unless $\lambda=0$ or $\forall~(v_i,v_j)\in E_T:~W^{ij}=0$, we conclude that $\fr{\partial \ell(f_W(x_k),y_k)}{\partial f_W(x_k)}\neq 0$ for all $k\in [m]$. In this case, we also obtain that $\{\fr{\partial f_W(x_k)}{\partial \textnormal{vec}(W^{ij})}\}^{m}_{k=1}$ are collinear vectors by \eqref{eq:j1j2}.
Therefore, any convergence point of training a neural network using mini-batch SGD along with weight decay is highly degenerate and does not fit any one of the training labels. To better understand the essence of this degeneracy, we provide the following proposition (proved in Appendix~\ref{app:proofs}), which is specialized for ReLU networks.
\begin{restatable}[$\lambda > 0$]{proposition}{collinear}\label{prop:propA}
Let $\ell(a,b)$ be a differentiable loss function, $\lambda > 0$, and let $f_W(x)$ be a ReLU neural network, where $\textnormal{succ}(1)=\{p\}$ and $(v_p,v_1)\in E_T$. Let $\{x^k_{i}\}^{N^{p1}}_{k=1}$ be the $N^{p1}$ patches of $x_i$ used by the layer $C^{p1}$. Let $W$ be a convergence point of mini-batch SGD for minimizing $L^{\lambda}_S(f_W)$ (see \eqref{eq:1}). Then, either $f_W \equiv 0$ or $\forall~ i,j\in [m]:\{x^k_i,x^k_j\}^{N^{p1}}_{k=1}$ are linearly dependent tensors.
\end{restatable}
The preceding proposition shows that unless the patches of any two training samples are linearly dependent, any convergence point of SGD corresponds to the zero function. When $N^{p1}$ is small, the linear dependence criterion is unrealistic. For example, if $C^{p1}$ is a fully-connected layer, $N^{p1}=1$, and the condition asserts that any two training samples $x_i,x_j$ are collinear. Since this is unrealistic, we conclude that convergence is impossible unless $f_W \equiv 0$.
As a next step, we consider convergence of SGD when training without weight decay.
\begin{restatable}[$\lambda = 0$]{proposition}{zerolambda}\label{prop:propB}
Let $\lambda=0$ and let $\ell$ be a differentiable loss function. Let $f_W(x)$ be a ReLU neural network, where $\textnormal{succ}(1)=\{p\}$ and $(v_p,v_1)\in E_T$ is fully-connected. Let $\{x^k_{i}\}^{N^{p1}}_{k=1}$ be the $N^{p1}$ patches of $x_i$ used by the layer $C^{p1}$. Let $W$ be a convergence point of mini-batch SGD for minimizing $L^{\lambda}_S(f_W)$. Then, $\forall~ i\in [m]:\fr{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})}=0$ or $f_W(x_j)=0$ or $\forall~ j\in [m]:\{x^k_i,x^k_j\}^{N^{p1}}_{k=1}$ are linearly dependent tensors. %In particular, if $\ell$ is convex and $\forall~b\in \R~\nexists~a^* \in \R:~ \ell(a^*,b)=\inf_a \ell(a,b)$, then, for all $i \in [m]$, $f_W(x_i)=0$ or $\{x^k_i,x^k_j\}^{N^{p1}}_{k=1}$ are linearly dependent tensors.
\end{restatable}
The preceding proposition (proved in Appendix~\ref{app:proofs}) provides conditions for SGD convergence of training a convolutional network without weight decay. It shows that at convergence, for every sample the network either perfectly fits the label (i.e., $\fr{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} = 0$) or outputs zero unless the patches of that sample are linearly dependent with the patches of any other sample. As mentioned, the linear dependence criteria is generally unrealistic when $(v_p,v_1)$ is fully-connected. Therefore, convergence with fully-connected networks is possible only when $f_W$ perfectly fits the training set labels.
We note that if $\ell(a,b)$ is convex and has no minima $a$ for any $b \in \R$ (e.g., binary cross-entropy, logistic loss or exponential loss), then $\forall~ i \in [m]:\fr{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \neq 0$. Therefore, the only possible convergence points of fully-connected networks are ones for which $\forall~i \in [m]:f_W(x_i)=0$. Since this is in general absurd, we argue that perfect convergence of training a network with exponential-type loss functions is generally impossible. While convergence to a non-zero function is not guaranteed, in practice training without weight decay may still fall into the regime of `almost convergence', in which $\max_{i \in [m]} \big\vert \fr{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \big\vert$ is tiny and as a result the training steps $- \mu \cdot \fr{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \cdot \fr{\partial f_W(x_i)}{\partial W^{ij}}$ are very small. This is usually the case when an overparameterized network has been properly trained. For the squared loss, convergence may occur when the network perfectly fits the training labels, i.e., $\forall~i\in [m]:~f_W(x_i)=y_i$.
It is worth noting that training with mini-batches during optimization and weight decay are critical to our analysis. While many papers examine the training dynamics and critical points of GD, as we show, SGD convergence points are highly degenerate and, in general, behave differently than GD solutions. %For example, it is impossible to derive \eqref{eq:j1j2} without using mini-batches. Without weight decay, we cannot claim that $W^{ij}$ is proportional to the low-rank matrix $M^{ij}$.
Surprisingly, this analysis is unaffected by batch size, and therefore, the presence of SGD noise occurs regardless of batch size, as long as it is strictly smaller than the full dataset's size.
\section{Experiments}\label{sec:experiments}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_03_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_05_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_3_batch_size.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_03_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_05_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_3_batch_size.pdf}
\\
{\small $\mu=0.03$} & {\small $\mu=0.05$} & {\small $\mu=0.1$} & {\small $\mu=0.3$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of MLP-BN-10-100 trained on CIFAR10 with various batch sizes.} The top row shows the average rank across layers, while the bottom row shows the train and test accuracy rates for each setting. Weight decay $\lambda=\expnumber{5}{-4}$ was used to train each model. To calculate the rank, we used an $\epsilon=0.001$ threshold.}\label{fig:mlp_bn_10_100_lr/bs}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/rank_resnet_bn_bs_4_lr_15_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/rank_resnet_bn_bs_8_lr_15_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/rank_resnet_bn_bs_16_lr_15_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/rank_resnet_bn_bs_32_lr_15_weight_decay.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/acc_resnet_bn_bs_4_lr_15_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/acc_resnet_bn_bs_8_lr_15_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/acc_resnet_bn_bs_16_lr_15_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/acc_resnet_bn_bs_32_lr_15_weight_decay.pdf}
\\
{\small $B=4$} & {\small $B=8$} & {\small $B=16$} & {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of ResNet-18 trained on CIFAR10 with varying weight decay.} In this experiment: $\mu=1.5$ and $\epsilon=0.001$.}
\label{fig:wd_resnet18}
\end{figure*}
\begin{figure*}[t] % vit batch size
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/vit/batch_size/varyB_lr4e-3_avgrank.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/batch_size/varyB_lr8e-3_avgrank.pdf} &
%\includegraphics[width=0.20\linewidth]{plots/vit/batch_size/varyB_lr1e-2_avgrank.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/batch_size/varyB_lr2e-2_avgrank.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/batch_size/varyB_lr4e-2_avgrank.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/vit/batch_size/varyB_lr4e-3_accuracy.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/batch_size/varyB_lr8e-3_accuracy.pdf} &
%\includegraphics[width=0.20\linewidth]{plots/vit/batch_size/varyB_lr1e-2_accuracy.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/batch_size/varyB_lr2e-2_accuracy.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/batch_size/varyB_lr4e-2_accuracy.pdf}
\\
{\small $\mu=0.004$} & {\small $\mu=0.008$} %& {\small $\mu=0.01$}
& {\small $\mu=0.02$} & {\small $\mu=0.04$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of ViT trained on CIFAR10 with various batch sizes.} In this experiment: $\lambda=\expnumber{5}{-4}$ and $\epsilon=0.01$.}\label{fig:bs_vit}
%(no-momentum, patch size=4, depth=6, heads=8, dropout=0.1, $mlp_dim$=512, weight decay=5e-4)
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c@{}c}
\rotatebox[origin=c]{90}{{\tiny CIFAR10, CE}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/cifar10_dist_ce_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/cifar10_acc_ce_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/cifar10_loss_ce_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/cifar10_rank_ce_mlp_2000_lr_01.pdf}}
\\
\rotatebox[origin=c]{90}{{\tiny CIFAR10, MSE}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/cifar10_dist_mse_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/cifar10_acc_mse_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/cifar10_loss_mse_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/cifar10_rank_mse_mlp_2000_lr_01.pdf}}
\\
& {\small (a)} & {\small (b)} & {\small (c)} & {\small (d)}
\end{tabular}
\caption{{\bf Convergence of MLP-5-2000 trained on CIFAR10.} In {\bf (a)} we plot the averaged distance between the weight matrices at epoch $t$ and epoch $t+1$, captured by, $d(W_{t+1},W_{t})$. In {\bf (b)} we plot the train accuracy rates, in {\bf (c)} we plot the averaged train loss and in {\bf (d)} we plot the average rank across the trainable matrices.}
\label{fig:convergence}
\end{figure*}
In this section we empirically study the implicit bias towards rank minimization in deep ReLU networks. Throughout the experiments we extensively vary different hyperparameters (e.g., the learning rate, weight decay, and the batch size) and study their effect on the rank of the various matrices in the network. In addition, we experimentally analyze the convergence (and non-convergence) of training neural networks with SGD\footnote{The plots are best viewed when zooming into the pictures.}. %{\color{blue} Each of the runs were trained using a single GPU on a computing cluster with several available GPU types. The length of the runs did not exceed 60 hours, and the available GPUs were NVIDIA's GeForce RTX 2080, GeForce RTX 2080 Ti, Quadro RTX 6000, Tesla V-100, GeForce RTX A6000, A100, and GeForce GTX 1080 Ti.}
\subsection{Setup}
\paragraph{Evaluation process.} We consider $k$-class classification problems and train a multilayered neural network $f_W :\R^n \to \R^k$ on some balanced training dataset $S$. The model is trained using CE/~MSE loss minimization between its logits and the one-hot encodings of the labels. After each epoch, we compute the averaged rank across the network's weight matrices and its train and test accuracy rates. For a convolutional layer $C^{ij}$, we use $W^{ij}$ as its weight matrix. To estimate the rank of a given matrix $M$, we count how many of the singular values of $M/\|M\|_2$ are $\notin [-\epsilon, \epsilon]$, where $\epsilon$ is a small tolerance value. In these experiments we consider the MNIST and CIFAR10 datasets. %During training, we used random cropping, random horizontal flips and random rotations (by $15k$ degrees for $k$ uniformly sampled from $[24]$) and standardize the data. %For all datasets we standardized the data.
\paragraph{Architectures.} We consider several network architectures. %The first architecture is a convolutional network (ConvNet), beginning with a stack of a $2\times 2$ convolutional layer with stride $2$, batch normalization, a convolution of the same structure, batch normalization and ReLU. Following that we have a set of $L$ stacks of $3\times 3$ convolutional layer with stride $1$ and padding $1$, batch normalization and ReLU. The output tensors of these layers shares the same shape as their input's shape. The number of channels in each convolutional layer is $H$. The last layer is linear. The network is denoted by Conv-$L$-$H$.
{\bf (i)} The first architecture is an MLP, denoted by MLP-BN-$L$-$H$, which consists of $L$ hidden layers, where each layer contains a fully-connected layer of width $H$, followed by batch normalization and ReLU activations. On top of that, we compose a fully-connected output layer. {\bf (ii)} The second architecture, denoted by RES-BN-$L$-$H$, consists of a linear layer of width $H$, followed by $L$ residual blocks, ending with a fully-connected layer. Each block computes a function of the form $z+\sigma(n_2(W_2 \sigma(n_1(W_1z))))$, where $W_1,W_2 \in \R^{H\times H}$, $n_1,n_2$ are batch normalization layers and $\sigma$ is the ReLU function. We denote by MLP-$L$-$H$ and RES-$L$-$H$ the same architectures without applying batch normalization. {\bf(iii)} The third, denoted VGG-16, is the convolutional network proposed by \cite{DBLP:journals/corr/SimonyanZ14a}, but with dropout replaced by batch normalization layers and with only one fully-connected layer at the end. {\bf(iv)} The fourth architecture is the residual network proposed in \citep{7780459}, denoted ResNet-18. {\bf (v)} The fifth layer is a small visual transformer \citep{dosovitskiy2020image}, denoted by ViT. Our ViT splits the input images into patches of size $4\times 4$ includes $8$ self-attention heads, where each one of them consists of $6$ self-attention layers. The self-attention layers are followed by two fully-connected layers with dropout probability 0.1 and a GELU activation in between them.
%Unless mentioned otherwise,
% {\bf Datasets.\enspace} Throughout the experiments, we consider three datasets: (i) MNIST; (ii) CIFAR10; (iii) STL-10. For CIFAR10 we used random cropping, random horizontal flips and random rotations (by $15k$ degrees for $k$ uniformly sampled from $[24]$). For all datasets we standardized the data. {\color{red}Update}
\subsection{Experiments on Rank Minimization}\label{sec:rank_exp}
In each experiment we trained various models while varying one hyperparameter (e.g., batch size) and leaving the other hyperparameters constant. The models were trained with SGD for cross-entropy loss minimization along with weight decay. For MLP-BN-10-100, ResNet-18 and VGG-16, we decayed the learning rate three times by a factor of $0.1$ at epochs 60, 100, and 200 and training is stopped after 500 epochs. We train instances of ViT using SGD and the the learning rate is decayed by a factor of $0.2$ three times at epochs 60, 100 and training is stopped after 200 epochs. By default, we trained the models with weight decay $\lambda=\expnumber{5}{-4}$.
%{\bf Rank, batch size and the learning rate.\enspace}
As can be seen in Figs.~\ref{fig:mlp_bn_10_100_lr/bs} and~\ref{fig:bs_vit} by decreasing the batch size, we essentially strengthen the low-rank constraint over the network's matrices, which eventually leads to matrices of lower ranks. This is consistent with the prediction made in Sec.~\ref{sec:lowrank} that we learn matrices with lower ranks when training the network with smaller batch sizes. Interestingly, we also notice a regularizing effect for the learning rate; the average rank tends to decrease when increasing the learning rate. %{\color{red}EXPERIMENT WITH FIXED BS.}
%{\bf Rank minimization and weight decay.\enspace}
As can be seen in Fig.~\ref{fig:wd_resnet18}, by increasing $\lambda$ we typically impose stronger rank minimization constraints. Interestingly, it appears that the batch size have little effect on the ranks of the weight matrices when training with $\lambda=0$ which is exactly the case when our bound is infinite. This empirically validates that weight decay is a necessary to obtain a significant low-rank bias.
%To study the effect of the weight decay on rank minimization, we train instances of MLP-BN-10-100 using SGD with various batch sizes $B \in \{16,32,64,128\}$ and values of $\lambda \in \{0, \expnumber{1}{-4}, \expnumber{2}{-4}, \expnumber{4}{-4}, \expnumber{8}{-4}, \expnumber{6}{-3}\}$. The learning rate is initially $\mu=0.1$ and is decayed by a factor of $0.1$ three times at epochs 60, 100, and 200 and training is stopped after 500 epoch and the weight decay is $\lambda=\expnumber{5}{-4}$. {\color{blue} {\bf (ii)} We additionally train instances of VGG-16 and ResNet-18 with the previous hyperparamaters except with momentum $0$.}
%The results are reported in Fig. .
\subsection{Experiments on SGD Noise}\label{sec:sgd_noise_exp}
In Sec.~\ref{sec:convergence} we showed that convergence to a non-zero function is impossible when training a fully-connected neural network with SGD and weight decay. To validate this prediction, we trained MLP-5-2000 instances examined their convergence as a function of $\lambda$. Each model was trained for CIFAR10 classification using SGD with batch size 128 and learning rate 0.1 for 2000 epochs.
To investigate the convergence of the networks, we measure the averaged distance between the network's matrices between consecutive epochs: $d(W_{t+1},W_{t}) := \fr{1}{\vert E_T\vert} \sum_{(p,q)\in E_T} \|W^{ij}_{t+1} - W^{ij}_{t}\|$, where $\{W^{ij}_t\}_{(p,q)\in E_T}$ are the various trainable matrices in the network at epoch $t$. As mentioned in Sec.~\ref{sec:theory}, convergence is possible only when $\lim\limits_{t\to \infty} d(W_{t+1},W_{t}) = 0$.
In Fig.~\ref{fig:convergence} we monitor {\bf (a)} $d(W_{t+1},W_{t})$, {\bf (b)} the train accuracy rates, {\bf (c)} the train losses and {\bf (d)} the averaged rank of the trainable matrices. As predicted in Prop.~\ref{prop:propA}, when training with $\lambda>0$, $W_t$ either converges to zero and $f_{W_t}$ to the zero function (e.g., see the results with $\lambda=\expnumber{1}{-4}, \expnumber{1}{-3}$) or $W_t$ does not converge (i.e., $d(W_{t+1},W_{t})$ does not tend to zero). Furthermore, we observe that for $\lambda=0$, $d(W_{t+1},W_{t})$ is smaller by orders of magnitude compared to using $\lambda>0$ (except for cases when the selection of $\lambda > 0$ leads to $W^{ij}_t \to 0$). Interestingly, even though for certain values of $\lambda > 0$, the training loss and accuracy converged, the network's parameters do not converge.
Finally, when training for cross-entropy loss minimization without weight decay we encounter the `almost convergence' regime discussed in Sec.~\ref{sec:convergence}. Namely, even though perfect convergence is impossible, the term $d(W_{t+1},W_{t})$ may become as small as we wish by increasing the size of the neural network. Therefore, since the MLP-5-2000 is relatively large (compared to the dataset's size), we may inaccurately get the impression that $d(W_{t+1},W_{t})$ tends to zero.
% We repeated the same experiment with instances of MLP-10-100 without the batch normalization as well. In this case, we used $\lambda \in \{0, \expnumber{1}{-3}, \expnumber{2}{-3}, \expnumber{4}{-3}, \expnumber{8}{-3}, \expnumber{6}{-2}, \expnumber{2}{-1}\}$. For each run we trained the network for 500 epochs using SGD with momentum $0.9$, initial learning rate $0.01$ and batch size 64.
\section{Conclusions}\label{sec:conclusions}
A mathematical characterization of the biases associated with SGD-trained neural networks is regarded as a significant open problem in the theory of deep learning~\citep{NIPS2017_10ce03a1}. In addition to its independent interest, a low-rank bias -- though probably not necessary for generalization -- may be a key ingredient in an eventual characterization of the generalization properties of deep networks. In fact, recent results~\citep{huh2022the} and our preliminary experiments (see Figs.~\ref{fig:gen_1}-\ref{fig:gen_2} in the appendix) suggest that low-rank bias in neural networks improves generalization.
By investigating the ``SGD Near-Convergence Regime'', we proved that SGD together with weight decay induces a low-rank bias in a variety of network architectures. Our result also shows that the batch size used by SGD influences the rank of the learned matrices. This means that the batch size plays an active role in regularizing the learned function.% and should ideally be tuned in accordance with the dataset's complexity.
We also prove that when training a fully-connected neural network, SGD noise must always be present, even asymptotically, regardless of batch size, as long as weight decay is used. Weight decay may not be strictly necessary for SGD noise and low-rank bias to appear. This is the case, for example, when training with exponential-type loss functions and Weight Normalization~\citep{SalDied16}.
%SGD has several other interesting properties. ~\cite{Wu2018} showed that SGD selects flat vs sharp minima -- which is arguably a significant advantage in the training of overparametrized neural networks because global zero-loss minima are highly degenerate under these conditions~\cite{doi:10.1137/19M1308943}.
We hope that our work will spark further research into the near-convergence regime. For instance, it may provide an interesting algorithmic approach to the long-standing problem of developing low-rank regularizers during optimization. It would also be interesting to study whether additional structures (e.g., neural collapse~\citep{Papyan24652}, sparsity) emerge during the near-convergence regime and to extend our analysis to more sophisticated learning algorithms (e.g., Adam \citep{kingma:adam}) and learning settings (e.g., unsupervised learning, self-superviSsed learning, etc).
% We make several theoretical and empirical observations:
% \begin{itemize}[leftmargin=0.5cm]
% \item We show that training with mini-batch SGD along with weight decay induces a low-rank bias in the weight matrices of neural networks.
% \item We provide theoretical intuition for the effect of weight decay and the batch size on the rank of the learned weight matrices. Our experiments also show that weight decay and training with small mini-batches are necessary to obtain rank minimization.
% \item We theoretically show that SGD noise is inevitable when applying weight decay, regardless of the batch size and learning rate (unless the network converges to a trivial solution).
% \end{itemize}
%It has not escaped our attention that these results have a rich set of implications with respect to other mysteries of deep learning that are yet to be solved.
%They include questions such as why SGD training of networks generalizes well, the relationship between rank minimization and architectural choices (residual networks~\citep{He_2017_ICCV}, transformers~\citep{NIPS2017_3f5ee243} and convolutional networks~\citep{3a6e6cc3ac6e4326b35a6315d7cbddf7}) and whether rank minimization extends beyond the standard settings of supervised learning.
\newpage
\bibliography{refs}
\bibliographystyle{iclr2023_conference}
\appendix
\onecolumn
% \section{Additional Architectures}
% \paragraph{Optimization.} {\color{red}edit the detaching} In this work, we consider optimizing $f_W$ using mini-batch stochastic gradient descent (SGD), possibly with detaching certain layers. Namely, we aim at minimizing the regularized empirical risk $L^{\lambda}_{S}(f_{W})$ by applying SGD for a certain number of iterations. We initialize $W_0$ using a standard initialization procedure and iteratively update $W_t$. At each iteration, we sample a subset $\tiS = \{(x_{ij},y_{ij})\}^{B}_{j=1} \subset S$ uniformly at random and update $W_{t+1} \leftarrow W_t - \mu_t \cdot \bar{\nabla}_W L^{\lambda}_{\tiS}(f_{W_t}(~\cdot~;\tiX))$, where $\tiX = \{x_{ij}\}^{B}_{j=1}$ and $\mu_t > 0$ is our learning rate at iteration $t$. Here, $f_{W}(~\cdot~;\tiX) : \R^{n} \to \R$ is the function given by $f_W(x;\tiX)$ for a fixed value of $\tiX$. Furthermore, for a given function $g_{\theta} = g^L_{\theta} \circ \dots g^1_{\theta}:\R^{n} \to \R^k$ where $\{g^{i_j}_{\theta}\}^{r}_{j=1}$ are detached, $\frac{\bar{\partial} g_{\theta}(z)}{\partial \theta} = \bar{\nabla}_{\theta} g_{\theta}(z)$ denotes the {\color{red}pseudo gradient} of $g_{\theta}$, which is computed using the chain rule, while treating $\{g^{i_j}_{\theta}\}^{r}_{j=1}$ as independent of $\theta$.
% \paragraph{Network architecture.} Formally, $f_W$ is as a directed acyclic graph (DAG) $G=(V,E)$, where $V= \{v_1,\dots,v_L\}$ consists of the various layers of the network and each edge $e_{ij}=(v_i,v_j) \in E$ specifies a connection between two layers. Each layer is a function $v_{i}:\R^{c_{1} \times h_1 \times w_1} \to \R^{c_{i} \times h_i \times w_i}$ and each connection holds a linear transformation $C^{ij} : \R^{c_{j} \times h_j \times w_j} \to \R^{c_i \times h_i \times w_i}$. The layers are divided into three categories: the input layer $v_1$, the output layer $v_L$ and intermediate layers. In this setting, we do not have connections directed towards the input layer nor connections coming out of the output layer (i.e., $\forall~i \in [L]:~(v_i,v_L), (v_0,v_i) \notin E$). Given an input $x$ and batch $\tiX$, the output of a given layer $v_i$ is evaluated as follows
% \begin{equation*}
% v_i(x;\tiX) ~:=~ \sigma\left(n_i\left[\sum_{j \in \textnormal{pred}(i)} C^{ij}(v_{j}(x;\tiX));\tiX\right]\right),
% \end{equation*}
% except for the output layer $v_L$ that computes
% \begin{equation*}
% f_W(x;\tiX) ~:=~ v_L(x;\tiX) ~:=~ \sum_{j \in \textnormal{pred}(L)}C^{Lj} (v_{j}(x;\tiX)),
% \end{equation*}
% where $\textnormal{pred}(i)$ is the set of indices $j$, such that $(v_i,v_j) \in E$. Here, $\sigma$ is an element-wise activation function (e.g., ReLU, Leaky ReLU, sigmoid).
% In this setting, the transformations $C^{ij}$ could be trainable, $(v_p,v_q) \in E_T$ (e.g., convolutional layers), or constant, $(v_p,v_q) \notin E_T$ (e.g., residual connections or a pooling layer). Throughout the analysis, we consider paths within the graph $G$, denoted by $\pi = (\pi_0,\dots,\pi_T)$, where $\pi_0=1$, $\pi_T = L$ and for all $i=0,\dots,T-1:~ (v_{\pi_{i}},v_{\pi_{i+1}}) \in E$.
% \paragraph{Convolutional layers.} {\color{red}add a formal def maybe?} Throughout the paper we consider convolutional layers~\citep{6795724,10.5555/646469.691875} (see also~\citep{GoodBengCour16}) $C:\R^{c_1 \times h_1 \times w_1} \to \R^{c_2 \times h_2 \times w_2}$ with kernel sizes $(k_1,k_2)$, padding $p$ and stride $s$. The formula for the output dimensions are $h_2 = \left(\lfloor(h_1 - k_1 +2p)\rfloor/s+1\right)$ and $w_2 = \left(\lfloor(w_1 - k_2 +2p)\rfloor/s+1\right)$ and the number of patches on which the layer acts upon is $ h_2w_2$. For a given convolutional layer $C^{ij}$ in the network $f_W$, we denote $N^{ij}:=h_pw_p$. For a given convolutional layer $C^{ij}$ we define a matrix $W^{ij}$ that computes $\forall~ x \in \R^{c_q \times h_q \times w_q}:~ W^{ij} \vect(x) = \vect(C^{ij}(x))$. This is guaranteed since both the padding and convolution operations of the layer are linear. We also consider $V^{ij} \in \R^{c_{p} \times c_{q} k_1k_2}$ which is a matrix whose rows are $\vect(V^{ij}_{c,:,:})$.
% \paragraph{Pooling Layers.} A pooling layer (see~\citep{Goodfellow-et-al-2016}) $C$ with kernel dimensions $(k_1,k_2)$ stride $s$ and padding $p$ takes an input $c_{in}\times h_1 \times w_1$ and computes an output with dimensions $h_2 = \left(\lfloor(h_1 - k_1 +2p)\rfloor/s+1\right)$ and $w_2 = \left(\lfloor(w_1 - k_2 +2p)\rfloor/s+1\right)$. In this paper we consider average pooling and max-pooling layers.
% \paragraph{Batch normalization layers.} A batch normalization layer~\citep{pmlr-v37-ioffe15} is defined in the following manner,
% \begin{small}
% \begin{equation*}
% \begin{aligned}
% n_l(v_l(x;\tiX);\tiX) ~&=~ \gamma_l \cdot A_l(\tiX) \odot (v_l(x;\tiX) - b_l(\tiX)) \\
% &\quad + \one_{c\times h_l \times w_l}\cdot \beta_l \\
% A_l(\tiX) ~&=~ \left(\left(\sqrt{\Var_{\hat{x} \sim U[\tiX]}(v_l(\hat{x};\tiX)_{j})}\right)^{-1}\right)_{j} \\
% b_l(\tiX) ~&=~ \E_{\hat{x} \sim U[\tiX]}[v_l(\hat{x};\tiX)],
% \end{aligned}
% \end{equation*}
% \end{small}
% where $v_l(\hat{x};\tiX)_{j}$ is the $j$'th coordinate of $v_l(\hat{x};\tiX)$ for $j \in [c_l]\times [h_l]\times [w_l]$. In addition, $\gamma_l,\beta_l \in \R$ are trainable scale and shift parameters and $\one_{a_1 \times \dots \times a_k} \in \R^{a_1\times \dots \times a_k}$ is a tensor whose all of its cells are equal $1$. To simplify the analysis, we detach (also known as stop-grad~\citep{NEURIPS2020_f3ada80d}) the matrices $A_l(\tiX)$ during backpropagation at any step during training. Namely, we treat $\frac{\bar{\partial} A_l(\tiX)}{\partial W}=0$ in order to compute the training updates (without changing the forward passes). We note that detaching is a commonly used approach for training a network with normalization layers~\citep{Wiesler2014MeannormalizedSG,renorm,pmlr-v22-raiko12,https://doi.org/10.48550/arxiv.1410.7455,NEURIPS2019_2f4fe03d,10.1007/s10994-020-05929-w,https://doi.org/10.48550/arxiv.2202.07304}. In fact,~\cite{NEURIPS2019_2f4fe03d} showed that detaching the variance term of layer normalization is competitive and even improves the performance on various tasks.
\newpage
\section{Proofs}\label{app:proofs}
\ranklemma*
\begin{proof} Let $x \in \R^{c_1\times h_1\times w_1}$ be an input tensor and $C^{ij}$ be a certain convolutional layer with kernel size $(k_1,k_2)$, stride $s$ and padding $p$. We would like to show that $\rank\left(\nabla_{W^{ij}} f_W(x)\right) \leq N^{ij}$. We begin by writing the output of $f_W$ as a sum over paths that pass through $C^{ij}$ and paths that do not. We note that the output can be written as follows,
\begin{equation*}
f_W(x) ~=~ \sum_{l_1 \in \textnormal{pred}(l_0)} C^{l_0l_1} \circ v_{l_1}(x),
\end{equation*}
where $l_0 = L$ and $C^{l_0l_1} \circ z := C^{l_0l_1}(z)$. In addition, each layer $v_{l}$ can also be written as
\begin{equation*}
v_{l_1}(x) ~=~ D_{l_1} \odot \sum_{l_2\in \textnormal{pred}(l_1)} C^{l_1l_2} \circ v_{l_2}(x),
\end{equation*}
where $D_l := D_l(x) := \sigma'(v_l(x))) \in \R^{c_l \times h_l \times w_l}$.
A path $\pi$ within the network's graph $G$ is a sequence $\pi = (\pi_0,\dots,\pi_T)$, where $\pi_0=1$, $\pi_T = L$ and for all $i=0,\dots,T-1:~ (v_{\pi_{i}},v_{\pi_{i+1}}) \in E$. We can write $f_W(x)$ as the sum of matrix multiplications along paths $\pi$ from $v_1$ to $v_{l_0}$. Specifically, we can write $f_W(x)$ as a follows
%\begin{small}
\begin{equation*}
\begin{aligned}
f_W(x) ~&=~ \sum_{\pi \textnormal{ from $i$ to $l_0$}} C^{\pi_{T}\pi_{T-1}} \circ D_{\pi_{T-1}} \cdots D_{\pi_2} \odot C^{\pi_2\pi_1} \circ D_{\pi_1} \odot C^{ij} \circ v_j(x) \\
&\quad\quad + \sum_{\substack{\pi \textnormal{ from $1$ to $l_0$} \\ (i,j) \notin \pi }} C^{\pi_{T}\pi_{T-1}} \circ D_{\pi_{T-1}} \odot C^{\pi_{T-1}\pi_{T-2}} \cdots D_{\pi_2} \odot C^{\pi_2 \pi_1} \circ x, \\
~&=:~ A_W(x)+B_W(x)
\end{aligned}
\end{equation*}
%\end{small}
where $T=T(\pi)$ denotes the length of the path $\pi$. Since $\sigma$ is a piece-wise linear function with a finite number of pieces, for any $x\in \R^{c_1\times h_1\times w_1}$, with measure $1$ over $W$, the matrices $\{D_{l}(x)\}^{L-1}_{l=1}$ are constant in the neighborhood of $W$. Furthermore, $W^{ij}$ does not appear in the multiplications along the paths $\pi$ from $1$ to $l_0$ that exclude $(i,j)$. Therefore, we conclude that $\frac{\partial B_W(x)}{\partial W^{ij}}=0$.
As a next step we would like to analyze the rank of $\frac{\partial A_W(x)}{\partial W^{ij}}$. For this purpose, we rewrite the convolutional layers and the multiplications by the matrices $D_{l}(x)$ as matrix multiplications.
{\bf Representing $C^{ij}$.\enspace} We begin by representing the layer $C^{ij}$ as a linear transformation of its input with $N^{ij}$ blocks of $W^{ij}$. For this purpose, we define a representation of a given $3$-dimensional tensor input $z \in \R^{c_j \times h_j \times w_j}$ as a vector $\vect^{ij}(z) \in \R^{N^{ij} c_{j} k_1 k_2}$. First, we pad $z$ with $p$ rows and columns of zeros and obtain $\textnormal{Pad}_p(z)$. We then vectorize each one of its patches (of dimensions $c_{j}\times k_1\times k_2$) that the convolutional layer is acting upon (potentially overlapping) and concatenate them. We can write the vectorized output of the convolutional layer as $U^{ij} \vect^{ij}(z)$, where
\begin{equation}\label{eq:Uij}
U^{ij} := \renewcommand{\arraystretch}{1.2}
\left(
\begin{array}{ c c c c c c }
\multicolumn{1}{c}{} & & 0 & 0 & 0 & 0 \\
\multicolumn{2}{c}{\raisebox{.6\normalbaselineskip}[0pt][0pt]{$W^{ij}$}} & 0 & 0 & 0 & 0 \\
0 & 0 & & & 0 & 0 \\
0 & 0 & \multicolumn{2}{c}{\raisebox{.6\normalbaselineskip}[0pt][0pt]{$\ddots$}} & 0 & 0 \\
0 & 0 & 0 & 0 & & \multicolumn{1}{c}{} \\
0 & 0 & 0 & 0 & \multicolumn{2}{c}{\raisebox{.6\normalbaselineskip}[0pt][0pt]{$W^{ij}$}}
\end{array}
\right)
\end{equation}
is a $(N^{ij} c_{i}) \times (N^{ij} c_{j} k_1k_2)$ matrix with $N^{ij}$ copies of $W^{ij}$. We note that this is a non-standard representation of the convolutional layer's operation as a linear transformation. Typically, we write the convolutional layer as a linear transformation $V^{ij}$ acting on the vectorized version $\vect(z) \in \R^{c_{j}k_1k_2}$ of its input $z$. Since $\vect^{ij}(z)$ consists of the same variables as in $\vect(z)$ with potentially duplicate items, there is a linear transformation that translates $\vect(z)$ into $\vect^{ij}(z)$. Therefore, we can simply write $V^{ij} \vect(z) = U^{ij} \vect^{ij}(z)$.
{\bf Representing convolutional layers.\enspace} Except of $C^{ij}$, we represent each one of the network's convolutional layers $C^{ld}$ in $f_W$ as linear transformations. As mentioned earlier, we can write $\vect(C^{ld}(z)) = V^{ld}\vect(z)$, for any input $z\in \R^{c_d \times h_d \times w_d}$.
{\bf Representing pooling and rearrangement layers.\enspace} An average pooling layer or a rearrangement layer $C^{ld}$ can be easily represented as a (non-trainable) linear transformation of its input. Namely, we can write $\vect(C^{ld}(z)) = V^{ld}\vect(z)$ for some constant matrix $V^{ld}$. A max-pooling layer can be written as a composition of ReLU activations and multiple (non-trainable) linear transformations, since $\max(x,y) = \sigma(x-y) + y$. Therefore, without loss of generality we can replace the pooling layers with non-trainable linear transformations and ReLU activations.
{\bf Computing the rank.\enspace} Finally, we note that $\vect(C^{ij} \circ z) = U^{ij} \vect^{ij}(z) = V^{ij} \vect(z)$, $\vect(D_l \odot z) = P_l \cdot \vect(z)$ for $P_{l} := \textnormal{diag}(\vect(D_{l}))$. Therefore, we can write
\begin{equation*}
\begin{aligned}
A_W(x) ~&=~ \sum_{\pi \textnormal{ from $i$ to $l_0$}} W^{\pi_{T}\pi_{T-1}} \cdot P_{\pi_{T-1}} \cdots P_{\pi_2} \cdot W^{\pi_2\pi_1} \cdot P_{\pi_1} \cdot U^{ij} \cdot \vect^{ij}(v_j(x)) \\
~&=:~ a(x)^{\top} \cdot U^{ij} \cdot b(x),
\end{aligned}
\end{equation*}
where $a(x)^{\top} ~:=~ \sum_{\pi \textnormal{ from $i$ to $l_0$}} W^{\pi_{T}\pi_{T-1}} \cdot P_{\pi_{T-1}} \cdots P_{\pi_2} \cdot W^{\pi_2\pi_1} \cdot P_{\pi_1}$ and $b(x) := \vect^{ij}(v_j(x))$.
We note that with measure $1$, the matrices $\{P_l\}^{L-1}_{l=1}$ are constant in the neighborhood of $W$. In addition, $a(x)$ and $b(x)$ are computed as multiplications of matrices $W^{ld}$ and $P_l$ excluding $(i,j)=(p,q)$. Therefore, with measure $1$ over the selection of $W$, the Jacobians of $a(x)$ and $b(x)$ with respect to $V^{ij}$ are $0$. Furthermore, due to \eqref{eq:Uij} and the definition of $U^{ij}$, we can write
\begin{equation*}
\begin{aligned}
a(x)^{\top} \cdot U^{ij} \cdot b(x) ~=~ \sum^{N^{ij}}_{t=1} a_t(x)^{\top} \cdot V^{ij} \cdot b_t(x),
\end{aligned}
\end{equation*}
where $a_t(x)$ and $b_t(x)$ are the slices of $a(x)$ and $b(x)$ that are multiplied by the $t$'th $V^{ij}$ block in $U^{ij}$. Since the Jacobians of $a_i(x)$ and $b_i(x)$ with respect to $V^{ij}$ are $0$ with measure $1$ over the selection of $W$, we have,
\begin{equation}
\frac{\partial a(x)^{\top} \cdot U^{ij} \cdot b(x)}{\partial V^{ij}} ~=~ \sum^{N^{ij}}_{t=1} a_t(x) \cdot b_t(x)^{\top}.
\end{equation}
Therefore, we conclude that, with measure $1$ over the selection of $W$, we have $\frac{\partial f_W(x)}{\partial V^{ij}} = \sum^{N^{ij}}_{t=1} a_t(x) \cdot b_t(x)^{\top}$ which is a matrix of rank $\leq N^{ij}$.
\end{proof}
\collinear*
\begin{proof}
Since $\{(v_p,v_1) \in E \mid v_p \in V\} \subset E_T$ is of size $1$, we denote this single layer by $(v_p,v_1)$. Following the proof of Lem.~\ref{lem:rank}, we define a representation of a given $3$-dimensional tensor input $x \in \R^{c_j \times h_j \times w_j}$ as a vector $\vect^{ij}(x) = (x^1\| \dots\| x^{N^{ij}}) \in \R^{N^{ij} c_{j} k_1 k_2}$, where $x^k$ is the vectorization of the $k$'th $c_j \times k_1 \times k_2$ patch of $x$. Similar to the proof of Lem.~\ref{lem:rank}, we can write
\begin{equation}\label{eq:rep_f_2}
\begin{aligned}
f_W(x) ~&=~ \sum_{\pi \textnormal{ from $1$ to $L$}} C^{\pi_{T}\pi_{T-1}} \circ D_{\pi_{T-1}} \odot C^{\pi_{T-1}\pi_{T-2}} \cdots D_{\pi_2} \odot C^{\pi_2 \pi_1} \circ x \\
~&=~ %\sum_{(v_p,v_1)\in E}
H(x) \cdot U^{p1} \cdot \vect^{p1}(x)% \\
%~&=~ \sum_{(v_p,v_1)\in E_T} E(x) \cdot U^{p1} \cdot \vect^{p1}(x), \\
\end{aligned}
\end{equation}
where $H(x) := \sum_{\pi \textnormal{ from $p$ to $L$}} V^{\pi_{T}\pi_{T-1}} \cdot P_{\pi_{T-1}}(x_i) \cdots V^{\pi_2\pi_1} \cdot P_{\pi_1}(x_i)$ and denote $H_{i}:=H(x_i)$ and $H_{i} = (H^1_{i},\dots,H^{N^{p1}}_{i})$, where $H^k_{i}$ is of dimension $c_1k_1k_2$. Hence, we can write,
\begin{equation*}
\frac{\partial f_W(x_i)}{\partial W^{p1}} = \sum^{N^{p1}}_{k=1}H^k_{i} \cdot (x^k_i)^{\top}.
\end{equation*}
We would like to show that $\{x^k_{i},x^k_{j}\}^{N^{p1}}_{k=1}$ are linearly dependent vectors or $f_W \equiv 0$. Assume the opposite by contradiction, i.e., that $\{x^k_{i},x^k_{j}\}^{N^{p1}}_{k=1}$ are not linearly dependent and that $f_W \not\equiv 0$. In particular, by \eqref{eq:rep_f_2}, we have, $W^{p1} \neq 0$.
According to the analysis in Sec.~\ref{sec:convergence}, $\{\frac{\partial f_W(x_i)}{\partial \textnormal{vec}(W^{p1})}\}^{m}_{i=1}$ are collinear vectors. Therefore, for any pair $i,j\in [m]$, there is a scalar $\alpha^{p}_{ij} \in \R$, such that, $\sum^{N^{p1}}_{k=1} H^k_{i} \cdot (x^k_i)^{\top} = \alpha_{ij} \sum^{N^{p1}}_{k=1} H^k_{j} \cdot (x^k_j)^{\top}$.
Consider a given index $i\in [m]$. We would like to show that either $H_{i}=0$ or $\{x^k_{i},x^k_{j}\}^{N^{p1}}_{k=1}$ is a set of linearly dependent vectors. Assume that $H_{i}=0$. Then, $\frac{\partial f_W(x_i)}{\partial W^{p1}}=0$ and $M^{p1}=0$ and $W^{p1}=0$ according to \eqref{eq:batchone}, which implies that $f_W\equiv 0$ (since $(v_p,v_1)$ is the only connection starting from $v_1$). Assume that $H_{i}\neq 0$. Then, there exist $k \in [N^{p1}]$ and $r \in [\textnormal{dim}(H_i)/N^{p1}]$, for which the $r$'th coordinate of $H^k_i$ is non-zero. Therefore, the $r$'th row of $\sum^{N^{p1}}_{k=1} H^k_{i} \cdot (x^k_i)^{\top} - \alpha_{ij} \sum^{N^{p1}}_{k=1} H^k_{j} \cdot (x^k_j)^{\top}=0$ is a non-trivial linear combination of the vectors $\{x^k_i,x^k_j\}^{N^{p1}}_{k=1}$.
% Then, $E_{p,i}\cdot x^{\top}_i=0$, which implies that either $E_{p,i}=0$ or $x_i=0$ (which is false by assumption). Otherwise ($\forall~j \in [m]:~\alpha^p_{ij} \neq 0$), either $E_{p,i}=0$ or $x_i$ and $x_j$ are collinear. Since this argument holds for all $j \in [m]$, unless $E_{p,i}=0$, we obtain that $x_i,x_j$ are collinear for all $j\in [m]$. Note that if $\{x_i,x_j\}$ are collinear for all $j\in [m]$, then, $\{x_i\}^{m}_{i=1}$ is a set of collinear vectors. Therefore, we conclude that either $E_{p,i} = 0$ or $\{x_i\}^{m}_{i=1}$ is a set of collinear vectors (which we assumed to be false). We note that if $E_{p,i}=0$ for all $i\in [m]$, then, $M^{p1}=0$ (see Eq.~8 in the main text), and by Eq.~9 in the main text, we have, $W^{p1}=0$ in contradiction. Therefore, we conclude that $f_W\equiv 0$ or $\{x_i\}^{m}_{i=1}$ are collinear vectors.
\end{proof}
\zerolambda*
\begin{proof}
Let $i\in [m]$ be an index for which $\frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})}\neq 0$. %We consider two possibilities: {\bf (i)} there exist an index $j \in [m]$ for which $\frac{\partial \ell(f_W(x_{j}),y_{j})}{\partial f_W(x_{j})} = 0$ or {\bf (ii)} for all indices $j \in [m]$, we have $\frac{\partial \ell(f_W(x_{j}),y_{j})}{\partial f_W(x_{j})}\neq 0$. Assume that the former holds. Then, by \eqref{eq:}, we have,
%\begin{equation}
%\frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \cdot \frac{\partial f_W(x_{i})}{\partial W^{p1}} ~=~ 0.
%\end{equation}
%This implies that $\sum^{N^{p1}}_{k=1}H^k_{i} \cdot (x^k_i)^{\top} = \frac{\partial f_W(x_{i})}{\partial W^{p1}}=0$ (see the proof of Prop.~\ref{prop:propA}). If $H_i=0$, then according \eqref{eq:rep_f_2}, we have $f_W(x_i)=0$. Otherwise, we obtain that $\{x^k_i\}^{N^{p1}}_{k=1}$ are linearly dependent. Hence, $\{x^k_i,x^k_j\}^{N^{p1}}_{k=1}$ are linearly dependent.
Then, by \eqref{eq:j1j2}, for all $j\in [m]$, we have
\begin{equation}
\frac{\partial f_W(x_{i})}{\partial W^{p1}} ~=~ \left[\frac{\partial \ell(f_W(x_{j}),y_{j})}{\partial f_W(x_{j})}\right]/\left[\frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})}\right]^{-1} \cdot \frac{\partial f_W(x_{j})}{\partial W^{p1}}.
\end{equation}
In particular, $\frac{\partial f_W(x_i)}{\partial \textnormal{vec}(W^{p1})}$ and $\frac{\partial f_W(x_j)}{\partial \textnormal{vec}(W^{p1})}$ are collinear vectors (for all $j \in [m]$). Hence, by the proof of Prop.~\ref{prop:propA}, either $f_W(x_i)=0$ or $\{x^k_i,x^k_j\}^{m}_{i=1}$ is a set of linearly dependent vectors. %Therefore, since $\{x_i\}^{m}_{i=1}$ is not a set of collinear vectors, by \eqref{eq:rep_f_2}, we obtain that $\forall~i \in [m]:~f_W(x_i)=0$.
%Finally, if $\ell$ is convex and $\forall~b\in \R~\nexists~a^* \in \R:~ \ell(a^*,b)=\inf_a \ell(a,b)$, we obtain that $\forall~ i \in [m]:~\frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \neq 0$. Hence, by the above, we conclude that, $\forall~ i \in [m]:~f_W(x_i)=0$ or $\{x^k_i,x^k_j\}^{N^{p1}}_{k=1}$ are linearly dependent vectors (for all $j \in [m]$).
\end{proof}
\newpage
\section{Additional Experiments}
{\bf Experiments on rank minimization.\enspace} To further demonstrate the bias towards rank minimization of SGD with weight decay, we conducted a series of experiments with different learning settings. We follow the same training and evaluation protocol described in Sec.~\ref{sec:rank_exp}. The results are summarized in Figs.~\ref{fig:bs}-\ref{fig:vit_var_wd}.
{\bf Experiments on SGD noise.\enspace} We repeated the experiment in Fig.~\ref{fig:convergence} for training the models on MNIST. As can be seen in Fig.~\ref{fig:convergence2}, similar to the previous experiment, when the models were trained with $\lambda>0$, the weights $W_t$ were not able to converge (i.e., $d(W_{t+1},W_{t})$ does not tend to zero), even though for certain values of $\lambda > 0$ the training accuracy and loss converged. On the other hand, when $\lambda=0$, the distance $d(W_{t+1},W_{t})$ tends to zero.
{\bf Low-rank bias and generalization.\enspace} We looked into the connection between low-rank bias and generalization. In Figs.~\ref{fig:gen_1}-\ref{fig:gen_2} we trained ResNet-18 and VGG-16 instances on CIFAR10 while varying the batch size and keeping $\lambda$ and $\mu$ constant. To provide a fair comparison, we chose $\lambda$ and $\mu$ in each setting to ensure that all models fit the training data perfectly. As can be seen, models trained with smaller batch sizes, i.e. models with lower rank in their weights, tend to generalize better. Based on these findings, we hypothesize that when two neural networks of the same architecture are trained with SGD with different hyperparameters and perfectly fit the data, the one with a lower average rank will outperform the other at test time.
% Finally, when training for cross-entropy loss minimization without weight decay we encounter the `almost convergence' regime discussed in Sec.~\ref{sec:convergence}. Namely, even though perfect convergence is impossible (see Prop.~\ref{prop:propB}) for a fixed sized neural network, the term $d(W_{t+1},W_{t})$ may become as small as we wish by increasing the size of the neural network. Therefore, since the MLP-5-2000 is relatively large (compared to the dataset's size), we may inaccurately get the impression that $d(W_{t+1},W_{t})$ tends to zero.
\begin{figure*}[h]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c@{}c}
\rotatebox[origin=c]{90}{{\tiny MNIST, CE}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/mnist_dist_ce_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/mnist_acc_ce_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/mnist_loss_ce_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/mnist_rank_ce_mlp_2000_lr_01.pdf}}
\\
\rotatebox[origin=c]{90}{{\tiny MNIST, MSE}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/mnist_dist_mse_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/mnist_acc_mse_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/mnist_loss_mse_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/mnist_rank_mse_mlp_2000_lr_01.pdf}}
\\
% \rotatebox[origin=c]{90}{{\tiny CIFAR10, CE}} &
% \raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/cifar10_dist_ce_mlp_2000_lr_01.pdf}} &
% \raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/cifar10_acc_ce_mlp_2000_lr_01.pdf}} &
% \raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/cifar10_loss_ce_mlp_2000_lr_01.pdf}} &
% \raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/cifar10_rank_ce_mlp_2000_lr_01.pdf}}
% \\
% \rotatebox[origin=c]{90}{{\tiny CIFAR10, MSE}} &
% \raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/cifar10_dist_mse_mlp_2000_lr_01.pdf}} &
% \raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/cifar10_acc_mse_mlp_2000_lr_01.pdf}} &
% \raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/cifar10_loss_mse_mlp_2000_lr_01.pdf}} &
% \raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/cifar10_rank_mse_mlp_2000_lr_01.pdf}}
% \\
& {\small (a)} & {\small (b)} & {\small (c)} & {\small (d)}
\end{tabular}
\caption{{\bf Convergence of MLP-5-2000 trained on MNIST with CE/~MSE loss.} In {\bf (a)} we plot the averaged distance between the weight matrices at epoch $t$ and epoch $t+1$, captured by, $\frac{1}{\vert E_T\vert}\sum_{(p,q) \in E_{T}} \|W^{ij}_{t+1} - W^{ij}_{t}\|$. In {\bf (b)} we plot the train accuracy rates, in {\bf (c)} we plot the averaged train loss and in {\bf (d)} we plot the average rank across the trainable matrices.}
\label{fig:convergence2}
\end{figure*}
\begin{figure*}[h]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/rank_batch_size_mlp_100_10_lr_0001.pdf}&
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/rank_batch_size_mlp_100_10_lr_001.pdf}&
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/rank_batch_size_mlp_100_10_lr_01.pdf}&
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/rank_batch_size_mlp_100_10_lr_05.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/acc_batch_size_mlp_100_10_lr_0001.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/acc_batch_size_mlp_100_10_lr_001.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/acc_batch_size_mlp_100_10_lr_01.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/acc_batch_size_mlp_100_10_lr_05.pdf}
\\
{\small $\mu=0.001$} & {\small $\mu=0.01$} & {\small $\mu=0.1$} & {\small $\mu=0.5$}
\end{tabular}
\caption{{\bf Average rank of MLP-BN-10-100 trained on CIFAR10 with various batch sizes.} Each model was trained with a $\lambda=\expnumber{5}{-4}$ weight decay. To estimate the rank, we used an $\epsilon=0.001$ threshold.}
\label{fig:bs}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/rank_wd_mlp_100_10_bs_16.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/rank_wd_mlp_100_10_bs_32.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/rank_wd_mlp_100_10_bs_64.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/rank_wd_mlp_100_10_bs_128.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/acc_wd_mlp_100_10_bs_16.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/acc_wd_mlp_100_10_bs_32.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/acc_wd_mlp_100_10_bs_64.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/acc_wd_mlp_100_10_bs_128.pdf}
\\
{\small $B=16$} & {\small $B=32$} & {\small $B=64$} & {\small $B=128$}
\end{tabular}
\caption{{\bf Average rank of MLP-BN-10-100 trained on CIFAR10 with varying $\lambda$.} Each model was trained with a $\mu=0.1$ initial learning rate and $0.9$ momentum. To estimate the rank, we used an $\epsilon=0.001$ threshold.}
\label{fig:momentum}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/rank_batch_size_resnet_100_5_lr_00125.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/rank_batch_size_resnet_100_5_lr_0025.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/rank_batch_size_resnet_100_5_lr_005.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/rank_batch_size_resnet_100_5_lr_01.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/acc_batch_size_resnet_100_5_lr_00125.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/acc_batch_size_resnet_100_5_lr_0025.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/acc_batch_size_resnet_100_5_lr_005.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/acc_batch_size_resnet_100_5_lr_01.pdf}
\\
{\small $\mu=0.0125$} & {\small $\mu=0.025$} & {\small $\mu=0.05$} & {\small $\mu=0.1$}
\end{tabular}
\caption{{\bf Average rank of RES-BN-5-500 trained on CIFAR10 with various batch sizes.} Each model was trained with a $\lambda=\expnumber{5}{-4}$ weight decay. To estimate the rank, we used an $\epsilon=0.001$ threshold.}
\label{fig:bs_resnet}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/rank_batch_size_mlp_nobn_500_5_lr_00125.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/rank_batch_size_mlp_nobn_500_5_lr_0025.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/rank_batch_size_mlp_nobn_500_5_lr_005.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/rank_batch_size_mlp_nobn_500_5_lr_01.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/acc_batch_size_mlp_nobn_500_5_lr_00125.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/acc_batch_size_mlp_nobn_500_5_lr_0025.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/acc_batch_size_mlp_nobn_500_5_lr_005.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/acc_batch_size_mlp_nobn_500_5_lr_01.pdf}
\\
{\small $\mu=0.0125$} & {\small $\mu=0.025$} & {\small $\mu=0.05$} & {\small $\mu=0.1$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of MLP-5-500 trained on CIFAR10 with various batch sizes.} Each model was trained with a $\lambda=\expnumber{5}{-4}$ weight decay. To estimate the rank, we used an $\epsilon=0.001$ threshold.}
\label{fig:bs_mlp_nobn}
\end{figure*}
% \begin{figure*}[t]
% \centering
% \begin{tabular}{c@{}c@{}c@{}c@{}c}
% \includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_01_batch_size.pdf} &
% \includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_05_batch_size.pdf} &
% \includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_1_batch_size.pdf} &
% \includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/rank_mlp_bn_lr_3_batch_size.pdf}
% \\
% \includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_01_batch_size.pdf} &
% \includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_05_batch_size.pdf} &
% \includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_1_batch_size.pdf} &
% \includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/batch_size/acc_mlp_bn_lr_3_batch_size.pdf}
% \\
% {\small $\mu=0.01$} & {\small $\mu=0.05$} & {\small $\mu=0.1$} & {\small $\mu=0.3$}
% \end{tabular}
% \caption{{\bf Average ranks and accuracy rates of {\color{blue}MLP-BN-10-100 (no-momentum)} trained on CIFAR10 with various batch sizes {\color{blue}($\epsilon=0.001$)}.} See Fig.~\ref{fig:bs} for details.}
% \end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/rank_mlp_bn_bs_4_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/rank_mlp_bn_bs_8_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/rank_mlp_bn_bs_16_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/rank_mlp_bn_bs_32_lr_1_weight_decay.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/acc_mlp_bn_bs_4_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/acc_mlp_bn_bs_8_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/acc_mlp_bn_bs_16_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_bn/weight_decay/acc_mlp_bn_bs_32_lr_1_weight_decay.pdf}
\\
{\small $B=4$} & {\small $B=8$} & {\small $B=16$} & {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of MLP-BN-10-100 trained on CIFAR10 with varying $\lambda$.} Each model was trained with a $\mu=0.1$ initial learning rate. To estimate the rank, we used an $\epsilon=0.001$ threshold.}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/rank_mlp_no_bn_lr_05_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/rank_mlp_no_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/rank_mlp_no_bn_lr_2_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/rank_mlp_no_bn_lr_3_batch_size.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/acc_mlp_no_bn_lr_05_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/acc_mlp_no_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/acc_mlp_no_bn_lr_2_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/batch_size/acc_mlp_no_bn_lr_3_batch_size.pdf}
\\
{\small $\mu=0.05$} & {\small $\mu=0.1$} & {\small $\mu=0.2$} & {\small $\mu=0.3$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of MLP-5-500 trained on CIFAR10 with various batch sizes.} The models were trained with a $\lambda=\expnumber{5}{-4}$ weight decay. To estimate the rank, we used an $\epsilon=0.001$ threshold.}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/rank_mlp_no_bn_bs_4_lr_025_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/rank_mlp_no_bn_bs_8_lr_025_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/rank_mlp_no_bn_bs_16_lr_025_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/rank_mlp_no_bn_bs_32_lr_025_weight_decay.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/acc_mlp_no_bn_bs_4_lr_025_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/acc_mlp_no_bn_bs_8_lr_025_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/acc_mlp_no_bn_bs_16_lr_025_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_new_no_bn/weight_decay/acc_mlp_no_bn_bs_32_lr_025_weight_decay.pdf}
\\
{\small $B=4$} & {\small $B=8$} & {\small $B=16$} & {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of MLP-5-500 trained on CIFAR10 with varying $\lambda$.} The models were trained with a $\mu=0.025$ initial learning rate. To estimate the rank, we used an $\epsilon=0.001$ threshold.}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/rank_resnet_bn_lr_5_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/rank_resnet_bn_lr_8_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/rank_resnet_bn_lr_10_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/rank_resnet_bn_lr_20_batch_size.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/acc_resnet_bn_lr_5_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/acc_resnet_bn_lr_8_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/acc_resnet_bn_lr_10_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/acc_resnet_bn_lr_20_batch_size.pdf}
\\
{\small $\mu=0.5$} & {\small $\mu=0.8$} & {\small $\mu=1.0$} & {\small $\mu=2.0$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of ResNet-18 trained on CIFAR10 with various batch sizes.} The models were trained with a $\lambda=\expnumber{5}{-4}$ weight decay. To estimate the rank, we used an $\epsilon=0.001$ threshold.}
\label{fig:bs_resnet18}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/rank_resnet_bn_batch_size_4_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/rank_resnet_bn_batch_size_8_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/rank_resnet_bn_batch_size_16_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/rank_resnet_bn_batch_size_32_learning_rate.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/acc_resnet_bn_batch_size_4_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/acc_resnet_bn_batch_size_8_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/acc_resnet_bn_batch_size_16_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/acc_resnet_bn_batch_size_32_learning_rate.pdf}
\\
{\small $B=4$} & {\small $B=8$} & {\small $B=16$} & {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of ResNet-18 trained on CIFAR10 with varying $\mu$.} The models were trained with a $\mu=\expnumber{5}{-4}$ initial learning rate. To estimate the rank, we used an $\epsilon=0.001$ threshold.}
\label{fig:lr_resnet18}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/rank_vgg_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/rank_vgg_bn_lr_2_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/rank_vgg_bn_lr_3_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/rank_vgg_bn_lr_5_batch_size.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/acc_vgg_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/acc_vgg_bn_lr_2_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/acc_vgg_bn_lr_3_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/acc_vgg_bn_lr_5_batch_size.pdf}
\\
{\small $\mu=0.1$} & {\small $\mu=0.2$} & {\small $\mu=0.3$} & {\small $\mu=0.5$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of VGG-16 trained on CIFAR10 with various batch sizes.} The models were trained with a $\mu=\expnumber{5}{-4}$ initial learning rate. To estimate the rank, we used an $\epsilon=0.001$ threshold.}
\label{fig:bs_vgg}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/rank_vgg_bn_batch_size_4_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/rank_vgg_bn_batch_size_8_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/rank_vgg_bn_batch_size_16_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/rank_vgg_bn_batch_size_32_lr_1_weight_decay.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/acc_vgg_bn_batch_size_4_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/acc_vgg_bn_batch_size_8_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/acc_vgg_bn_batch_size_16_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/acc_vgg_bn_batch_size_32_lr_1_weight_decay.pdf}
\\
{\small $B=4$} & {\small $B=8$} & {\small $B=16$} & {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of VGG-16 trained on CIFAR10 with varying $\lambda$.} The models were trained with a $\mu=0.1$ initial learning rate. To estimate the rank, we used an $\epsilon=0.01$ threshold.}
\label{fig:wd_vgg}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/rank_vgg_bn_batch_size_4_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/rank_vgg_bn_batch_size_8_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/rank_vgg_bn_batch_size_16_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/rank_vgg_bn_batch_size_32_learning_rate.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/acc_vgg_bn_batch_size_4_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/acc_vgg_bn_batch_size_8_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/acc_vgg_bn_batch_size_16_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/acc_vgg_bn_batch_size_32_learning_rate.pdf}
\\
{\small $B=4$} & {\small $B=8$} & {\small $B=16$} & {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of VGG-16 trained on CIFAR10 with varying $\mu$.} The models were trained with a $\mu=\expnumber{5}{-4}$ weight decay. To estimate the rank, we used an $\epsilon=0.001$ threshold.}
\label{fig:lr_vgg}
\end{figure*}
\begin{figure*}[t] % vit weight decay
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
%\includegraphics[width=0.20\linewidth]{plots/vit/weight_decay/varywd_B8_avgrank.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/weight_decay/varywd_B16_avgrank.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/weight_decay/varywd_B32_avgrank.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/weight_decay/varywd_B64_avgrank.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/weight_decay/varywd_B128_avgrank.pdf}
\\
%\includegraphics[width=0.20\linewidth]{plots/vit/weight_decay/varywd_B8_accuracy.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/weight_decay/varywd_B16_accuracy.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/weight_decay/varywd_B32_accuracy.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/weight_decay/varywd_B64_accuracy.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vit/weight_decay/varywd_B128_accuracy.pdf}
\\
%{\small $B=8$} &
{\small $B=16$} & {\small $B=32$} & {\small $B=64$} & {\small $B=128$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of ViT trained on CIFAR10 with varying $\lambda$.} The models were trained with a $\mu=\expnumber{4}{-2}$ initial learning rate. To estimate the rank, we used an $\epsilon=0.01$ threshold.}
\label{fig:vit_var_wd}
\end{figure*}
\begin{figure*}[t] % resnet generalization
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.33\linewidth]{plots/generalization/resnet/rank_resnet_bn_lr_001_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/generalization/resnet/rank_resnet_bn_lr_005_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/generalization/resnet/rank_resnet_bn_lr_01_batch_size.pdf}
\\
\includegraphics[width=0.33\linewidth]{plots/generalization/resnet/acc_resnet_bn_lr_001_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/generalization/resnet/acc_resnet_bn_lr_005_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/generalization/resnet/acc_resnet_bn_lr_01_batch_size.pdf}
\\
{\bf {(a)}} & {\bf {(b)}} & {\bf {(c)}}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of ResNet-18 trained on CIFAR10 with different batch sizes.} We show the relationship between average rank and test performance. {\bf(a)} was trained with $\mu=0.001, \lambda=\expnumber{6}{-3}$, {\bf(b)} was trained with $\mu=0.005, \lambda=\expnumber{6}{-3}$, and {\bf (c)} was trained with $\mu=0.01, \lambda=4e-4$. To estimate the rank, we used an $\epsilon=0.05$ threshold.}
\label{fig:gen_1}
\end{figure*}
\begin{figure*}[t] % vgg generalization
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.33\linewidth]{plots/generalization/vgg/rank_vgg_bn_lr_001_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/generalization/vgg/rank_vgg_bn_lr_005_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/generalization/vgg/rank_vgg_bn_lr_01_batch_size.pdf}
\\
\includegraphics[width=0.33\linewidth]{plots/generalization/vgg/acc_vgg_bn_lr_001_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/generalization/vgg/acc_vgg_bn_lr_005_batch_size.pdf} &
\includegraphics[width=0.33\linewidth]{plots/generalization/vgg/acc_vgg_bn_lr_01_batch_size.pdf}
\\
{\bf {(a)}} & {\bf {(b)}} & {\bf {(c)}}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of VGG-16 trained on CIFAR10 with different batch sizes.} We show the relationship between average rank and test performance. {\bf(a)} was trained with $\mu=0.001, \lambda=6e-3$, {\bf(b)} was trained with $\mu=0.005, \lambda=5e-4$, and {\bf(c)} was trained with $\mu=0.01, \lambda=4e-4$. To estimate the rank, we used an $\epsilon=0.04$ threshold.}
\label{fig:gen_2}
\end{figure*}
\end{document}
\def\year{2022}\relax
%File: formatting-instructions-latex-2022.tex
%release 2022.1
\documentclass[letterpaper]{article} % DO NOT CHANGE THIS
\usepackage{aaai22} % DO NOT CHANGE THIS
\usepackage{times} % DO NOT CHANGE THIS
\usepackage{helvet} % DO NOT CHANGE THIS
\usepackage{courier} % DO NOT CHANGE THIS
\usepackage[hyphens]{url} % DO NOT CHANGE THIS
\usepackage{graphicx} % DO NOT CHANGE THIS
\urlstyle{rm} % DO NOT CHANGE THIS
\def\UrlFont{\rm} % DO NOT CHANGE THIS
\usepackage{natbib} % DO NOT CHANGE THIS AND DO NOT ADD ANY OPTIONS TO IT
\usepackage{caption} % DO NOT CHANGE THIS AND DO NOT ADD ANY OPTIONS TO IT
\DeclareCaptionStyle{ruled}{labelfont=normalfont,labelsep=colon,strut=off} % DO NOT CHANGE THIS
\frenchspacing % DO NOT CHANGE THIS
\setlength{\pdfpagewidth}{8.5in} % DO NOT CHANGE THIS
\setlength{\pdfpageheight}{11in} % DO NOT CHANGE THIS
%
% These are recommended to typeset algorithms but not required. See the subsubsection on algorithms. Remove them if you don't have algorithms in your paper.
\usepackage{algorithm}
\usepackage{algorithmic}
\usepackage{amsmath,amssymb}
\usepackage{thm-restate,enumitem}
\usepackage{xspace}
\usepackage{bbm}
\usepackage{xcolor}
\usepackage{amsthm}
\newcommand{\expnumber}[2]{{#1}\mathrm{e}{#2}}
\newcommand{\cC}{\mathcal{C}}
\newcommand{\cD}{\mathcal{D}}
\newcommand{\cF}{\mathcal{F}}
\newcommand{\cG}{\mathcal{G}}
\newcommand{\cH}{\mathcal{H}}
\newcommand{\cN}{\mathcal{N}}
\newcommand{\cP}{\mathcal{P}}
\newcommand{\cU}{\mathcal{U}}
\newcommand{\cX}{\mathcal{X}}
\newcommand{\cY}{\mathcal{Y}}
\newcommand{\E}{\mathbb{E}}
\newcommand{\R}{\mathbb{R}}
\renewcommand{\P}{\mathbb{P}}
\newcommand{\B}{\mathbb{B}}
\newcommand{\bI}{\mathbb{I}}
\newcommand{\Var}{\mathrm{Var}}
\newcommand{\EE}[1]{\E\left[ #1\right]}
\newcommand{\PP}[1]{\P\left[ #1\right]}
\newcommand{\I}[1]{\bI\left[ #1\right]}
\newcommand{\tx}{\tilde{x}}
\newcommand{\ty}{\tilde{y}}
\newcommand{\tS}{\tilde{S}}
\newcommand{\tX}{\tilde{\cX}}
\newcommand{\tU}{\tilde{U}}
\newcommand{\tT}{\tilde{T}}
\newcommand{\tcG}{\tilde{\cG}}
\newcommand{\tP}{\tilde{P}}
\newcommand{\tcP}{\tilde{\cP}}
\newcommand{\tcU}{\tilde{\cU}}
\newcommand{\tf}{\tilde{f}}
\newcommand{\tig}{\tilde{g}}
\newcommand{\tih}{\tilde{h}}
\newcommand{\clD}{\cD_{\cC}}
\newcommand{\one}{\mathbbm{1}}
\newcommand{\rank}{\textnormal{rank}}
\newcommand{\vect}{\textnormal{vec}}
\newcommand{\threepartdef}[6]
{
\left\{
\begin{array}{lll}
#1 & \mbox{if } #2 \\
#3 & \mbox{if } #4 \\
#5 & \mbox{if } #6
\end{array}
\right.
}
\newcommand{\twopartdef}[4]
{
\left\{
\begin{array}{lll}
#1 & \mbox{if } #2 \\
#3 & \mbox{if } #4 \\
\end{array}
\right.
}
\newtheorem{theorem}{Theorem}
\newtheorem{lemma}{Lemma}
\newtheorem{proposition}{Proposition}
\newtheorem{corollary}{Corollary}
\newtheorem{definition}{Definition}
\newtheorem{remark}{Remark}
% These are are recommended to typeset listings but not required. See the subsubsection on listing. Remove this block if you don't have listings in your paper.
\usepackage{newfloat}
\usepackage{listings}
\lstset{%
basicstyle={\footnotesize\ttfamily},% footnotesize acceptable for monospace
numbers=left,numberstyle=\footnotesize,xleftmargin=2em,% show line numbers, remove this entire line if you don't want the numbers.
aboveskip=0pt,belowskip=0pt,%
showstringspaces=false,tabsize=2,breaklines=true}
\floatstyle{ruled}
\newfloat{listing}{tb}{lst}{}
\floatname{listing}{Listing}
%
%\nocopyright
%
% PDF Info Is REQUIRED.
% For /Title, write your title in Mixed Case.
% Don't use accents or commands. Retain the parentheses.
% For /Author, add all authors within the parentheses,
% separated by commas. No accents, special characters
% or commands are allowed.
% Keep the /TemplateVersion tag as is
\pdfinfo{
/Title (AAAI Press Formatting Instructions for Authors Using LaTeX -- A Guide)
/Author (AAAI Press Staff, Pater Patel Schneider, Sunil Issar, J. Scott Penberthy, George Ferguson, Hans Guesgen, Francisco Cruz, Marc Pujol-Gonzalez)
/TemplateVersion (2022.1)
}
% DISALLOWED PACKAGES
% \usepackage{authblk} -- This package is specifically forbidden
% \usepackage{balance} -- This package is specifically forbidden
% \usepackage{color (if used in text)
% \usepackage{CJK} -- This package is specifically forbidden
% \usepackage{float} -- This package is specifically forbidden
% \usepackage{flushend} -- This package is specifically forbidden
% \usepackage{fontenc} -- This package is specifically forbidden
% \usepackage{fullpage} -- This package is specifically forbidden
% \usepackage{geometry} -- This package is specifically forbidden
% \usepackage{grffile} -- This package is specifically forbidden
% \usepackage{hyperref} -- This package is specifically forbidden
% \usepackage{navigator} -- This package is specifically forbidden
% (or any other package that embeds links such as navigator or hyperref)
% \indentfirst} -- This package is specifically forbidden
% \layout} -- This package is specifically forbidden
% \multicol} -- This package is specifically forbidden
% \nameref} -- This package is specifically forbidden
% \usepackage{savetrees} -- This package is specifically forbidden
% \usepackage{setspace} -- This package is specifically forbidden
% \usepackage{stfloats} -- This package is specifically forbidden
% \usepackage{tabu} -- This package is specifically forbidden
% \usepackage{titlesec} -- This package is specifically forbidden
% \usepackage{tocbibind} -- This package is specifically forbidden
% \usepackage{ulem} -- This package is specifically forbidden
% \usepackage{wrapfig} -- This package is specifically forbidden
% DISALLOWED COMMANDS
% \nocopyright -- Your paper will not be published if you use this command
% \addtolength -- This command may not be used
% \balance -- This command may not be used
% \baselinestretch -- Your paper will not be published if you use this command
% \clearpage -- No page breaks of any kind may be used for the final version of your paper
% \columnsep -- This command may not be used
% \newpage -- No page breaks of any kind may be used for the final version of your paper
% \pagebreak -- No page breaks of any kind may be used for the final version of your paperr
% \pagestyle -- This command may not be used
% \tiny -- This is not an acceptable font size.
% \vspace{- -- No negative value may be used in proximity of a caption, figure, table, section, subsection, subsubsection, or reference
% \vskip{- -- No negative value may be used to alter spacing above or below a caption, figure, table, section, subsection, subsubsection, or reference
\setcounter{secnumdepth}{2} %May be changed to 1 or 2 if section numbers are desired.
% The file aaai22.sty is the style file for AAAI Press
% proceedings, working notes, and technical reports.
%
% Title
% Your title must be in mixed case, not sentence case.
% That means all verbs (including short verbs like be, is, using,and go),
% nouns, adverbs, adjectives should be capitalized, including both words in hyphenated terms, while
% articles, conjunctions, and prepositions are lower case unless they
% directly follow a colon or long dash
\title{Weight Decay and SGD Provably Induce Low-Rank Bias in Neural Networks}
\author{
%Authors
% All authors must be in the same font size and format.
Written by AAAI Press Staff\textsuperscript{\rm 1}\thanks{With help from the AAAI Publications Committee.}\\
AAAI Style Contributions by Pater Patel Schneider,
Sunil Issar,\\
J. Scott Penberthy,
George Ferguson,
Hans Guesgen,
Francisco Cruz\equalcontrib,
Marc Pujol-Gonzalez\equalcontrib
}
\affiliations{
%Afiliations
\textsuperscript{\rm 1}Association for the Advancement of Artificial Intelligence\\
% If you have multiple authors and multiple affiliations
% use superscripts in text and roman font to identify them.
% For example,
% Sunil Issar, \textsuperscript{\rm 2}
% J. Scott Penberthy, \textsuperscript{\rm 3}
% George Ferguson,\textsuperscript{\rm 4}
% Hans Guesgen, \textsuperscript{\rm 5}.
% Note that the comma should be placed BEFORE the superscript for optimum readability
2275 East Bayshore Road, Suite 160\\
Palo Alto, California 94303\\
% email address must be in roman text type, not monospace or sans serif
publications22@aaai.org
%
% See more examples next
}
%Example, Single Author, ->> remove \iffalse,\fi and place them surrounding AAAI title to use it
\iffalse
\title{Weight Decay and SGD Provably Induce Low-Rank Bias in Neural Networks}
\author {
Author Name
}
\affiliations{
Affiliation\\
Affiliation Line 2\\
name@example.com
}
\fi
\iffalse
%Example, Multiple Authors, ->> remove \iffalse,\fi and place them surrounding AAAI title to use it
\title{Weight Decay and SGD Provably Induce Low-Rank Bias in Neural Networks}
\author {
% Authors
First Author Name,\textsuperscript{\rm 1}
Second Author Name, \textsuperscript{\rm 2}
Third Author Name \textsuperscript{\rm 1}
}
\affiliations {
% Affiliations
\textsuperscript{\rm 1} Affiliation 1\\
\textsuperscript{\rm 2} Affiliation 2\\
firstAuthor@affiliation1.com, secondAuthor@affilation2.com, thirdAuthor@affiliation1.com
}
\fi
% % REMOVE THIS: bibentry
% % This is only needed to show inline citations in the guidelines document. You should not need it and can safely delete it.
% \usepackage{bibentry}
% % END REMOVE bibentry
\begin{document}
\maketitle
\begin{abstract}
We analyze deep ReLU neural networks trained with mini-batch Stochastic Gradient Descent (SGD) and weight decay. We show, both theoretically and empirically, that when training a neural network using SGD with weight decay and small batch size, the resulting weight matrices are expected to be of small rank. Our analysis relies on a minimal set of assumptions and the neural networks may be arbitrarily wide or deep, and may include residual connections, as well as batch normalization and convolutional layers. Furthermore, we suggest that the asymptotic presence of SGD "noise" is also due to the bias towards small rank. In particular, we prove that when training with weight decay, SGD yields low rank minimizers that cannot interpolate all the training data.
%TP%Furthermore, we study the source of SGD noise and prove that when training with weight decay, the only solutions of SGD at convergence are zero functions.
\end{abstract}
\section{Introduction}\label{sec:intro}
Over the past few years, neural networks have challenged machine learning theory with several mysteries. We consider two of them here that are, to our knowledge, so far unsolved. The first is about the mathematical reasons for a bias in Stochastic Gradient Descent (SGD) towards neural networks with low-rank weight matrices, a fact which has been empirically observed (e.g., \cite{DBLP:journals/corr/abs-2201-12760,DBLP:journals/corr/abs-2103-10427}). The second is regarding the role and origin of SGD noise (e.g., \cite{DBLP:journals/corr/abs-2110-06914}). The origin of the SGD noise represents a puzzling and rather neglected question: {\em why does SGD, under certain conditions, never stops making updates?} In this paper, we report a simple observation that solves both mysteries and, somewhat surprisingly, shows that they have the same mathematical root. The observation is that $\frac{\partial f_W}{\partial W^i}$ is a matrix of rank $\leq 1$, where $f_W$ is a deep ReLU network with a scalar output and weight matrices $\{W^i\}_{i \in I}$ ($I$ is some finite set of indices)\footnote{When $f_W$ outputs a $k$-dimensional vector, each output component of the tensor corresponding to $\frac{\partial f}{\partial W^i}$ is a matrix of rank at most $1$.}.
{\color{red}Why is low-rank/SGD noise important?}
\subsection{Related Work}\label{sec:related_work}
% Stochastic gradient descent (SGD) is one of the standard workhorses for optimizing deep models~\citep{bottou-91c}. Though initially proposed to remedy the computational bottleneck of gradient descent (GD), recent studies suggest SGD also induces a crucial implicit regularization, which prevents the overparameterized models from converging to the minima that cannot generalize well~\cite{DBLP:journals/corr/ZhangBHRV16,https://doi.org/10.48550/arxiv.1711.04623,https://doi.org/10.48550/arxiv.1609.04836,zhu}. Empirical studies suggest that (i) SGD outperforms GD~\cite{zhu}, (ii) small batch SGD generalizes better than large batch SGD~\citep{NIPS2017_a5e0ff62,https://doi.org/10.48550/arxiv.1609.04836}, and (iii) gradient descent with additional noise cannot compete with SGD~\cite{zhu}. However, despite many efforts, the potentially important, yet implicit, regularization induced by SGD noise has never been fully understood. Furthermore, we are unaware of any study of the noise in the output of the network during SGD training, a noise which seems to be present even at very long training times, suggesting that true convergence to a single set of weight matrices matrices for all data points is never reached.
A prominent thread in the recent literature is devoted to the study of the implicit bias in neural networks towards rank minimization. Most of the interest was on the matrix factorization problem, which corresponds to training a depth-2 linear neural network with multiple outputs w.r.t. the square loss. As described by~\cite{DBLP:journals/corr/abs-2201-12760},~\cite{https://doi.org/10.48550/arxiv.1705.09280} initially conjectured that the implicit regularization in matrix factorization can be characterized by the nuclear norm of the corresponding linear predictor. This conjecture was formally refuted by~\cite{DBLP:journals/corr/abs-2012-09839}. \cite{DBLP:journals/corr/abs-2005-06398} conjectured that the implicit regularization in matrix factorization can be explained by rank minimization, and also hypothesized that some notion of rank minimization may be key to explaining generalization in deep learning. \cite{DBLP:journals/corr/abs-2012-09839} established evidence that the implicit regularization in matrix factorization is a heuristic for rank minimization. Beyond factorization problems, \cite{DBLP:journals/corr/abs-2006-06657} showed that in linear networks of output dimension 1, gradient flow (GF) w.r.t. exponentially-tailed classification losses converges to networks where the weight matrix of every layer is of rank 1.
However, with nonlinear neural networks, things are less clear. Empirically, a series of papers~\citep{NIPS2014_2afe4567,10.5555/3294771.3294853,s21165599,8099498,pmlr-v80-arora18b} showed that replacing the weight matrices by low-rank approximations results in only a small drop in accuracy. This suggests that the weight matrices in practice may be close to low-rank matrices. However, whether they provably behave this way remains unclear. \cite{DBLP:journals/corr/abs-2201-12760} show that for nonlinear ReLU networks, GF does not minimize rank. They also argue that ReLU networks of sufficient depth can have low-rank solutions under $\ell_2$ norm minimization. This interesting result, however, applies to layers added to a network that already solves the problem and may not have any low-rank bias. It is not directly related to the mechanism described in this paper, which applies to all layers in the network but only in the presence of regularization and SGD, unlike the phenomenon described in \cite{DBLP:journals/corr/abs-2201-12760}.
\subsection{Contributions}
We provide a mathematical analysis of the origin of SGD noise and of the implicit rank-minimization of SGD. We analyze deep ReLU networks trained with mini-batch stochastic gradient descent and regularization (i.e., weight decay). Our analysis is fairly generic: the neural networks may include linear layers, residual connections, and batch normalization layers. In addition, the loss function is only assumed to be differentiable. The central contributions are:
\begin{itemize}[leftmargin=0.5cm]
\item In Props.~\ref{prop:propA}-\ref{prop:propB}, we study the source of SGD noise. The result shows that SGD noise must be always present, even asymptotically, regardless of the batch size, as long as we incorporate weight decay. Thus, SGD noise is a generic property; this also means that, generically, there is never convergence of SGD. These predictions are later validated empirically in Sec.~\ref{sec:sgd_noise_exp}.
\item In Thm.~\ref{thm:batches} we prove that mini-batch SGD has an implicit bias towards networks with low-rank matrices. This theorem connects batch size, weight decay, optimization and rank.
\item In particular, in Sec.~\ref{sec:lowrank}, we theoretically predict that the batch size, weight decay and learning rate act as low-rank regularizers. For instance, we predict that the presence of weight decay is crucial to obtain low-rank solutions. These predictions are later validated empirically in Sec.~\ref{sec:rank_exp}.
\end{itemize}
\section{Problem Setup}\label{sec:setup}
We consider the problem of training a model for a standard classification or regression problem. Formally, the target task is defined by a distribution $P$ over samples $(x,y)\in \cX\times \cY$, where $\cX \subset \R^{c_{1} \times h_1 \times w_1}$ is the instance space (e.g., images), and $\cY \subset \R^k$ is a label space. %To simplify the presentation, we use one-hot encoding for the label space, that is, the labels are represented by the unit vectors in $\R^C$, and $\cY_C=\{e_c: c=1,\ldots,C\}$ where $e_c \in \R^C$ is the $c$th standard unit vector in $\R^C$; with a slight abuse of notation, sometimes we will also write $y=c$ instead of $y=e_c$.
%For a pair $(x,y)$ distributed by $P$, we denote by $P_c$ the class conditional distribution of $x$ given $y=c$ (i.e., $P_c(\boldsymbol{\cdot})=\P[x \in \boldsymbol{\cdot} \mid y=c]$).
A function $f_W \in \cF \subset \{f':\cX\to \R^k\}$ {\color{red}say something about $W$} assigns a prediction to an input point $x \in \cX$, and its performance on the distribution $P$ is measured by the expected risk
\begin{equation}
\label{eq:loss}
L_P(f_W)~:=~\E_{(x,y)\sim P}[\ell(f_W(x),y)],
\end{equation}
where $\ell:\R^k \times \cY \to [0,\infty)$ is a non-negative, differentiable, loss function (e.g., MSE or cross-entropy losses). For simplicity, throughout the paper, we focus on the case where $k=1$.
Since we do not have direct access to the full population distribution $P$, the goal is to learn a predictor, $f_W$, from some balanced training data $S = \{(x_i,y_i)\}_{i=1}^m$ of independent and identically distributed (i.i.d.) samples drawn from $P$ along with regularization that controls the complexity of the learned model and typically helps reducing overfitting. Namely, we intend to minimize the regularized empirical risk
\begin{equation}
\label{eq:loss_emp}
L^{\lambda}_S(f_W)~:=~\frac{1}{m}\sum^{m}_{i=1}\ell(f_W(x_i),y_i) + \lambda \|W\|^2_2,
\end{equation}
where $\lambda > 0$ is predefined hyperparameter, controlling the amount of regularization. In order to minimize this objective, we typically use mini-batch SGD, which is described below in more details.
% Finally, the performance of the trained model is evaluated using the train and test accuracy rates, which are computed as follows:
% $err_S(f_W) = \sum^{m}_{i=1}\ell(f_W(x_i),y_i)$ and $err_P(f_W) = \E_{(x,y)\sim P}[\ell(f_W(x),y)]$.
%We assume that the $m_0$ samples are taken directly from the classes $P_c$, following the observation in~\cite{Papyan24652} that the data should be balanced in order to achieve a high degree of neural collapse.
% \paragraph{Notation.} Throughout the paper, we use the following notations. For an integer $k\geq 1$, $[k]=\{1,\ldots,k\}$. For any real vector $z$, $\|z\|$ denotes its Euclidean norm. For a finite set $A$, we denote by $U[A]$ the uniform distribution over $A$. We denote by $\bI:\{\textnormal{True},\textnormal{False}\} \to \{0,1\}$ the indicator function.
\subsection{Architectures and Training}
In this work, the function $f_W$ represents a neural network, consisting of a set of layers of weights interlaced with activation units. We employ a fairly generic definition of a neural network, that includes convolutional layers and residual connections. %To include normalization layers, we extend $f_W$'s input to be $(x;\tX)$, where $x\in \R^{n}$ and $\tX = \{x_{ij}\}^{B}_{j=1} \subset \R^{n}$ is a batch of samples (see details below). For simplicity, we denote by $f_W(x)$ a network without normalization layers. When $\tX$ is clear from context, we simply write $\tf_W(x) := f_W(x;\tX)$.
\paragraph{Optimization.} {\color{red}edit the detaching} In this work, we consider optimizing $f_W$ using mini-batch stochastic gradient descent (SGD), possibly with detaching certain layers. Namely, we aim at minimizing the regularized empirical risk $L^{\lambda}_{S}(f_{W})$ by applying SGD for a certain number of iterations. We initialize $W_0$ using a standard initialization procedure and iteratively update $W_t$. At each iteration, we sample a subset $\tS = \{(x_{ij},y_{ij})\}^{B}_{j=1} \subset S$ uniformly at random and update $W_{t+1} \leftarrow W_t - \eta_t \cdot \bar{\nabla}_W L^{\lambda}_{\tS}(f_{W_t}(~\cdot~;\tX))$, where $\tX = \{x_{ij}\}^{B}_{j=1}$ and $\eta_t > 0$ is our learning rate at iteration $t$. Here, $f_{W}(~\cdot~;\tX) : \R^{n} \to \R$ is the function given by $f_W(x;\tX)$ for a fixed value of $\tX$. Furthermore, for a given function $g_{\theta} = g^L_{\theta} \circ \dots g^1_{\theta}:\R^{n} \to \R^k$ where $\{g^{i_j}_{\theta}\}^{r}_{j=1}$ are detached, $\frac{\bar{\partial} g_{\theta}(z)}{\partial \theta} = \bar{\nabla}_{\theta} g_{\theta}(z)$ denotes the {\color{red}pseudo gradient} of $g_{\theta}$, which is computed using the chain rule, while treating $\{g^{i_j}_{\theta}\}^{r}_{j=1}$ as independent of $\theta$.
\paragraph{Network architecture.} Formally, $f_W$ is as a directed acyclic graph (DAG) $G=(V,E)$, where $V= \{v_1,\dots,v_L\}$ consists of the various layers of the network and each edge $e_{ij}=(v_i,v_j) \in E$ specifies a connection between two layers. Each layer is a function $v_{i}:\R^{c_{1} \times h_1 \times w_1} \to \R^{c_{i} \times h_i \times w_i}$ and each connection holds a linear transformation $C^{ij} : \R^{c_{j} \times h_j \times w_j} \to \R^{c_i \times h_i \times w_i}$. The layers are divided into three categories: the input layer $v_1$, the output layer $v_L$ and intermediate layers. In this setting, we do not have connections directed towards the input layer nor connections coming out of the output layer (i.e., $\forall~i \in [L]:~(v_i,v_L), (v_0,v_i) \notin E$). Given an input $x$ and batch $\tX$, the output of a given layer $v_i$ is evaluated as follows
\begin{equation*}
v_i(x;\tX) ~:=~ \sigma\left(n_i\left[\sum_{j \in \textnormal{pred}(i)} C^{ij}(v_{j}(x;\tX));\tX\right]\right),
\end{equation*}
except for the output layer $v_L$ that computes
\begin{equation*}
f_W(x;\tX) ~:=~ v_L(x;\tX) ~:=~ \sum_{j \in \textnormal{pred}(L)}C^{Lj} (v_{j}(x;\tX)),
\end{equation*}
where $\textnormal{pred}(i)$ is the set of indices $j$, such that $(v_i,v_j) \in E$. Here, $\sigma$ is an element-wise activation function (e.g., ReLU, Leaky ReLU, sigmoid).
In this setting, the transformations $C^{pq}$ could be trainable, $(v_p,v_q) \in E_T$ (e.g., convolutional layers), or constant, $(v_p,v_q) \notin E_T$ (e.g., residual connections or a pooling layer). Throughout the analysis, we consider paths within the graph $G$, denoted by $\pi = (\pi_0,\dots,\pi_T)$, where $\pi_0=1$, $\pi_T = L$ and for all $i=0,\dots,T-1:~ (v_{\pi_{i}},v_{\pi_{i+1}}) \in E$.
\paragraph{Convolutional layers.} {\color{red}add a formal def maybe?} Throughout the paper we consider convolutional layers~\citep{6795724,10.5555/646469.691875} (see also~\citep{GoodBengCour16}) $C:\R^{c_1 \times h_1 \times w_1} \to \R^{c_2 \times h_2 \times w_2}$ with kernel sizes $(k_1,k_2)$, padding $p$ and stride $s$. The formula for the output dimensions are $h_2 = \left(\lfloor(h_1 - k_1 +2p)\rfloor/s+1\right)$ and $w_2 = \left(\lfloor(w_1 - k_2 +2p)\rfloor/s+1\right)$ and the number of patches on which the layer acts upon is $ h_2w_2$. For a given convolutional layer $C^{pq}$ in the network $f_W$, we denote $N^{pq}:=h_pw_p$. For a given convolutional layer $C^{pq}$ we define a matrix $W^{pq}$ that computes $\forall~ x \in \R^{c_q \times h_q \times w_q}:~ W^{pq} \vect(x) = \vect(C^{pq}(x))$. This is guaranteed since both the padding and convolution operations of the layer are linear. We also consider $V^{pq} \in \R^{c_{out} \times c_{in}k_1k_2}$ which is a matrix whose rows are the vectorized filters of $C^{pq}$.
\paragraph{Pooling Layers.} A pooling layer~\citep{Zhou1988ComputationOO} (see also~\citep{GoodBengCour16}) $C$ with kernel dimensions $(k_1,k_2)$ stride $s$ and padding $p$ takes an input $c_{in}\times h_1 \times w_1$ and computes an output with dimensions $h_2 = \left(\lfloor(h_1 - k_1 +2p)\rfloor/s+1\right)$ and $w_2 = \left(\lfloor(w_1 - k_2 +2p)\rfloor/s+1\right)$. In this paper we consider average pooling and max-pooling layers.
\paragraph{Batch normalization layers.} A batch normalization layer~\citep{pmlr-v37-ioffe15} is defined in the following manner,
\begin{small}
\begin{equation*}
\begin{aligned}
n_l(v_l(x;\tX);\tX) ~&=~ \gamma_l \cdot A_l(\tX) \odot (v_l(x;\tX) - b_l(\tX)) \\
&\quad + \one_{c\times h_l \times w_l}\cdot \beta_l \\
A_l(\tX) ~&=~ \left(\left(\sqrt{\Var_{\hat{x} \sim U[\tX]}(v_l(\hat{x};\tX)_{j})}\right)^{-1}\right)_{j} \\
b_l(\tX) ~&=~ \E_{\hat{x} \sim U[\tX]}[v_l(\hat{x};\tX)],
\end{aligned}
\end{equation*}
\end{small}
where $v_l(\hat{x};\tX)_{j}$ is the $j$'th coordinate of $v_l(\hat{x};\tX)$ for $j \in [c_l]\times [h_l]\times [w_l]$. In addition, $\gamma_l,\beta_l \in \R$ are trainable scale and shift parameters and $\one_{a_1 \times \dots \times a_k} \in \R^{a_1\times \dots \times a_k}$ is a tensor whose all of its cells are equal $1$. To simplify the analysis, we detach (also known as stop-grad~\citep{NEURIPS2020_f3ada80d}) the matrices $A_l(\tX)$ during backpropagation at any step during training. Namely, we treat $\frac{\bar{\partial} A_l(\tX)}{\partial W}=0$ in order to compute the training updates (without changing the forward passes). We note that detaching is a commonly used approach for training a network with normalization layers~\citep{Wiesler2014MeannormalizedSG,renorm,pmlr-v22-raiko12,https://doi.org/10.48550/arxiv.1410.7455,NEURIPS2019_2f4fe03d,10.1007/s10994-020-05929-w,https://doi.org/10.48550/arxiv.2202.07304}. In fact,~\cite{NEURIPS2019_2f4fe03d} showed that detaching the variance term of layer normalization is competitive and even improves the performance on various tasks.
\section{Theoretical Results}\label{sec:theory}
We begin our theoretical analysis with the observation that the rank of the gradients of the output of the network with respect to the parameters matrix of a given convolutional layer is upper bounded by the the number of patches $N^{pq}$ that the network acts upon.
\begin{lemma}\label{lem:rank}
Let $f_W$ be a convolutional neural network and let $C^{pq}$ be a given convolutional layer within $f_W$ with kernel size $(k_1,k_2)$, stride $s$, padding $p$, $c_{q}$ input channels and $c_{p}$ output channels. Let $V^{pq} \in \R^{c_{q} \times (c_{p} k_1k_2)}$ be the parameters matrix of $C^{pq}$. Then, the rank of $\frac{\partial f_W(x)}{\partial V^{pq}}$ is at most $N^{pq}$.
\end{lemma}
\begin{proof} Let $x \in \R^{c_1h_1w_1}$. We would like to show that $\frac{\partial f_W(x)}{\partial V^{pq}}$. We begin by writing the output of $f_W$ as a sum over paths that pass through $C^{pq}$ and paths that do not. We note that the output can be written as follows,
\begin{equation*}
f_W(x) ~=~ \sum_{l_1 \in \textnormal{pred}(l_0)} C^{l_0l_1} \circ v_{l_1}(x),
\end{equation*}
where $l_0 = L$ and $C^{ij} \circ z := C^{ij}(z)$. In addition, each layer $v_{l}$ can also be written as
\begin{equation*}
\vect(v_{l_1}(x)) ~=~ D_{l_1} \odot \sum_{l_2\in \textnormal{pred}(l_1)} C^{l_1l_2} \circ v_{l_2}(x),
\end{equation*}
where $D_l := D_l(x) := \sigma'(v_l(x)))$. Therefore, for any input $x$, we can write $f_W(x)$ as the sum of matrix multiplications along paths $\pi$ from $v_1$ to $v_{l_0}$. Specifically, we can write $f_W(x)$ as a follows
\begin{small}
\begin{equation*}
\begin{aligned}
f_W(x) ~&=~ \sum_{\pi \textnormal{ from $p$ to $l_0$}} C^{\pi_{T}\pi_{T-1}} \circ D_{\pi_{T-1}} \cdots D_{\pi_2} \odot C^{\pi_2\pi_1} \\
&\quad\quad\quad\quad\quad\quad \circ D_{\pi_1} \odot C^{pq} \circ v_q(x) \\
&\quad\quad + \sum_{\substack{\pi \textnormal{ from $1$ to $l_0$} \\ (p,q) \notin \pi }} C^{\pi_{T}\pi_{T-1}} \circ D_{\pi_{T-1}} \odot C^{\pi_{T-1}\pi_{T-2}} \\
&\quad\quad\quad\quad\quad\quad\quad \cdots D_{\pi_2} \odot C^{\pi_2 \pi_1} \circ x, \\
% ~&=~ \left[\sum_{\pi \textnormal{ from $p$ to $l_0$}} C^{\pi_{T}\pi_{T-1}} \circ D_{\pi_{T-1}} \cdots C^{\pi_2\pi_1} \circ D_{\pi_1} \right] \\
% &\quad\quad\quad\quad\quad\quad\quad \odot C^{pq} \circ v_q(x) \\
% &\quad\quad + \sum_{\substack{\pi \textnormal{ from $1$ to $l_0$} \\ (p,q) \notin \pi }} C^{\pi_{T}\pi_{T-1}} \circ D_{\pi_{T-1}} \odot C^{\pi_{T-1}\pi_{T-2}} \\
% &\quad\quad\quad\quad\quad\quad\quad \cdots D_{\pi_2} \odot C^{\pi_2 \pi_1} \circ x \\
% ~&=:~ u_q(x)^{\top} \cdot C^{pq} \circ (v_q(x)) \\
% &\quad\quad + \sum_{\substack{\pi \textnormal{ from $1$ to $l_0$} \\ (p,q) \notin \pi }} W^{\pi_{T}\pi_{T-1}} \cdot D_{\pi_{T-1}} \cdot W^{\pi_{T-1}\pi_{T-2}} \\
% &\quad\quad\quad\quad\quad\quad\quad \cdots D_{\pi_2} \cdot W^{\pi_2 \pi_1}\cdot \vect(x) \\
\end{aligned}
\end{equation*}
\end{small}
where $T := T(\pi)$ denotes the length of the path $\pi$. We denote the first sum by $A$ and the second one by $B$. We note that for any $x\in \R^{c_1h_1w_1}$, with measure $1$ over $W$, the matrices $\{D_{i}(x)\}^{L-1}_{i=1}$ are constant in the neighborhood of $W$. Furthermore, $C^{pq}$ does not appear in the multiplications along the paths $\pi$ from $1$ to $l_0$ that exclude $(p,q)$. Therefore, we conclude that $\frac{\partial B}{\partial V^{pq}}=0$.
As a next step we would like to analyze the rank of the gradient of $A$. For this purpose, we rewrite the convolutional layers and the multiplications by the matrices $D_{l}$ as linear operations.
\paragraph{Representing $C^{pq}$.} We begin by representing the layer $C^{pq}$ as a linear transformation of its input with $N^{pq}$ blocks of $V^{pq}$. Let $C^{pq}$ be a $c_{out} \times c_{in} \times k_1 \times k_2$ convolutional layer in $f_W$ that operates on $c_{in} \times k_1 \times k_2$ patches, with stride $s$ and padding $p$ on an input $x \in \R^{c_{in} \times h \times w}$. We let $N^{pq}$ be the number of input patches of the convolutional layer on its input. We begin by defining a flattening of a given $3$-dimensional tensor input $x \in \R^{c_{in} \times h\times w}$ into a vector $\vect^{pq}(x) \in \R^{N^{pq} c_{in} k_1 k_2}$ depending on the convolutional layer $C^{pq}$. We pad $x$ with $p$ rows and columns of zeros and vectorize each one of its patches (of dimensions $c_{in}\times k_1\times k_2$) that the convolutional layer is acting on (potentially overlapping) and concatenate them. Furthermore, we define a matrix $V^{pq} \in \R^{c_{out} \times (c_{in}k_1k_2)}$, where each one of its rows is the vectorization of each one of $C^{pq}$'s $c_{in} \times k_1 \times k_2$ filters. We can write the vectorized output of the convolutional layer as $U^{pq}\cdot \vect^{pq}(x)$, where
\begin{equation}\label{eq:Upq}
U^{pq} := \renewcommand{\arraystretch}{1.2}
\left(
\begin{array}{ c c c c c c }
\multicolumn{1}{c}{} & & 0 & 0 & 0 & 0 \\
\multicolumn{2}{c}{\raisebox{.6\normalbaselineskip}[0pt][0pt]{$V^{pq}$}} & 0 & 0 & 0 & 0 \\
0 & 0 & & & 0 & 0 \\
0 & 0 & \multicolumn{2}{c}{\raisebox{.6\normalbaselineskip}[0pt][0pt]{$\ddots$}} & 0 & 0 \\
0 & 0 & 0 & 0 & & \multicolumn{1}{c}{} \\
0 & 0 & 0 & 0 & \multicolumn{2}{c}{\raisebox{.6\normalbaselineskip}[0pt][0pt]{$V^{pq}$}}
\end{array}
\right)
\end{equation}
is a $(N^{pq} c_{p}) \times (N^{pq} c_{q} k_1k_2)$ matrix with $N^{pq}$ copies of $V^{pq}$. We note that this is a non-standard representation of the convolutional layer's operation as a linear transformation. Typically, we write the convolutional layer as a linear transformation $W^{pq}$ acting on the vectorized version $\vect(x) \in \R^{c_{in}k_1k_2}$ of its input $x$. Since $\vect^{pq}(x)$ consists of the same variables as $\vect(x)$ with potential duplicates, there is a linear transformation that translates $\vect(x)$ into $\vect^{pq}(x)$. Therefore, we can simply write $W^{pq} \cdot \vect(x) = U^{pq} \cdot \vect^{pq}(x)$.
\paragraph{Representing convolutional layers.} Except of $C^{pq}$, we represent each one of the network's convolutional layers $C^{ij}$ as the simple linear transformations $W^{ij}$. As mentioned earlier, for any input $x\in \R^{c_j \times h_j \times w_j}$ of $C^{ij}$, we can write $\vect(C^{ij}(x)) = W^{ij}(\vect(x))$.
\paragraph{Representing pooling layers.} An average pooling layer $C^{ij}$ can be represented as a (non-trainable) linear transformation $W^{ij}$ of its input. A max-pooling layer can be written as a composition of ReLU activations and a (non-trainable) linear transformation, since $\max(x,y) = \sigma(x-y) + y$. Therefore, without loss of generality we assume that the network is free of max-pooling layers.
Finally, we note that $\vect(C^{ij} \circ z) = U^{ij} \cdot \vect^{ij}(z) = W^{ij} \cdot \vect(z)$, $\vect(D_l \odot z) = P_l \cdot \vect(z)$ for $P_{l} := \textnormal{diag}(\vect(D_{l}))$. Therefore, by replacing $C^{pq}$ with $U^{pq}$, $C^{ij}$ with $W^{ij}$ (for all $(i,j)\neq (p,q)$) and $D_l$ with $P_l$, we have
\begin{equation*}
\begin{aligned}
A ~&=~ \sum_{\pi \textnormal{ from $p$ to $l_0$}} W^{\pi_{T}\pi_{T-1}} \cdot P_{\pi_{T-1}} \cdots P_{\pi_2} \cdot W^{\pi_2\pi_1} \\
&\quad\quad\quad\quad\quad\quad\quad\cdot P_{\pi_1} \cdot U^{pq} \cdot \vect^{pq}(v_q(x)) \\
~&=:~ a(x) \cdot U^{pq} \cdot b(x),
\end{aligned}
\end{equation*}
where $a(x) ~:=~ \sum_{\pi \textnormal{ from $p$ to $l_0$}} W^{\pi_{T}\pi_{T-1}} \cdot P_{\pi_{T-1}} \cdots P_{\pi_2} \cdot W^{\pi_2\pi_1} \cdot P_{\pi_1}$ and $b(x) := \vect^{pq}(v_q(x))$.
We note that with measure $1$, the matrices $\{P_l\}^{L-1}_{l=1}$ are constant in the neighborhood of $W$. In addition, $a(x)$ and $b(x)$ are computed as multiplications of matrices $W^{ij}$ and $P_l$ excluding $(i,j)=(p,q)$. Therefore, with measure $1$ over the selection of $W$, the gradients of $a(x)$ and $b(x)$ with respect to $V^{pq}$ are $0$. Furthermore, due to \eqref{eq:Upq} and the definition of $U^{pq}$, we can write
\begin{equation*}
\begin{aligned}
u_q(x)^{\top} \cdot W^{pq} \cdot \vect(v_q(x)) ~&=~ u_q(x)^{\top} \cdot U^{pq} \cdot \vect^{pq}(v_q(x)) \\
~&=~ \sum^{N^{pq}}_{i=1} a_i(x)^{\top} \cdot V^{pq} \cdot b_i(x),
\end{aligned}
\end{equation*}
where $a_i(x)$ and $b_i(x)$ are the slices of $a(x)$ and $a(x)$ that are multiplied by the $i$'th $V^{pq}$ block in $U^{pq}$. Since with measure $1$ over the selection of $W$ the differentials of $a_i(x)$ and $b_i(x)$ with respect to $V^{pq}$ are $0$, we have,
\begin{equation}
\frac{\partial u_q(x)^{\top} \cdot U^{pq} \cdot \vect^{pq}(v_q(x))}{\partial V^{pq}} ~=~ \sum^{N^{pq}}_{i=1} a_i(x) \cdot b_i(x)^{\top}.
\end{equation}
Therefore, we conclude that, with measure $1$ over the selection of $W$, we have $\frac{\partial f_W(x)}{\partial V^{pq}} = \sum^{N^{pq}}_{i=1} a_i(x) \cdot b_i(x)^{\top}$ which is a matrix of rank $\leq N^{pq}$.
\end{proof}
\begin{restatable}{lemma}{lowrankSec}\label{lem:sec_rank}
Let $\tf_W$ be a convolutional neural network with batch-normalization. Let $C^{pq}$ be a given convolutional layer within $\tf_W$ with kernel size $(k_1,k_2)$, stride $s$, padding $p$, $c_{q}$ input channels and $c_{p}$ output channels. Let $V^{pq} \in \R^{c_{q} \times (c_{p} k_1k_2)}$ be the parameters matrix of $C^{pq}$. Let $\tX = \{x_i\}^{B}_{i=1}$ be a batch of samples. Then, the ranks of $\frac{\partial f_W(x;\tX)}{\partial V^{pq}}$ and $\sum^{B}_{i=1}\alpha_i \frac{\partial f_W(x_i;\tX)}{\partial V^{pq}}$ are at most $BN^{pq}$.
\end{restatable}
\subsection{Low-Rank Bias in Convolutional Layers}
In this section we study the low-rank bias of convolutional layers in a neural network.
\begin{restatable}{theorem}{batches}\label{thm:batches}
Let $\|\cdot \|$ be any matrix norm and $\ell$ any differentiable loss function. Let $f_W(x)$ be a ReLU neural network and $C^{pq}$ be a given matrix within $f_W$ and let $B^* \in [m]$. Then,
\begin{equation*}
\min_{\substack{V \in R^{d_p \times d_q}:\\ \rank(V) \leq kN^{pq}B^*}} \|V^{pq} - V\| ~\leq~ \frac{1}{2\lambda}\min_{\substack{\tS \subset S: \\ \vert \tS\vert = B^*}} \|\bar{\nabla}_{W^{pq}} L^{\lambda}_{\tS}(f_W)\|
\end{equation*}
\end{restatable}
\begin{proof}
Let $\tS = \{(x_{ij},y_{ij})\}^{B^*}_{j=1} \subset S$ be a batch of $B^*$ samples and $\tX = \{x_{ij}\}^{B^*}_{j=1}$ be the corresponding unlabeled batch. By the chain rule, we can write the gradient of the loss function as follows
\begin{equation*}
\begin{aligned}
\frac{\bar{\partial} L_{\tS}(\tf_W)}{\partial W^{pq}} ~=~ & \frac{1}{B^*}\sum^{B^*}_{j=1}\frac{\partial \ell(f_W(x_{ij};\tX),y_{ij})}{\partial f_W(x_{ij};\tX)} \cdot \frac{\bar{\partial} f_W(x_{ij};\tX)}{\partial W^{pq}} \\
&+ 2\lambda W^{pq}.
\end{aligned}
\end{equation*}
According to Lems.~\ref{lem:rank} and~\ref{lem:sec_rank}, the matrix $U_{\tS} := -\frac{1}{2\lambda B^*}\sum^{B^*}_{j=1}\frac{\ell(f_W(x_{ij};\tX),y_{ij})}{\partial f_W(x_{ij};\tX)} \cdot \frac{\partial f_W(x_{ij};\tX)}{\partial W^{pq}}$ is of rank $\leq B^*N^{pq}$. Therefore, we obtain the following
\begin{equation*}
\begin{aligned}
&\min_{V:~ \rank(V) \leq B^*N^{pq}}\|W^{pq} - V\| \\
~&\leq~ \min_{\tS \subset S:~\vert \tS\vert = B^*}\|W^{pq} - U_{\tS}\| \\
~&=~ \frac{1}{2\lambda}\min_{\tS \subset S:~\vert \tS\vert = B^*}\left\|\frac{\bar{\partial} L^{\lambda}_{\tS}(\tf_W)}{\partial W^{pq}} \right\|,
\end{aligned}
\end{equation*}
which completes the proof of this theorem.
\end{proof}
In this section, we describe our main theoretical results. In Sec.~\ref{sec:convergence}, we prove that SGD noise is inevitable when applying weight decay. In Sec.~\ref{sec:lowrank} we provide a brief analysis showing that mini-batch SGD implicitly learns neural networks with low-rank matrices. Proofs are provided in the supplementary material. %Appendix~\ref{app:proofs}.
\subsection{Degeneracy and the Origin of SGD Noise}\label{sec:convergence}
We use the term `SGD noise' to refer to the inherent inability of SGD to converge to a stationary solution for the weight matrices that is the same across different mini-batches. In this section, we characterize the convergence of mini-batch SGD. Our results are essentially impossibility results that show that convergence of SGD takes place only when the network is the zero function. This implies that asymptotic noise is generically unavoidable (when training is successful). For simplicity, we focus on normalization-free univariate neural networks $f_W:\R^{n}\to \R$ and assume that $\forall~i\in [m]:~x_i\neq 0$.
{\color{red}Why is it the correct definition?}
We start by spelling out the conditions for what we call in this paper convergence of SGD. We note that at convergence, we have $\nabla_{W^{pq}} L_{\tS}(f_W)=0$ across all pairs $(v_p,v_q) \in E_T$ and mini-batches $\tS \subset S$ of size $B < m$. Specifically, we we can write,
\begin{equation}\label{eq:1}
\begin{aligned}
0 ~&=~ \nabla_{W^{pq}} L^{\lambda}_{\tS}(f_W) \\
~&=~ \frac{1}{B} \sum_{(x,y)\in \tS} \frac{\partial \ell(f_W(x),y)}{\partial f_W(x)} \cdot \frac{\partial f_W(x)}{\partial W^{pq}} + 2\lambda W^{pq}.
\end{aligned}
\end{equation}
Suppose we have two batches $\tS_1,\tS_2 \subset S$ of size $B$ that differ by only one sample. We denote the unique sample of each batch by $(x_i,y_i)$ and $(x_j,y_j)$ respectively. We notice that,
\begin{equation}\label{eq:substract}
\begin{aligned}
0 ~&=~ \nabla_{W^{pq}} L^{\lambda}_{\tS_1}(f_W) - \nabla_{W^{pq}} L^{\lambda}_{\tS_2}(f_W) \\
~&=~ \frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \cdot \frac{\partial f_W(x_{i})}{\partial W^{pq}} - \frac{\partial \ell(f_W(x_{j}),y_{j})}{\partial f_W(x_{j})} \cdot \frac{\partial f_W(x_{j})}{\partial W^{pq}}.
\end{aligned}
\end{equation}
Therefore, we conclude that for all $i,j\in [m]$,
\begin{equation}\label{eq:ij}
\begin{aligned}
M^{pq} ~&=~ \frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \cdot \frac{\partial f_W(x_{i})}{\partial W^{pq}}\\
~&=~ \frac{\partial \ell(f_W(x_{j}),y_{j})}{\partial f_W(x_{j})} \cdot \frac{\partial f_W(x_{j})}{\partial W^{pq}}.
\end{aligned}
\end{equation}
Hence, for all $(v_p,v_q) \in E_T$ and $i \in [m]$,
\begin{equation}\label{eq:batchone}
\frac{\partial \ell(f_W(x_i),y_i)}{\partial f_W(x_i)} \cdot \frac{\partial f_W(x_i)}{\partial W^{pq}} + 2\lambda W^{pq} ~=~ M^{pq} + 2\lambda W^{pq} ~=~ 0.
\end{equation}
Therefore, unless $\lambda=0$ or $\forall~(v_p,v_q)\in E_T:~W^{pq}=0$, we conclude that $\frac{\partial \ell(f_W(x_i),y_i)}{\partial f_W(x_i)}\neq 0$ for all $i\in [m]$. In this case, we also obtain that $\{\frac{\partial f_W(x_i)}{\partial \textnormal{vec}(W^{pq})}\}^{m}_{i=1}$ are collinear vectors by Eq.~\ref{eq:ij}.
Therefore, any convergence point of training a neural network using mini-batch SGD along with weight decay is highly degenerate and does not perfectly fit any one of the training labels. To better understand the essence of this degeneracy, we provide the following proposition, which is specialized for ReLU networks.
\begin{restatable}[$\lambda > 0$]{proposition}{collinear}\label{prop:propA}
Let $\ell(a,b)$ be a differentiable loss function, $\lambda > 0$, and let $f_W(x)$ be a ReLU neural network, where $\{(v_p,v_1)\in E \mid v_p \in V\} \subset E_T$. Let $W$ be a convergence point of mini-batch SGD for minimizing $L^{\lambda}_S(f_W)$ (see Eq.~\ref{eq:1}). Then, either $f_W \equiv 0$ or $\{x_i\}^{m}_{i=1}$ are collinear vectors.
\end{restatable}
The above proposition shows that unless the training samples lie on a one-dimensional linear space, any convergence point of SGD corresponds to the zero function. Since both conditions are unrealistic, we conclude that convergence is impossible in any practical scenario.
In the following proposition, we look at the convergence of SGD when training without weight decay.
\begin{restatable}[$\lambda = 0$]{proposition}{zerolambda}\label{prop:propB}
Let $\ell(a,b)$ be a differentiable loss function, $\lambda = 0$, and let $f_W(x)$ be a ReLU neural network, where $\{(v_p,v_1)\in E \mid v_p \in V\} \subset E_T$. Let $W$ be a convergence point of mini-batch SGD for minimizing $L^{\lambda}_S(f_W)$. Assume that $\{x_i\}^{m}_{i=1}$ are not collinear vectors. Then, for any $i\in [m]$, for which $\frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \neq 0$, we have, $f_W(x_i)=0$. In particular, if $\ell$ is convex and $\forall~b\in \R~\nexists~a^* \in \R:~ \ell(a^*,b)=\inf_a \ell(a,b)$, then, $\forall~i \in [m]:~f_W(x_i)=0$.
\end{restatable}
The above proposition shows that at convergence, the neural network outputs zero for any training sample that it does not perfectly fit (i.e., $\frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \neq 0$). In particular, if $\ell$ is an exponential-type loss function (e.g., binary cross-entropy, logistic loss and exponential loss), $\ell(a,b)$ has no minima $a$ for any $b \in \R$, and therefore, the only possible convergence points are ones for which $f_W(x_i)=0$ across all $i \in [m]$. While theoretically convergence to a non-zero function is not guaranteed, in practice training without weight decay can still fall into the regime of `almost convergence', in which $\max_{i \in [m]} \big\vert \frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \big\vert$ is tiny and the training steps $- \eta \cdot \frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \cdot \frac{\partial f_W(x_i)}{\partial W^{pq}}$ are very small as a result. As a general tendency, when increasing the size of the network (i.e., the width or the depth), we should expect the term $\max_{i \in [m]} \big\vert \frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} \big\vert$ to tend to zero (if the network is trained successfully). Furthermore, with the MSE loss we may actually have convergence that happens when the network perfectly fits the training labels, i.e., $\forall~i\in [m]:~f_W(x_i)=y_i$.% (in this case $\forall~i \in [m]:~\frac{\partial \ell(f_W(x_{i}),y_{i})}{\partial f_W(x_{i})} = 0$).
% To summarize, when training a network for cross-entropy minimization with mini-batch SGD, convergence is not guaranteed (with or without regularization). That is because the convergence points are highly degenerate. On the other hand, when training a network to minimize the squared loss without regularization, any point $W$ for which $\forall~i \in [m]:~f_W(x_i)=y_i$ is a convergence point, which is achievable with overparameterized networks.
It is worth noting that our analysis critically depends on using SGD instead of GD during optimization. While many papers analyze the training dynamics and critical points of GD, the solutions of SGD are highly degenerate and will, in general, look different from the solutions of GD. For instance, it is impossible to derive Eq.~\ref{eq:ij} without relying on optimization with SGD, and without weight decay, we could not argue that $W^{pq}$ is proportional to the low-rank matrix $M^{pq}$. Finally, the analysis above is independent of the batch size as long as it is strictly smaller than the full dataset's size. %As follows from Eq.~\ref{eq:batchone} the convergence points of mini-batch SGD are equivalent to the convergence points of SGD with batch size $1$.
\section{Experiments}\label{sec:experiments}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c@{}c}
\rotatebox[origin=c]{90}{{\tiny MNIST, CE}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/mnist_dist_ce_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/mnist_acc_ce_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/mnist_loss_ce_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/mnist_rank_ce_mlp_2000_lr_01.pdf}}
\\
\rotatebox[origin=c]{90}{{\tiny MNIST, MSE}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/mnist_dist_mse_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/mnist_acc_mse_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/mnist_loss_mse_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/mnist_rank_mse_mlp_2000_lr_01.pdf}}
\\
\rotatebox[origin=c]{90}{{\tiny CIFAR10, CE}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/cifar10_dist_ce_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/cifar10_acc_ce_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/cifar10_loss_ce_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/ce/cifar10_rank_ce_mlp_2000_lr_01.pdf}}
\\
\rotatebox[origin=c]{90}{{\tiny CIFAR10, MSE}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/cifar10_dist_mse_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/cifar10_acc_mse_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/cifar10_loss_mse_mlp_2000_lr_01.pdf}} &
\raisebox{-0.5\height}{\includegraphics[width=0.24\linewidth]{plots/convergence/mse/cifar10_rank_mse_mlp_2000_lr_01.pdf}}
\\
& {\small (a)} & {\small (b)} & {\small (c)} & {\small (d)}
\end{tabular}
\caption{{\bf Convergence of MLP-5-2000 trained on MNIST and CIFAR10 with CE/~MSE loss.} In {\bf (a)} we plot the averaged distance between the weight matrices at epoch $t$ and epoch $t+1$, captured by, $\frac{1}{\vert E_T\vert}\sum_{(p,q) \in E_{T}} \|W^{pq}_{t+1} - W^{pq}_{t}\|$. In {\bf (b)} we plot the train accuracy rates, in {\bf (c)} we plot the averaged train loss and in {\bf (d)} we plot the average rank across the trainable matrices.}
\label{fig:convergence}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/rank_batch_size_mlp_100_10_lr_0001.pdf}&
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/rank_batch_size_mlp_100_10_lr_001.pdf}&
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/rank_batch_size_mlp_100_10_lr_01.pdf}&
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/rank_batch_size_mlp_100_10_lr_05.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/acc_batch_size_mlp_100_10_lr_0001.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/acc_batch_size_mlp_100_10_lr_001.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/acc_batch_size_mlp_100_10_lr_01.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/batch_size/acc_batch_size_mlp_100_10_lr_05.pdf}
\\
{\small $\eta=0.001$} & {\small $\eta=0.01$} & {\small $\eta=0.1$} & {\small $\eta=0.5$}
\end{tabular}
\caption{{\bf Average rank of MLP-BN-10-100 trained on CIFAR10 with varying batch sizes.} Each line stands for the results of a different batch size. In the (top) row we plot the average rank across layers and in the (bottom) row we plot the train and test accuracy rates for each setting.}
\label{fig:bs}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/rank_wd_mlp_100_10_bs_16.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/rank_wd_mlp_100_10_bs_32.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/rank_wd_mlp_100_10_bs_64.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/rank_wd_mlp_100_10_bs_128.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/acc_wd_mlp_100_10_bs_16.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/acc_wd_mlp_100_10_bs_32.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/acc_wd_mlp_100_10_bs_64.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp/weight_decay/acc_wd_mlp_100_10_bs_128.pdf}
\\
{\small $B=16$} & {\small $B=32$} & {\small $B=64$} & {\small $B=128$}
\end{tabular}
\caption{{\bf Average rank of MLP-BN-10-100 trained on CIFAR10 with varying weight decay.} See Fig.~\ref{fig:bs} for details.}
\label{fig:wd}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/rank_batch_size_resnet_100_5_lr_00125.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/rank_batch_size_resnet_100_5_lr_0025.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/rank_batch_size_resnet_100_5_lr_005.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/rank_batch_size_resnet_100_5_lr_01.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/acc_batch_size_resnet_100_5_lr_00125.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/acc_batch_size_resnet_100_5_lr_0025.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/acc_batch_size_resnet_100_5_lr_005.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet/batch_size/acc_batch_size_resnet_100_5_lr_01.pdf}
\\
{\small $\eta=0.0125$} & {\small $\eta=0.025$} & {\small $\eta=0.05$} & {\small $\eta=0.1$}
\end{tabular}
\caption{{\bf Average rank of RES-BN-5-500 trained on CIFAR10 with varying batch sizes.} See Fig.~\ref{fig:bs} for details.}
\label{fig:bs_resnet}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/rank_batch_size_mlp_nobn_500_5_lr_00125.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/rank_batch_size_mlp_nobn_500_5_lr_0025.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/rank_batch_size_mlp_nobn_500_5_lr_005.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/rank_batch_size_mlp_nobn_500_5_lr_01.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/acc_batch_size_mlp_nobn_500_5_lr_00125.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/acc_batch_size_mlp_nobn_500_5_lr_0025.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/acc_batch_size_mlp_nobn_500_5_lr_005.pdf} &
\includegraphics[width=0.25\linewidth]{plots/mlp_no_bn/learning_rate/acc_batch_size_mlp_nobn_500_5_lr_01.pdf}
\\
{\small $\eta=0.0125$} & {\small $\eta=0.025$} & {\small $\eta=0.05$} & {\small $\eta=0.1$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of MLP-5-500 trained on CIFAR10 with varying batch sizes.} See Fig.~\ref{fig:bs} for details.}
\label{fig:bs_mlp_nobn}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/rank_resnet_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/rank_resnet_bn_lr_5_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/rank_resnet_bn_lr_8_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/rank_resnet_bn_lr_10_batch_size.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/acc_resnet_bn_lr_1_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/acc_resnet_bn_lr_5_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/acc_resnet_bn_lr_8_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/batch_size/acc_resnet_bn_lr_10_batch_size.pdf}
\\
{\small $\eta=0.1$} & {\small $\eta=0.5$} & {\small $\eta=0.8$} & {\small $\eta=1.0$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of {\color{blue}RES18} trained on CIFAR10 with varying batch sizes.} See Fig.~\ref{fig:bs} for details.}
\label{fig:bs_mlp_nobn}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/rank_resnet_bn_bs_8_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/rank_resnet_bn_bs_16_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/rank_resnet_bn_bs_32_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/rank_resnet_bn_bs_64_lr_1_weight_decay.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/acc_resnet_bn_bs_8_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/acc_resnet_bn_bs_16_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/acc_resnet_bn_bs_32_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/weight_decay/acc_resnet_bn_bs_64_lr_1_weight_decay.pdf}
\\
{\small $B=8$} & {\small $B=16$} & {\small $B=32$} & {\small $B=64$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of {\color{blue}RES18} trained on CIFAR10 with varying weight decay.} See Fig.~\ref{fig:bs} for details.}
\label{fig:bs_mlp_nobn}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/rank_resnet_bn_batch_size_4_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/rank_resnet_bn_batch_size_8_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/rank_resnet_bn_batch_size_16_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/rank_resnet_bn_batch_size_32_learning_rate.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/acc_resnet_bn_batch_size_4_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/acc_resnet_bn_batch_size_8_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/acc_resnet_bn_batch_size_16_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/resnet18/learning_rate/acc_resnet_bn_batch_size_32_learning_rate.pdf}
\\
{\small $B=4$} & {\small $B=8$} & {\small $B=16$} & {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of {\color{blue}RES18} trained on CIFAR10 with varying learning rate.} See Fig.~\ref{fig:bs} for details.}
\label{fig:bs_mlp_nobn}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/rank_vgg_bn_lr_001_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/rank_vgg_bn_lr_005_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/rank_vgg_bn_lr_01_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/rank_vgg_bn_lr_02_batch_size.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/acc_vgg_bn_lr_001_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/acc_vgg_bn_lr_005_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/acc_vgg_bn_lr_01_batch_size.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/batch_size/acc_vgg_bn_lr_02_batch_size.pdf}
\\
{\small $\eta=0.001$} & {\small $\eta=0.005$} & {\small $\eta=0.01$} & {\small $\eta=0.02$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of {\color{blue}VGG16} trained on CIFAR10 with varying batch sizes.} See Fig.~\ref{fig:bs} for details.}
\label{fig:bs_mlp_nobn}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/rank_vgg_bn_batch_size_16_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/rank_vgg_bn_batch_size_32_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/rank_vgg_bn_batch_size_64_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/rank_vgg_bn_batch_size_128_lr_1_weight_decay.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/acc_vgg_bn_batch_size_16_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/acc_vgg_bn_batch_size_32_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/acc_vgg_bn_batch_size_64_lr_1_weight_decay.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/weight_decay/acc_vgg_bn_batch_size_128_lr_1_weight_decay.pdf}
\\
{\small $B=16$} & {\small $B=32$} & {\small $B=64$} & {\small $B=128$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of {\color{blue}VGG16} trained on CIFAR10 with varying weight decay.} See Fig.~\ref{fig:bs} for details.}
\label{fig:bs_mlp_nobn}
\end{figure*}
\begin{figure*}[t]
\centering
\begin{tabular}{c@{}c@{}c@{}c@{}c}
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/rank_vgg_bn_batch_size_4_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/rank_vgg_bn_batch_size_8_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/rank_vgg_bn_batch_size_16_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/rank_vgg_bn_batch_size_32_learning_rate.pdf}
\\
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/acc_vgg_bn_batch_size_4_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/acc_vgg_bn_batch_size_8_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/acc_vgg_bn_batch_size_16_learning_rate.pdf} &
\includegraphics[width=0.25\linewidth]{plots/vgg16/learning_rate/acc_vgg_bn_batch_size_32_learning_rate.pdf}
\\
{\small $B=4$} & {\small $B=8$} & {\small $B=16$} & {\small $B=32$}
\end{tabular}
\caption{{\bf Average ranks and accuracy rates of {\color{blue}VGG16} trained on CIFAR10 with varying learning rate.} See Fig.~\ref{fig:bs} for details.}
\label{fig:bs_mlp_nobn}
\end{figure*}
In this section we experimentally analyze the convergence (and non-convergence) of training neural networks with SGD. In addition, we empirically study the implicit bias towards rank minimization in deep ReLU networks. Throughout the experiments we extensively vary different hyperparameters (e.g., the learning rate, weight decay, and the batch size) and study their effect on the rank of the various matrices in the network\footnote{The plots are best viewed when zooming into the pictures.}. For each run we used one Tesla-k80 GPU for 5-20 hours.
\subsection{Setup}
\paragraph{ Evaluation process.} Following the problem setup, we consider $k$-class classification problems and train a multilayered neural network $f_W = e \circ h = e \circ g^L \circ \dots \circ g^1 :\R^n \to \R^k$ on some balanced training dataset $S$. The model is trained using CE/~MSE loss minimization between its logits and the one-hot encodings of the labels. Here, $g^1,\dots,g^L$ are the various hidden layers of the network (e.g., fully-connected layers or residual blocks), where $e$ is its top linear layer. After each epoch, we compute the averaged rank across the network's weight matrices and its train and test accuracy rates. To estimate the rank of a given matrix, we normalize it and count how many of its singular values are $\notin [-\expnumber{1}{-4},\expnumber{1}{-4}]$. Training was done on the MNIST and CIFAR10 datasets. %During training, we used random cropping, random horizontal flips and random rotations (by $15k$ degrees for $k$ uniformly sampled from $[24]$) and standardize the data. %For all datasets we standardized the data.
\paragraph{Architectures.} We consider two types of architectures. %The first architecture is a convolutional network (ConvNet), beginning with a stack of a $2\times 2$ convolutional layer with stride $2$, batch normalization, a convolution of the same structure, batch normalization and ReLU. Following that we have a set of $L$ stacks of $3\times 3$ convolutional layer with stride $1$ and padding $1$, batch normalization and ReLU. The output tensors of these layers shares the same shape as their input's shape. The number of channels in each convolutional layer is $H$. The last layer is linear. The network is denoted by Conv-$L$-$H$.
The first architecture is an MLP. Our MLPs, denoted by MLP-BN-$L$-$H$ consist of $L$ hidden layers, where each layer contains a linear layer of width $H$, followed by batch normalization and ReLU activations. On top of that, we compose a linear output layer. The second architecture, denoted by RES-BN-$L$-$H$, consist of a linear layer of width $H$, followed by $L$ residual blocks, each computing a function of the form $z+\sigma(n_2(W_2 \sigma(n_1(W_1z))))$, where $W_1,W_2 \in \R^{H\times H}$, $n_1,n_2$ are batch normalization layers and $\sigma$ is the element-wise ReLU function. Again, the network is ended with a linear output layer. We denote by MLP-$L$-$H$ and RES-$L$-$H$ the same architectures without applying batch normalization.
%Unless mentioned otherwise,
% {\bf Datasets.\enspace} Throughout the experiments, we consider three datasets: (i) MNIST; (ii) CIFAR10; (iii) STL-10. For CIFAR10 we used random cropping, random horizontal flips and random rotations (by $15k$ degrees for $k$ uniformly sampled from $[24]$). For all datasets we standardized the data. {\color{red}Update}
\subsection{Experiments on SGD Noise}\label{sec:sgd_noise_exp}
In Sec.~\ref{sec:convergence} we showed that training a neural network with SGD along with weight decay, cannot converge to a non-zero function. To verify this prediction, we trained MLP-5-2000 networks with varying degrees of weight decay $\lambda \in \{0, \expnumber{1}{-6}, \expnumber{1}{-5}, \expnumber{1}{-4}, \expnumber{1}{-3}\}$ and investigated their convergence. Each model was trained using SGD with batch size 128 and learning rate 0.1 for 2000 epochs (without learning rate scheduling or momentum). We trained the models once with the cross-entropy loss and once with the squared loss.
To evaluate the convergence of the networks, we use the following quantity: $d(W_{t+1},W_{t}) := \frac{1}{\vert E_T\vert} \sum_{(p,q)\in E_T} \|W^{pq}_{t+1} - W^{pq}_{t}\|$, where $\{W^{pq}_t\}_{(p,q)\in E_T}$ are the various trainable matrices in the network at epoch $t$. This quantity measures the averaged distance between the network's matrices at epochs $t$ and $t+1$. Convergence is possible only if $\lim_{t\to \infty} d(W_{t+1},W_{t}) = 0$.
In Fig.~\ref{fig:convergence} we monitor (a) $d(W_{t+1},W_{t})$, (b) the train accuracy rates, (c) the train losses and (d) the averaged rank of the trainable matrices. As predicted in Prop.~\ref{prop:propA}, when training with $\lambda>0$, $W_t$ either converges to zero and $f_{W_t}$ to the zero function (e.g., see the results of training on CIFAR10 with the cross-entropy loss and $\lambda=\expnumber{1}{-4}, \expnumber{1}{-3}$) or $W_t$ does not converge (i.e., $d(W_{t+1},W_{t})$ does not tend to zero). On the other hand, the term $d(W_{t+1},W_{t})$ is smaller by orders of magnitude when using $\lambda=0$ in comparison to using $\lambda>0$ (except for cases when the selection of $\lambda > 0$ leads to $W^{pq}_t \to 0$). Interestingly, even though for certain values of $\lambda > 0$, the training loss and accuracy converged, the network's parameters do not converge.
Finally, when training for cross-entropy loss minimization without weight decay we encounter the `almost convergence' regime discussed in Sec.~\ref{sec:convergence}. Namely, even though perfect convergence is impossible (see Prop.~\ref{prop:propB}) for a fixed sized neural network, the term $d(W_{t+1},W_{t})$ may become as small as we wish by increasing the size of the neural network. Therefore, since the MLP-5-2000 is relatively large (compared to the dataset's size), we may inaccurately get the impression that $d(W_{t+1},W_{t})$ tends to zero.
\subsection{Experiments on Rank Minimization}\label{sec:rank_exp}
\paragraph{Rank, batch size and the learning rate.} We train instances of MLP-BN-10-100, RES-BN-5-500 and MLP-5-500 using SGD with different batch sizes $B$ and initial learning rates $\eta$. The learning rate is decayed by a factor of $0.1$ three times at epochs 60, 100, and 200 and training is stopped after 500 epoch. We trained the models with weight decay $\lambda=\expnumber{5}{-4}$ and momentum $0.9$.
As can be seen in Figs.~\ref{fig:bs},~\ref{fig:bs_resnet} and~\ref{fig:bs_mlp_nobn}, by decreasing the batch size, we essentially strengthen the low-rank constraint over the network's matrices, which eventually leads to matrices of lower ranks. This is consistent with the prediction made in Sec.~\ref{sec:lowrank} that we expect to learn matrices of smaller ranks when training the network with smaller batch sizes. In addition, we also notice that when increasing the learning rate, the rank minimization constraint strengthens as well, which leads to lower ranks and generally worse training performances. Therefore, to properly train a neural network with smaller batch sizes, it is necessary to decrease the learning rate accordingly.
\paragraph{Rank minimization and weight decay.} To study the effect of the weight decay on rank minimization, we trained instances of MLP-BN-10-100 using SGD with different batch sizes $B \in \{16,32,64,128\}$ and values of $\lambda \in \{0, \expnumber{1}{-4}, \expnumber{2}{-4}, \expnumber{4}{-4}, \expnumber{8}{-4}, \expnumber{6}{-3}\}$. The learning rate is initially $\eta=0.1$ and is decayed by a factor of $0.1$ three times at epochs 60, 100, and 200 and training is stopped after 500 epoch and the weight decay is $\lambda=\expnumber{5}{-4}$.
The results are reported in Fig.~\ref{fig:wd}. As can be seen, by increasing $\lambda$ we typically impose stronger rank minimization constraints, as predicted in Sec.~\ref{sec:lowrank}. In addition, regardless of the batch size, we lose the rank minimization constraint when training with $\lambda=0$ as predicted in Sec.~\ref{sec:lowrank}.
% We repeated the same experiment with instances of MLP-10-100 without the batch normalization as well. In this case, we used $\lambda \in \{0, \expnumber{1}{-3}, \expnumber{2}{-3}, \expnumber{4}{-3}, \expnumber{8}{-3}, \expnumber{6}{-2}, \expnumber{2}{-1}\}$. For each run we trained the network for 500 epochs using SGD with momentum $0.9$, initial learning rate $0.01$ and batch size 64.
\section{Conclusions}\label{sec:conclusions}
{\color{red}needs to be rewritten} In this work we theoretically studied the source of SGD noise and the implicit bias towards low-rank weight matrices in deep neural networks. We make several theoretical and empirical observations:
\begin{itemize}[leftmargin=0.5cm]
\item SGD noise is inevitable when applying weight decay, regardless of the batch size;
\item We show that the batch size, weight decay and learning rate contribute to a regularizing effect on the rank of the trainable matrices of neural networks;
\item Weight decay and training with mini-batches are necessary to obtain rank minimization.
\end{itemize}
It has not escaped our attention that these results have a rich set of implications with respect to other mysteries of deep learning that are yet to be solved. They include questions such as why SGD training of networks generalizes well, the relationship between rank minimization and architectural choices (residual networks~\citep{He_2017_ICCV}, transformers~\citep{NIPS2017_3f5ee243} and convolutional networks~\citep{3a6e6cc3ac6e4326b35a6315d7cbddf7}) and whether rank minimization extends beyond the standard settings of supervised learning.
\bibliography{aaai22}
%\bibliographystyle{natbib}
\appendix
\onecolumn
\section{Proofs}
\lowrankSec*
\begin{proof}
We would like to show that $\frac{\bar{\partial} f_W(x_i;\tX)}{\partial V^{pq}}$ and $\sum^{B}_{i=1}\beta_i \cdot \frac{\bar{\partial} f_W(x_i;\tX)}{\partial V^{pq}}$ are matrices of rank $\leq N^{pq}B$. We follow the proof of Lem.~2 in~\citep{https://doi.org/10.48550/arxiv.2206.05794} and extend it to the case where the network includes convolutional layers. For this purpose, we represent the top layers of the network as a sequence of linear transformations and the layer $C^{pq}$ as a block matrix with blocks $V^{pq}$ by the same technique as in the proof of Lem.~\ref{lem:rank}. The layers below $C^{pq}$ are composed together with the input into a vector that is independent of $V^{pq}$.
We note that for any fixed vector $x \in \R^{c_1 \times h_1 \times w_1}$, the network can be written as follows
\begin{equation*}
f(x;\tX) ~=~ \sum_{l_1 \in \textnormal{pred}(L)} C^{Ll_1} (v_{l_1}(x;\tX)).
\end{equation*}
In addition, each layer $v_{l_1}$ can be written as follows
\begin{equation*}
\begin{aligned}
v_{l_1}(x;\tX) ~&=~ D_{l_1}(x;\tX) \odot n_{l_1}(u_{l_1}(x;\tX);\tX),
\end{aligned}
\end{equation*}
where $D_{l_1}(x;\tX) := \sigma'(n_{l_1}(u_{l_1}(x;\tX);\tX))$ and $u_{l_1}(x;\tX) := \sum_{l_2\in \textnormal{pred}(l_1)} C^{l_1l_2} (v_{l_2}(x;\tX))$. Therefore, we can write
\begin{equation*}
n_{l_1}(u_{l_1}(x;\tX);\tX) ~=~ \gamma_{l_1} \cdot A_{l_1}(\tX) \odot \left(u_{l_1}(x;\tX) - b_{l_1}(\tX) \right) + \one_{c_{l_1} \times h_{l_1} \times w_{l_1}} \cdot \beta_{l_1},
\end{equation*}
where $b_{l}(\tX) := \E_{\hat{x}\sim U[\tX]}[u_{l}(\hat{x};\tX)] = \sum_{t\in \textnormal{pred}(l)} C^{lt} (\mu_t(\tX))$ and $\mu_t(\tX) := \E_{x \sim U[\tX]}[v_{t}(x;\tX)]$. Hence, we can write
\begin{equation*}
\begin{aligned}
v_{l_1}(x;\tX) ~&=~ \sum_{l_2\in \textnormal{pred}(l_1)} Q_{l_1}(x;\tX) \odot \left[ C^{l_1l_2}(v_{l_2}(x;\tX)) - C^{l_1l_2}(\mu_{l_2}(\tX)) \right] + D_{l_1}(x;\tX) \cdot \beta_{l_1} \\
~&=~ \sum_{l_2\in \textnormal{pred}(l_1)} Q_{l_1}(x;\tX) \odot C^{l_1l_2}\left( v_{l_2}(x;\tX) - \mu_{l_2}(\tX) \right) + D_{l_1}(x;\tX) \cdot \beta_{l_1}, \\
\end{aligned}
\end{equation*}
where $Q_{l}(x;\tX) := D_{l}(x;\tX) \odot A_{l}(\tX)$ and the last equation follows from the linearity of convolutional layers.
% In addition, the flattened version of each layer $v_{l_1}$ can be written as
% \begin{equation*}
% \begin{aligned}
% \vect(v_{l_1}(x;\tX)) ~&=~ \bar{D}_{l_1}(x;\tX) \cdot \vect(n_l(u_l(x;\tX);\tX)) %\left[\gamma_{l_1} \cdot \bar{A}_{l_1}(\tX) \cdot \left(\sum_{l_2\in \textnormal{pred}(l_1)} W^{l_1l_2} \cdot \vect(v_{l_2}(x;\tX)) - \vect(b_{l_2}(\tX)) \right) + \one_{c_{l_1} h_{l_1} w_{l_1}} \cdot \beta_{l_1}\right]\\
% %~&=~ \bar{D}_{l_1}(x;\tX) \cdot \left[ \gamma_{l_1}\cdot \bar{A}_{l_1}(\tX) \cdot \left(\sum_{l_2\in \textnormal{pred}(l_1)} W^{l_1l_2} \cdot \left[\vect(v_{l_2}(x;\tX)) - \vect(\mu_{l_2}(\tX)) \right] \right) +\one_{c_{l_1} h_{l_1} w_{l_1}} \cdot \beta_{l_1} \right]\\
% %~&=~ \sum_{l_2\in \textnormal{pred}(l_1)} Q_{l_1}(x;\tX) \cdot W^{l_1l_2}\cdot \left( \vect(v_{l_2}(x;\tX)) - \vect(\mu_{l_2}(\tX)) \right) + \bar{D}_{l_1}(x;\tX) \cdot \one_{c_{l_1} h_{l_1} w_{l_1}} \cdot \beta_{l_1},
% \end{aligned}
% \end{equation*}
% where $\bar{D}_l(x;\tX) := \textnormal{diag}[D_l(x;\tX)]$, $D_l(x;\tX) := \sigma'(n_l(u_l(x;\tX);\tX))$ and
% \begin{equation*}
% \vect(n_l(u_l(x;\tX);\tX)) ~=~ \left[\gamma_{l} \cdot A_{l}(\tX) \cdot \left(\sum_{l'\in \textnormal{pred}(l)} W^{ll'} \cdot \vect(v_{l'}(x;\tX)) - \vect(b_{l'}(\tX)) \right) + \one_{c_{l_1} h_{l_1} w_{l_1}} \cdot \beta_{l}\right].
% \end{equation*}
% \begin{equation*}
% \begin{aligned}
% \vect(v_{l_1}(x;\tX)) ~&=~ \bar{D}_{l_1}(x;\tX) \cdot \left[\gamma_{l_1} \cdot \bar{A}_{l_1}(\tX) \cdot \left(\sum_{l_2\in \textnormal{pred}(l_1)} W^{l_1l_2} \cdot \vect(v_{l_2}(x;\tX)) - \vect(b_{l_2}(\tX)) \right) + \one_{c_{l_1} h_{l_1} w_{l_1}} \cdot \beta_{l_1}\right]\\
% ~&=~ \bar{D}_{l_1}(x;\tX) \cdot \left[ \gamma_{l_1}\cdot \bar{A}_{l_1}(\tX) \cdot \left(\sum_{l_2\in \textnormal{pred}(l_1)} W^{l_1l_2} \cdot \left[\vect(v_{l_2}(x;\tX)) - \vect(\mu_{l_2}(\tX)) \right] \right) +\one_{c_{l_1} h_{l_1} w_{l_1}} \cdot \beta_{l_1} \right]\\
% ~&=~ \sum_{l_2\in \textnormal{pred}(l_1)} Q_{l_1}(x;\tX) \cdot W^{l_1l_2}\cdot \left( \vect(v_{l_2}(x;\tX)) - \vect(\mu_{l_2}(\tX)) \right) + \bar{D}_{l_1}(x;\tX) \cdot \one_{c_{l_1} h_{l_1} w_{l_1}} \cdot \beta_{l_1},
% \end{aligned}
% \end{equation*}
% Here,
% \begin{equation*}
% b_l(\tX) ~:=~ \sum_{t\in \textnormal{pred}(l)} W^{lt} \cdot \mu_t(\tX) ~:=~ \sum_{t\in \textnormal{pred}(l)} W^{lt} \cdot \E_{x \sim U[\tX]}\left[ v_{t}(x;\tX) \right]
% \end{equation*}
% and $Q_l(x;\tX) := \gamma_l \cdot D_{l}(x;\tX) \cdot A_l(x;\tX)$. For simplicity, we write $l_0 := L$, $\mu_l := \mu_l(x;\tX)$, $v^j_{l} := v_{l}(x_j;\tX)$ and $Q^{j}_l := Q_l(x_j;\tX)$.
We denote $\mu_l := \mu_l(\tX)$, $v^j_l := v_l(x_j;\tX)$, $D^{j}_l := D_l(x_j;\tX)$, $A_l := A_l(\tX)$, $Q^{j}_l := Q_l(x_j;\tX)$ and define
\begin{equation*}
\Delta^{j}_l := \twopartdef
{x_j} {l=1}
{v^{j}_l - \mu_{l}} {l\neq 1}.
\end{equation*}
Hence, we can write $f(x_{j_0};\tX)$ as a follows
\begin{equation*}
\begin{aligned}
f(x_{j_0};\tX) ~&=~ \sum_{l_1 \in \textnormal{pred}(l_0)} C^{l_0l_1} \cdot v^{j_0}_{l_1} \\
~&=~ \sum_{l_1 \in \textnormal{pred}(l_0)} C^{l_0l_1} \left(Q^{j_0}_{l_1} \odot \left( \sum_{l_2 \in \textnormal{pred}(l_1)} C^{l_1l_2} (v^{j_0}_{l_2} - \mu_{l_2}) \right) +D^{j_0}_{l_1}\cdot \beta_{l_1}\right) \\
~&=~ \sum_{l_1 \in \textnormal{pred}(l_0)} C^{l_0l_1} (D^{j_0}_{l_1}\cdot \beta_{l_1}) \\
&\quad + \sum_{l_1 \in \textnormal{pred}(l_0)} \sum_{l_2 \in \textnormal{pred}(l_1)} C^{l_0l_1} (Q^{j_0}_{l_1} \odot C^{l_1l_2} (v^{j_0}_{l_2} - \mu_{l_2})) \\
~&=~ \sum_{l_1 \in \textnormal{pred}(l_0)} C^{l_0l_1} (D^{j_0}_{l_1}\cdot \beta_{l_1}) \\
&\quad + \sum_{l_1 \in \textnormal{pred}(l_0)} \sum_{l_2 \in \textnormal{pred}(l_1)} C^{l_0l_1} (Q^{j_0}_{l_1} \odot C^{l_1l_2} (\Delta^{j_0}_{l_2})) \\
\end{aligned}
\end{equation*}
% Similarly,
% \begin{equation}
% \begin{aligned}
% \Delta^{j_0}_{l_2} ~&=~ v^{j_0}_{l_2} - \mu_{l_2} \\
% ~&=~ D^{j_0}_{l_2} \cdot \one \cdot \beta_{l_2} \\
% &\quad+ Q^{j_0}_{l_2} \cdot \left(\sum_{l_3\in \textnormal{pred}(l_2)} W^{l_2l_3} \cdot (v^{j_0}_{l_3} - \mu_{l_3}) \right) \\
% &\quad- \frac{1}{B}\sum^{B}_{j_1=1} Q^{j_1}_{l_2} \cdot \left(\sum_{l_3\in \textnormal{pred}(l_2)} W^{l_2l_3} \cdot (v^{j_1}_{l_3} - \mu_{l_3}) \right) \\
% ~&=~ D^{j_0}_{l_2} \cdot \one \cdot \beta_{l_2} + \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{l_2} \cdot \left(\sum_{l_3\in \textnormal{pred}(l_2)} W^{l_2l_3} \cdot (v^{j_1}_{l_3} - \mu_{l_3}) \right) \\
% ~&=~ D^{j_0}_{l_2} \cdot \one \cdot \beta_{l_2} + \sum_{l_3\in \textnormal{pred}(l_2)}\sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{l_2} \cdot W^{l_2l_3} \cdot \Delta^{j_1}_{l_3}, \\
% \end{aligned}
% \end{equation}
For any $l_{t+2}\neq 1$, we can write
\begin{equation}\label{eq:delta}
\begin{aligned}
\Delta^{j_t}_{l_{t+2}} ~&=~ v^{j_t}_{l_{t+2}} - \mu_{l_{t+2}} \\
~&=~ \left(D^{j_{t}}_{l_{t+2}} - \frac{1}{B}\sum^{B}_{j_{t+1}=1} D^{j_{t+1}}_{l_{t+2}}\right) \cdot \beta_{l_{t+2}} \\
&\quad + Q^{j_t}_{l_{t+2}} \odot \left(\sum_{l_{t+3}\in \textnormal{pred}(l_{t+2})} C^{l_{t+2}l_{t+3}} (v^{j_t}_{l_{t+3}} - \mu_{l_{t+3}}) \right) \\
&\quad- \frac{1}{B}\sum^{B}_{j_{t+1}=1} Q^{j_{t+1}}_{l_{t+2}} \odot \left(\sum_{l_{t+3}\in \textnormal{pred}(l_{t+2})} C^{l_{t+2}l_{t+3}} (v^{j_{t+1}}_{l_{t+3}} - \mu_{l_{t+3}}) \right) \\
~&=~ \sum^{B}_{j_{t+1}=1}\alpha_{j_t,j_{t+1}} \cdot D^{j_{t+1}}_{l_{t+2}} \cdot \beta_{l_{t+2}} + \sum_{l_{t+3}\in \textnormal{pred}(l_{t+2})}\sum^{B}_{j_{t+1}=1} \alpha_{j_t,j_{t+1}} Q^{j_{t+1}}_{l_{t+2}} \odot C^{l_{t+2}l_{t+3}} (\Delta^{j_{t+1}}_{l_{t+3}}), \\
\end{aligned}
\end{equation}
where $\alpha_{i,j}=-1/B$ for all pairs $i\neq j$ and $\alpha_{i,i}=\frac{B-1}{B}$ otherwise. For simplicity, from now on we omit writing the clauses `$(~~)$' when applying a convolutional layer $C^{ij}$ and write $\circ$ instead, i.e., $C^{ij}\circ x := C^{ij}(x)$. Hence, by recursion we can write the $T$'th expansion of $f_W$,
\begin{equation*}
\begin{aligned}
f_W(x_{j_0};\tX) ~&=~ \sum_{l_1 \in \textnormal{pred}(l_0)} C^{l_0l_1} \circ D^{j_0}_{l_1}\cdot \beta_{l_1} \\
% &\quad + \sum_{l_1 \in \textnormal{pred}(l_0)} \sum_{l_2 \in \textnormal{pred}(l_1)} C^{l_0l_1} \circ Q^{j_0}_{l_1} \odot C^{l_1l_2} \circ D^{j_0}_{l_2} \cdot \beta_{l_2}\\
&\quad+\sum^{T}_{t=2} \sum_{l_1 \in \textnormal{pred}(l_0)} \dots \sum_{l_{t+1}\in \textnormal{pred}(l_{t})} C^{l_0l_1} \circ Q^{j_0}_{l_1} \odot C^{l_1l_2}\circ \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{l_2} \odot C^{l_2l_3} \\
&\quad\quad\quad \cdots C^{l_{t-1} l_{t}}\circ \sum^{B}_{j_{t-1}=1} \alpha_{j_{t-2},j_{t-1}} Q^{j_{t-1}}_{l_{t}} \odot C^{l_{t}l_{t+1}} \circ \sum^{B}_{j_t = 1} \alpha_{j_{t-1},j_t} \cdot D^{j_{t}}_{l_{t+1}} \cdot \beta_{l_{t+1}} \\
&\quad+\sum_{l_1 \in \textnormal{pred}(l_0)} \sum_{l_2 \in \textnormal{pred}(l_1)} \sum_{l_3 \in \textnormal{pred}(l_2)} \dots \sum_{l_{T+2}\in \textnormal{pred}(l_{T+1})} C^{l_0l_1} \circ Q^{j_0}_{l_1} \odot C^{l_1l_2}\circ \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{l_2} \odot C^{l_2l_3} \\
&\quad\quad\quad \cdots C^{l_T l_{T+1}}\circ \sum^{B}_{j_{T}=1} \alpha_{j_{T-1},j_T} Q^{j_T}_{l_{T+1}} \odot C^{l_{T+1}l_{T+2}}\circ \Delta^{j_{T}}_{l_{T+2}}. \\
\end{aligned}
\end{equation*}
Instead of recursively expanding $f_W$'s expression $T$ times by rewriting $\Delta^{j_{t}}_{l_{t+2}}$ using \eqref{eq:delta}, we can expand through each term $\Delta^{j_{t}}_{l_{t+2}}$ based on different kinds of criteria. For instance, we can keep $\Delta^{j_{t}}_{l_{t+2}}$ unchanged or expand it based on a whether $(l_{t+2},l_{t+3})=(p,q)$ or not. Therefore, we can write $f_W(x_{j_0};\tX)$ as the sum of multiplications along paths $\pi$ that start with $(p,q)$ and end by $L$ and paths that start with $1$ and end with $L$ that do not pass through $(p,q)$. Formally, let $j_0 \in [m]$ and $\pi$ be a path from some node to $\pi_T=L$ of length $T$,
\begin{equation*}
E^{j_0}_{1,\pi} ~:=~ \twopartdef
{C^{\pi_2\pi_1} \circ D^{j_0}_{\pi_1} \cdot \beta_{\pi_1}} {T=2}
%{C^{\pi_3\pi_2} \circ Q^{j_0}_{\pi_2} \odot C^{\pi_2\pi_1} \circ D^{j_0}_{\pi_1} \cdot \beta_{\pi_1}} {T=3}
{
\begin{aligned}
&C^{\pi_T\pi_{T-1}} \circ Q^{j_0}_{\pi_{T-1}} \odot C^{\pi_{T-1}\pi_{T-2}} \circ \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{\pi_{T-2}} \odot C^{\pi_{T-2}\pi_{T-3}}\\
&\quad \cdots C^{\pi_{3} \pi_{2}} \circ \sum^{B}_{j_{T-2}=1} \alpha_{j_{T-3},j_{T-2}} Q^{j_{T-2}}_{\pi_{1}} \odot C^{\pi_2 \pi_1} \circ \sum^{B}_{j_{T-1}=1}\alpha_{j_{T-2},j_{T-1}}\cdot D^{j_{T-1}}_{\pi_{1}} \cdot \beta_{\pi_1}
\end{aligned}
}
{T > 2}
\end{equation*}
and also
\begin{equation*}
E^{j_0}_{2,\pi} ~:=~ \twopartdef
{
\begin{aligned}
&C^{\pi_T \pi_{T-1}} \circ Q^{j_0}_{\pi_{T-1}} \odot C^{\pi_{T-1}\pi_{T-2}} \circ \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{\pi_{T-2}} \\
&\quad \cdots C^{\pi_3 \pi_2} \circ \sum^{B}_{j_{T-3}=1} \alpha_{j_{T-4},j_{T-3}} Q^{j_{T-3}}_{\pi_{2}} \odot C^{\pi_2 \pi_1} \circ \Delta^{j_{T-3}}_{\pi_1}
\end{aligned}
} {\pi_{1}\neq 1}
{
\begin{aligned}
&C^{\pi_T \pi_{T-1}} \circ Q^{j_0}_{\pi_{T-1}} \cdot C^{\pi_{T-1}\pi_{T-2}} \circ \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{\pi_{T-2}} \\
&\quad \cdots C^{\pi_3 \pi_2} \circ \sum^{B}_{j_{T-3}=1} \alpha_{j_{T-4},j_{T-3}} Q^{j_{T-3}}_{\pi_2} \odot C^{\pi_2 \pi_1} \circ x
\end{aligned}
}
{\pi_1=1}
\end{equation*}
With this terminology we can write
\begin{equation*}
\begin{aligned}
f_W(x_{j_0};\tX) ~&=~
\sum_{\substack{ \textnormal{$\pi$ from some node $\neq 1$ to $L$} \\ \textnormal{$(p,q) \notin \pi$}}} E^{j_0}_{1,\pi}
+ \sum_{\substack{ \textnormal{$\pi$ from 1 to $L$} \\ \textnormal{$(p,q) \notin \pi$}}} E^{j_0}_{2,\pi}
+\sum_{\substack{\pi \textnormal{ from $q$ to $L$} \\ \pi_2=p,\pi_1=q}} E^{j_0}_{2,\pi} ~=:~ Z^{j_0}_1+Z^{j_0}_2+Z^{j_0}_3.
\end{aligned}
\end{equation*}
We note that for any $x$ and $\tX$, with measure $1$ over $W$, the matrices $D_{l}(x;\tX)$ are constant in the neighborhood of $W$. Together with the assumption that $A_{l}(x;\tX)$ are detached, the pseudo-gradients of the matrices $Q^j_l$ with respect to $V^{pq}$ are zero. Therefore, since the terms $Z^{j_0}_1$ and $Z^{j_0}_2$ do not involve $V^{pq}$, we have $\frac{\bar{\partial} f_W(x_{j_0};\tX)}{\partial V^{pq}} = \frac{\bar{\partial} Z^{j_0}_3}{\partial V^{pq}}$. As a next step, we represent $\frac{\bar{\partial} Z^{j_0}_2}{\partial V^{pq}}$. For simplicity, we assume that each $\pi$ from $q$ to $L$ that crosses through $(p,q)$ is of length $T(\pi)\geq 3$. With a slight abuse of notation, we denote $T:=T(\pi)$. We can write
% \begin{equation*}
% \begin{aligned}
% Z^{j_0}_3 ~&=~ \sum_{\substack{\pi \textnormal{ from some node to $L$} \\ \pi_2=p,\pi_1=q}} C^{\pi_T\pi_{T-1}} \circ Q^{j_0}_{\pi_{T-1}} \odot C^{\pi_{T-1}\pi_{T-2}} \circ \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{\pi_{T-2}} \odot C^{\pi_{T-2}\pi_{T-3}} \\
% &\quad\quad\quad\quad\quad \cdots \sum^{B}_{j_{T-4}=1} \alpha_{j_{T-5},j_{T-4}} Q^{j_{T-4}}_{\pi_{3}} \odot C^{\pi_{3} \pi_{2}} \circ \sum^{B}_{j_{T-3}=1} \alpha_{j_{T-4},j_{T-3}} Q^{j_{T-3}}_{\pi_{2}} \odot C^{\pi_2 \pi_1} \circ D^{j_{T-2}}_{\pi_{1}} \cdot \beta_{\pi_1} \\
% ~&=~ \sum_{\substack{\pi \textnormal{ from $q$ to $L$} \\ \pi_2=p,\pi_1=q}} W^{\pi_T\pi_{T-1}} \cdot Q^{j_0}_{\pi_{T-1}} \cdot W^{\pi_{T-1}\pi_{T-2}} \cdot \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{\pi_{T-2}} \cdot W^{\pi_{T-2}\pi_{T-3}} \\
% &\quad\quad\quad\quad\quad \cdots \sum^{B}_{j_{T-3}=1} \alpha_{j_{T-4},j_{T-3}} Q^{j_{T-3}}_{\pi_{2}} \cdot W^{\pi_2\pi_1} \cdot \sum_{l \in \textnormal{pred}(\pi_1)}\sum^{B}_{j_{T-2}=1} \alpha_{j_{T-3},j_{T-2}} Q^{j_{T-2}}_{l} \cdot W^{\pi_1 l} \cdot D^{j_{T-2}}_{l} \cdot \beta_{l} \\
% ~&=~ \sum^{B}_{j_{T-3}=1} \left[\sum_{\substack{\pi \textnormal{ from $q$ to $L$} \\ \pi_2=p,\pi_1=q}} (u^{\pi,1}_{j_0,j_{T-3}})^{\top} \right] \cdot W^{pq} \cdot \sum_{l \in \textnormal{pred}(q)}\sum^{B}_{j_{T-2}=1} \alpha_{j_{T-3},j_{T-2}} Q^{j_{T-2}}_{q} \cdot W^{ql} \cdot D^{j_{T-2}}_{q} \cdot \beta_{q}\\
% ~&=~ \sum^{B}_{j_{T-3}=1} (u^{1}_{j_0,j_{T-3}})^{\top} \cdot W^{pq} \cdot u^{2}_{j_{T-3}},
% \end{aligned}
% \end{equation*}
\begin{equation*}
\begin{aligned}
Z^{j_0}_3 ~&=~ \sum_{\substack{\pi \textnormal{ from $p$ to $L$} \\ \pi_2=p,\pi_1=q}} C^{\pi_T \pi_{T-1}} \circ Q^{j_0}_{\pi_{T-1}} \odot C^{\pi_{T-1}\pi_{T-2}} \circ \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{\pi_{T-2}} \cdots C^{\pi_3 \pi_2} \circ \sum^{B}_{j_{T-3}=1} \alpha_{j_{T-4},j_{T-3}} Q^{j_{T-3}}_{\pi_{2}} \odot C^{\pi_2 \pi_1} \circ \Delta^{j_{T-3}}_{\pi_1} \\
~&=~ \sum_{\substack{\pi \textnormal{ from $p$ to $L$} \\ \pi_2=p,\pi_1=q}} C^{\pi_T \pi_{T-1}} \circ Q^{j_0}_{\pi_{T-1}} \odot C^{\pi_{T-1}\pi_{T-2}} \circ \sum^{B}_{j_1=1} \alpha_{j_0,j_1} Q^{j_1}_{\pi_{T-2}} \cdots C^{\pi_3 \pi_2} \circ \sum^{B}_{j_{T-3}=1} \alpha_{j_{T-4},j_{T-3}} Q^{j_{T-3}}_{\pi_2} \odot C^{pq} \circ \Delta^{j_{T-3}}_{q}. \\
%~&=~ \sum^{B}_{j_{T-3}=1} (z^{1}_{j_0,j_{T-3}})^{\top} \cdot W^{pq} \cdot z^{2}_{j_{T-3}},
\end{aligned}
\end{equation*}
As a next step we linearize the above expression. We notice that $\vect(C^{ij} \circ z) = U^{ij} \cdot \vect^{ij}(z) = W^{ij} \cdot \vect(z)$ and $\vect(Q^{j}_l \odot z) = \textnormal{diag}(\vect(Q^{j}_{l})) \cdot \vect(z)$. Therefore, we replace $C^{pq}$ by $U^{pq}$ and $\Delta^{j_{T-3}}_q$ with $\vect^{pq}(\Delta^{j_{T-3}}_q)$. In addition, each one of the convolutional layers $C^{ij}$ above $C^{pq}$ is replaced with $W^{ij}$ and each application of element-wise multiplication by $Q^{j}_l$ is replace with a (right side) matrix multiplication by $P^{j}_{l} := \textnormal{diag}(\vect(Q^{j}_{l}))$. We obtain that
\begin{equation*}
\begin{aligned}
Z^{j_0}_3 ~&=~ \sum_{\substack{\pi \textnormal{ from $p$ to $L$} \\ \pi_2=p,\pi_1=q}} W^{\pi_T \pi_{T-1}} \cdot P^{j_0}_{\pi_{T-1}} \cdot W^{\pi_{T-1}\pi_{T-2}} \cdot \sum^{B}_{j_1=1} \alpha_{j_0,j_1} P^{j_1}_{\pi_{T-2}} \cdots W^{\pi_3 \pi_2} \cdot \sum^{B}_{j_{T-3}=1} \alpha_{j_{T-4},j_{T-3}} P^{j_{T-3}}_{\pi_2} \cdot U^{pq} \cdot \vect^{pq}(\Delta^{j_{T-3}}_{q}) \\
~&=~ \sum^{B}_{j_{T-3}=1} (z^{1}_{j_0,j_{T-3}})^{\top} \cdot U^{pq} \cdot z^{2}_{j_{T-3}},
\end{aligned}
\end{equation*}
where
\begin{equation*}
\begin{aligned}
(z^{1}_{j_0,j_{T-3}})^{\top} ~:&=~ \sum_{\substack{\pi \textnormal{ from $p$ to $L$} \\ \pi_2=p,\pi_1=q}} W^{\pi_T \pi_{T-1}} \cdot P^{j_0}_{\pi_{T-1}} \cdot W^{\pi_{T-1}\pi_{T-2}} \cdot \sum^{B}_{j_1=1} \alpha_{j_0,j_1} P^{j_1}_{\pi_{T-2}} \cdots W^{\pi_3 \pi_2} \cdot \alpha_{j_{T-4},j_{T-3}} P^{j_{T-3}}_{\pi_2} \\
z^{2}_{j_{T-3}} ~:&=~ \vect^{pq}(\Delta^{j_{T-3}}_{q}).
\end{aligned}
\end{equation*}
We notice that $\frac{\bar{\partial} z^{1}_{j_0,j_{T-3}}}{\partial V^{pq}} = 0$ since it is decomposed of $W^{ij}$ for $(i,j)\neq (p,q)$ and $P^{j}_l$ which are reshapes of $Q^j_l = D^j_l \odot A_l$, where $\frac{\partial D^j_l}{\partial V^{pq}}=0$ since $D^j_l$ are locally constant as functions of $W$ and $\frac{\bar{\partial} A_l}{\partial V^{pq}}=0$ by assumption. Furthermore, $\frac{\partial z^{2}_{j_{T-3}}}{\partial V^{pq}}=0$ since $\Delta^{j_{T-3}}_q = v^{j_{T-3}}_q - \mu_q$, where $v_q$ is a layer below $C^{pq}$ (so its outputs are independent of $V^{pq}$). We note that $U^{pq}$ is a block matrix with $N^{pq}$ blocks of $V^{pq}$. Therefore, by we can write
\begin{equation*}
Z^{j_0}_{3} ~=~ \sum^{B}_{j_{T-3}=1} \sum^{N^{pq}}_{i=1} a^{\top}_{i,j_0,j_{T-3}} \cdot V^{pq} \cdot b_{i,j_{T-3}},
\end{equation*}
where $z^1_{j_0,j_{T-3}}= (a_{1,j_0,j_{T-3}} \| \dots \| a_{N^{pq},j_{T-3}})$ and $z^{2}_{j_{T-3}} = (b_{1,j_{T-3}} \| \dots \| b_{N^{pq},j_{T-3}})$ are partitions of the vectors $z^1_{j_0,j_{T-3}}$ and $z^2_{j_{T-3}}$ into $N^{pq}$ consecutive sub-vectors of the same size. In conclusion, we have
\begin{equation*}
\frac{\bar{\partial} f(x_{j_0};\tX)}{\partial V^{pq}} ~=~ \frac{\partial Z^{j_0}_{3}}{\partial V^{pq}} ~=~ \sum^{B}_{j_{T-3}=1} \sum^{N^{pq}}_{i=1} a_{i,j_0,j_{T-3}} \cdot b^{\top}_{i,j_{T-3}}
\end{equation*}
which is a matrix of rank $\leq BN^{pq}$.
% Therefore, we can write $Z^{j_0}_2 = \sum^{B}_{j_{T-3}=1} (u^1_{j_0,j_{T-3}})^{\top} \cdot W^{pq} \cdot u^2_{j_{T-3}}$ and we obtain that $\frac{\bar{\partial} Z^{j_0}_2}{\partial W^{pq}} = \sum^{B}_{j_{T-3}=1} (u^1_{j_0,j_{T-3}})\cdot (u^2_{j_{T-3}})^{\top}$.
% Next, we would like to analyze the pseudo gradient $\frac{\bar{\partial} Z^{j_0}_4}{\partial W^{pq}}$. We can write
We also notice that $b_{i,j_{T-3}}$ are independent of $j_0$. Therefore, we obtain that
\begin{equation*}
\begin{aligned}
\sum^{B}_{j_0=1}\alpha_{j_0} \cdot \frac{\bar{\partial} f_W(x_{j_0};\tX)}{\partial V^{pq}}
~&=~ \sum^{B}_{j_0=1} \alpha_{j_0}\cdot \frac{\bar{\partial} Z^{j_0}_3}{\partial
V^{pq}} \\
~&=~ \sum^{B}_{j_0=1} \alpha_{j_0} \sum^{B}_{j_{T-3}=1} \sum^{N^{pq}}_{i=1} a_{i,j_0,j_{T-3}} \cdot b^{\top}_{i,j_{T-3}}\\
~&=~ \sum^{B}_{j_{T-3}=1} \sum^{N^{pq}}_{i=1} \left(\sum^{B}_{j_0=1} \alpha_{j_0} \cdot a_{i,j_0,j_{T-3}}\right) \cdot b^{\top}_{i,j_{T-3}},
\end{aligned}
\end{equation*}
which is also a matrix of rank $\leq BN^{pq}$.
\end{proof}
\end{document}
\begin{thebibliography}{26}
\providecommand{\natexlab}[1]{#1}
\providecommand{\url}[1]{\texttt{#1}}
\expandafter\ifx\csname urlstyle\endcsname\relax
\providecommand{\doi}[1]{doi: #1}\else
\providecommand{\doi}{doi: \begingroup \urlstyle{rm}\Url}\fi
\bibitem[Alvarez \& Salzmann(2017)Alvarez and
Salzmann]{10.5555/3294771.3294853}
Alvarez, J.~M. and Salzmann, M.
\newblock Compression-aware training of deep networks.
\newblock In \emph{Proceedings of the 31st International Conference on Neural
Information Processing Systems}, NIPS'17, pp.\ 856–867, Red Hook, NY, USA,
2017. Curran Associates Inc.
\newblock ISBN 9781510860964.
\bibitem[Arora et~al.(2018)Arora, Ge, Neyshabur, and Zhang]{pmlr-v80-arora18b}
Arora, S., Ge, R., Neyshabur, B., and Zhang, Y.
\newblock Stronger generalization bounds for deep nets via a compression
approach.
\newblock In Dy, J. and Krause, A. (eds.), \emph{Proceedings of the 35th
International Conference on Machine Learning}, volume~80 of \emph{Proceedings
of Machine Learning Research}, pp.\ 254--263. PMLR, 10--15 Jul 2018.
\bibitem[Bottou(1991)]{bottou-91c}
Bottou, L.
\newblock Stochastic gradient learning in neural networks.
\newblock In \emph{Proceedings of Neuro-N\^imes 91}, Nimes, France, 1991. EC2.
\newblock URL \url{http://leon.bottou.org/papers/bottou-91c}.
\bibitem[Denton et~al.(2014)Denton, Zaremba, Bruna, LeCun, and
Fergus]{NIPS2014_2afe4567}
Denton, E.~L., Zaremba, W., Bruna, J., LeCun, Y., and Fergus, R.
\newblock Exploiting linear structure within convolutional networks for
efficient evaluation.
\newblock In Ghahramani, Z., Welling, M., Cortes, C., Lawrence, N., and
Weinberger, K. (eds.), \emph{Advances in Neural Information Processing
Systems}, volume~27. Curran Associates, Inc., 2014.
\newblock URL
\url{https://proceedings.neurips.cc/paper/2014/file/2afe4567e1bf64d32a5527244d104cea-Paper.pdf}.
\bibitem[Dosovitskiy et~al.(2020)Dosovitskiy, Beyer, Kolesnikov, Weissenborn,
Zhai, Unterthiner, Dehghani, Minderer, Heigold, Gelly,
et~al.]{dosovitskiy2020image}
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X.,
Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et~al.
\newblock An image is worth 16x16 words: Transformers for image recognition at
scale.
\newblock \emph{arXiv preprint arXiv:2010.11929}, 2020.
\bibitem[Goodfellow et~al.(2016)Goodfellow, Bengio, and
Courville]{Goodfellow-et-al-2016}
Goodfellow, I., Bengio, Y., and Courville, A.
\newblock \emph{Deep Learning}.
\newblock MIT Press, 2016.
\newblock \url{http://www.deeplearningbook.org}.
\bibitem[Gunasekar et~al.(2017)Gunasekar, Woodworth, Bhojanapalli, Neyshabur,
and Srebro]{https://doi.org/10.48550/arxiv.1705.09280}
Gunasekar, S., Woodworth, B., Bhojanapalli, S., Neyshabur, B., and Srebro, N.
\newblock Implicit regularization in matrix factorization, 2017.
\newblock URL \url{https://arxiv.org/abs/1705.09280}.
\bibitem[He et~al.(2016)He, Zhang, Ren, and Sun]{7780459}
He, K., Zhang, X., Ren, S., and Sun, J.
\newblock Deep residual learning for image recognition.
\newblock In \emph{2016 IEEE Conference on Computer Vision and Pattern
Recognition (CVPR)}, pp.\ 770--778, 2016.
\newblock \doi{10.1109/CVPR.2016.90}.
\bibitem[Hoffer et~al.(2017)Hoffer, Hubara, and Soudry]{NIPS2017_a5e0ff62}
Hoffer, E., Hubara, I., and Soudry, D.
\newblock Train longer, generalize better: closing the generalization gap in
large batch training of neural networks.
\newblock In Guyon, I., Luxburg, U.~V., Bengio, S., Wallach, H., Fergus, R.,
Vishwanathan, S., and Garnett, R. (eds.), \emph{Advances in Neural
Information Processing Systems}, volume~30. Curran Associates, Inc., 2017.
\newblock URL
\url{https://proceedings.neurips.cc/paper/2017/file/a5e0ff62be0b08456fc7f1e88812af3d-Paper.pdf}.
\bibitem[Huh et~al.(2022)Huh, Mobahi, Zhang, Cheung, Agrawal, and
Isola]{huh2022the}
Huh, M., Mobahi, H., Zhang, R., Cheung, B., Agrawal, P., and Isola, P.
\newblock The low-rank simplicity bias in deep networks, 2022.
\newblock URL \url{https://openreview.net/forum?id=dn4B7Mes2z}.
\bibitem[Jastrzebski et~al.(2017)Jastrzebski, Kenton, Arpit, Ballas, Fischer,
Bengio, and Storkey]{https://doi.org/10.48550/arxiv.1711.04623}
Jastrzebski, S., Kenton, Z., Arpit, D., Ballas, N., Fischer, A., Bengio, Y.,
and Storkey, A.
\newblock Three factors influencing minima in sgd, 2017.
\newblock URL \url{https://arxiv.org/abs/1711.04623}.
\bibitem[Ji \& Telgarsky(2020)Ji and
Telgarsky]{DBLP:journals/corr/abs-2006-06657}
Ji, Z. and Telgarsky, M.
\newblock Directional convergence and alignment in deep learning.
\newblock \emph{CoRR}, abs/2006.06657, 2020.
\newblock URL \url{https://arxiv.org/abs/2006.06657}.
\bibitem[Keskar et~al.(2017)Keskar, Mudigere, Nocedal, Smelyanskiy, and
Tang]{https://doi.org/10.48550/arxiv.1609.04836}
Keskar, N.~S., Mudigere, D., Nocedal, J., Smelyanskiy, M., and Tang, P. T.~P.
\newblock On large-batch training for deep learning: Generalization gap and
sharp minima.
\newblock In \emph{International Conference on Learning Representations
(ICLR)}, 2017.
\bibitem[Le \& Jegelka(2022)Le and Jegelka]{le2022training}
Le, T. and Jegelka, S.
\newblock Training invariances and the low-rank phenomenon: beyond linear
networks.
\newblock In \emph{International Conference on Learning Representations}, 2022.
\newblock URL \url{https://openreview.net/forum?id=XEW8CQgArno}.
\bibitem[Lecun et~al.(1998)Lecun, Bottou, Bengio, and Haffner]{726791}
Lecun, Y., Bottou, L., Bengio, Y., and Haffner, P.
\newblock Gradient-based learning applied to document recognition.
\newblock \emph{Proceedings of the IEEE}, 86\penalty0 (11):\penalty0
2278--2324, 1998.
\newblock \doi{10.1109/5.726791}.
\bibitem[Li et~al.(2020)Li, Luo, and Lyu]{DBLP:journals/corr/abs-2012-09839}
Li, Z., Luo, Y., and Lyu, K.
\newblock Towards resolving the implicit bias of gradient descent for matrix
factorization: Greedy low-rank learning.
\newblock \emph{CoRR}, abs/2012.09839, 2020.
\newblock URL \url{https://arxiv.org/abs/2012.09839}.
\bibitem[Neyshabur et~al.(2017)Neyshabur, Bhojanapalli, Mcallester, and
Srebro]{NIPS2017_10ce03a1}
Neyshabur, B., Bhojanapalli, S., Mcallester, D., and Srebro, N.
\newblock Exploring generalization in deep learning.
\newblock In Guyon, I., Luxburg, U.~V., Bengio, S., Wallach, H., Fergus, R.,
Vishwanathan, S., and Garnett, R. (eds.), \emph{Advances in Neural
Information Processing Systems}, volume~30. Curran Associates, Inc., 2017.
\newblock URL
\url{https://proceedings.neurips.cc/paper/2017/file/10ce03a1ed01077e3e289f3e53c72813-Paper.pdf}.
\bibitem[Ongie \& Willett(2022)Ongie and
Willett]{https://doi.org/10.48550/arxiv.2202.00856}
Ongie, G. and Willett, R.
\newblock The role of linear layers in nonlinear interpolating networks, 2022.
\newblock URL \url{https://arxiv.org/abs/2202.00856}.
\bibitem[Razin \& Cohen(2020)Razin and
Cohen]{DBLP:journals/corr/abs-2005-06398}
Razin, N. and Cohen, N.
\newblock Implicit regularization in deep learning may not be explainable by
norms.
\newblock \emph{CoRR}, abs/2005.06398, 2020.
\newblock URL \url{https://arxiv.org/abs/2005.06398}.
\bibitem[Simonyan \& Zisserman(2014)Simonyan and
Zisserman]{DBLP:journals/corr/SimonyanZ14a}
Simonyan, K. and Zisserman, A.
\newblock Very deep convolutional networks for large-scale image recognition.
\newblock \emph{CoRR}, abs/1409.1556, 2014.
\bibitem[Timor et~al.(2022)Timor, Vardi, and
Shamir]{DBLP:journals/corr/abs-2201-12760}
Timor, N., Vardi, G., and Shamir, O.
\newblock Implicit regularization towards rank minimization in relu networks.
\newblock \emph{CoRR}, abs/2201.12760, 2022.
\newblock URL \url{https://arxiv.org/abs/2201.12760}.
\bibitem[Tukan et~al.(2021)Tukan, Maalouf, Weksler, and Feldman]{s21165599}
Tukan, M., Maalouf, A., Weksler, M., and Feldman, D.
\newblock No fine-tuning, no cry: Robust svd for compressing deep networks.
\newblock \emph{Sensors}, 21\penalty0 (16), 2021.
\newblock ISSN 1424-8220.
\newblock \doi{10.3390/s21165599}.
\newblock URL \url{https://www.mdpi.com/1424-8220/21/16/5599}.
\bibitem[Yu et~al.(2017)Yu, Liu, Wang, and Tao]{8099498}
Yu, X., Liu, T., Wang, X., and Tao, D.
\newblock On compressing deep models by low rank and sparse decomposition.
\newblock In \emph{2017 IEEE Conference on Computer Vision and Pattern
Recognition (CVPR)}, pp.\ 67--76, 2017.
\newblock \doi{10.1109/CVPR.2017.15}.
\bibitem[Zhang et~al.(2016)Zhang, Bengio, Hardt, Recht, and
Vinyals]{DBLP:journals/corr/ZhangBHRV16}
Zhang, C., Bengio, S., Hardt, M., Recht, B., and Vinyals, O.
\newblock Understanding deep learning requires rethinking generalization.
\newblock \emph{CoRR}, abs/1611.03530, 2016.
\newblock URL \url{http://arxiv.org/abs/1611.03530}.
\bibitem[Zhou \& Chellappa(1988)Zhou and Chellappa]{Zhou1988ComputationOO}
Zhou, Y.-T. and Chellappa, R.
\newblock Computation of optical flow using a neural network.
\newblock \emph{IEEE 1988 International Conference on Neural Networks}, pp.\
71--78 vol.2, 1988.
\bibitem[Zhu et~al.(2019)Zhu, Wu, Yu, Wu, and Ma]{zhu}
Zhu, Z., Wu, J., Yu, B., Wu, L., and Ma, J.
\newblock The anisotropic noise in stochastic gradient descent: Its behavior of
escaping from sharp minima and regularization effects.
\newblock In \emph{Proceedings of the 36th International Conference on Machine
Learning}, volume 139 of \emph{Proceedings of Machine Learning Research}.
PMLR, 2019.
\end{thebibliography}