0 Research T. Athawale, L. Yan, B. Wang, and C. R. Johnson are with Scientific
Computing & Imaging (SCI) Institute, University of Utah. E-mail:
{tushar.athawale, linyan, beiwang<EMAIL_ADDRESS>A. Entezari is with the
Department of CISE at the University of Florida, Gainesville, FL, 32611.
E-mail<EMAIL_ADDRESS>
# Statistical Rendering for Visualization of Red Sea Eddy Simulation Data

Tushar Athawale  Alireza Entezari  Senior Member, IEEE
Bei Wang  and Chris R. Johnson  Fellow, IEEE

###### Abstract

Analyzing the effects of ocean eddies is important in oceanology for gaining
insights into transport of energy and biogeochemical particles. We present an
application of statistical visualization algorithms for the analysis of the
Red Sea eddy simulation ensemble. Specifically, we demonstrate the
applications of _statistical volume rendering_ and _statistical Morse complex
summary maps_ to a velocity magnitude field for studying the eddy positions in
the flow dataset. In statistical volume rendering, we model per-voxel data
uncertainty using noise models, such as parametric and nonparametric, and
study the propagation of uncertainty into the volume rendering pipeline. In
the statistical Morse complex summary maps, we derive histograms charactering
uncertainty of gradient flow destinations to understand Morse complex
topological variations across the ensemble. We demonstrate the utility of our
statistical visualizations for an effective analysis of the potential eddy
positions and their spatial uncertainty.

Human-centered computingVisualizationVisualization application
domainsScientific visualization; Mathematics of computingProbability and
statisticsProbabilistic algorithms

Introduction

In this work, we propose visualizations for the analysis of the Red Sea eddy
simulation dataset, which is available via the IEEE SciVis Contest 2020
111https://kaust-vislab.github.io/SciVis2020/. The randomness in data
acquisition is captured via an ensemble of simulations, in which each ensemble
member is generated based on the MIT ocean general circulation model (MITgcm)
and the Data Research Testbed (DART) [6] with varying initial conditions. Each
ensemble is sampled on a domain with a grid of resolution $500\times 500\times
50$, and ensembles are sampled for $60$ time steps to represent a time-varying
3D flow [9].

Although the uncertainty inherent in physical systems may be represented
through multiple simulations/sensors, the large size of the ensemble
simulations, e.g., 1.5 TB for the Red Sea dataset, can become a bottleneck for
data transmission and visualization. The recent work by Wang et al. [10]
reviewed the challenges as well as advances in visualization methods for
handling the complexity of large ensembles. They found that the most common
approach to alleviate the complexity of large ensembles is through statistical
summarizations. Visualizing statistical summaries of ensembles can help users
understand commonalities and differences in the features observed across the
ensemble members. Statistical summarizations can also help alleviate the
storage burden through reduced representation of data.

We employ two statistical summarization techniques for the visualization of
the Red Sea eddy simulation dataset, namely statistical volume rendering [8]
and statistical Morse complex summary maps [2]. We now demonstrate the results
of the application of statistical rendering techniques to the contest dataset.

##  1 Statistical 3D Volume Rendering

We employ statistical volume rendering frameworks [8, 7, 1] for 3D
visualizations of the Red Sea eddy simulation ensemble. In a statistical
volume rendering framework, per-voxel uncertainty is characterized using a
probability distribution, which is estimated from the ensemble members. The
probability distributions are then propagated through the direct volume
rendering pipeline to derive likely (expected) visualizations for the
ensemble.

We derive the expected visualizations for four noise models, namely, uniform
[8], Gaussian, Gaussian mixtures [7], and nonparametric models [1], as shown
in Fig. 1. The visualizations are derived for the velocity magnitude ensemble
with $20$ members over the domain 40∘E-50∘E and 10∘N-20∘N for the time step
$t=40$. Fig. 1a visualizes a single ensemble member using the arrow glyphs
color mapped by velocity magnitude. The high-velocity magnitude is generally
observed near the vortex rim. The transfer function shown in Fig. 1c maps the
regions with relatively high-, moderate-, and low-velocity magnitudes to red,
blue, and yellow, respectively. The same transfer function is used for all
statistical renderings in Fig. 1.

The expected visualizations derived using the uniform, Gaussian, Gaussian
mixture, and nonparametric statistical models (Fig. 1 (d-j)) appear
significantly different from the mean-field visualization (Fig. 1b). The mean
statistics exhibit high sensitivity to the outlier members, thus, they lack
reliable reconstructions of the expected vortical features for the ensemble.
In contrast, the distribution-based models display reconstructions with
relatively high resilience to the outlier members and indicate the presence of
eddies in the regions indicated by $e_{1}$, $e_{2}$, and $e_{3}$ (see Fig.
1d). As can be inferred from the statistical renderings, the eddy denoted by
$e_{1}$ can be observed across all distribution models, thus indicating a high
likelihood of its presence/position. The eddy indicated by $e_{2}$ is clearly
seen in the uniform, Gaussian, and Gaussian mixture (ordered) models, but not
in the remaining noise models. The eddy denoted by $e_{3}$ exhibits a high
level of uncertainty regarding its presence/position, as no noise model shows
a clear vortical structure in the same region.

Note that the statistical summarizations using the uniform and Gaussian noise
models consume only twice the amount of memory needed for the mean-field
statistical approach since they store mean and width/variance per voxel. We
use four Gaussians for uncertainty modeling with Gaussian mixtures (see [7]),
which means, they consume $12$ times the amount of memory needed for the mean
field (mean, variance, and weight per Gaussian). Quantile interpolation
consumes memory proportional to the number of quantiles (see [1] for more
details). The reduced data representation allows for statistical volume
rendering at interactive frame rates. (Refer to the supplementary video for
the data interaction demo.)

Fig. 2a visualizes a box-plot-like view for the velocity magnitude ensemble.
Specifically, we derive the lower quartile (lower 25%), middle quartile
(central 50%), and upper quartile (upper 25%) at each voxel of the dataset and
visualize each quartile with the uniform statistics. The quartile view [1]
gives us insight into variations in features across the three populations.
Fig. 2b analyzes the effects of sample size on nonparametric statistical
renderings. The dotted boxes in Fig. 2b illustrate the features with
relatively high sensitivity to underlying data. Fig. 3 depicts how
visualizations evolve for the time steps $t=36,\cdots,40$ for the mean and
parametric statistics.

##  2 Statistical 2D Morse Complex Summary Maps

Morse and Morse-Smale complexes are topological descriptors that provide
abstract representations of the gradient flow behavior of scalar fields [4].
We study the variability of Morse complexes for the Red Sea ensemble members
using the probabilistic maps [2] to extract the expected vortex structures as
well as to gain insight into the positional variability of expected vortex
structures. For our analysis, we use an ensemble of $10$ members, in which
each member corresponds to a 2D slice perpendicular to the z-axis ($z=1$) for
time step $40$. We again analyze the eddies over the domain 40∘E-50∘E and
10∘N-20∘N. Each ensemble member represents a velocity vector field, and Morse
complexes are computed from the negation of velocity magnitudes of each
ensemble member to focus on local minima of the vector fields. The
probabilistic map computation comprises three steps: persistence
simplification for each member, local maxima association across simplified
members via labeling, and Morse complex visualization.

### Persistence simplification.

Persistent homology is a tool in topological data analysis for quantifying the
significance of topological features. It is widely used for data de-noising
through persistence simplification [5]. We employ persistence simplification
to obtain a common label set across all ensemble members, guided by
persistence graphs and spaghetti plots in Fig. 4. In particular, at the
selected simplification scale (dotted red line) in Fig. 4a, $5$ of $10$
($50\%$) members agree on the number of maxima ($11$) after simplification.

We illustrate three ensemble members in Figs. 5a-c, respectively. For each
ensemble member, its corresponding simplified Morse complex contains 2-cells
that highlight vortical features of ocean eddies (white boxes). The mean field
Morse complex in Fig. 5d, however, does not give any insight into the
structural uncertainty, that is, the variabilities of these features across
the ensemble. The spaghetti plots of the simplified Morse complexes in Fig. 4b
do not display the topological consistency of 1-cells, thereby indicating the
high variability of simulations. For the simulations with high variability, we
benefit from the k-means and Morse mapping labeling strategies for deriving
associations among local maxima of ensemble members, as demonstrated below.

### Labeling.

In Fig. 6, we compare the three labeling strategies proposed in [2]. As
illustrated in Fig. 6d, the number of mandatory maxima [3] is small ($3$)
since ensemble members have large variations. Simplifying each ensemble member
to have $3$ maxima will miss most of the features of interest (Fig. 6e). The
Morse mapping (Fig. 6a) and the k-means clustering (Fig. 6b-c) strategies, on
the other hand, provide reasonable results. In the k-mean clustering, we set
$k=11$ since we simplified each ensemble member to contain $11$ maxima based
on the analysis of persistence graphs. The Morse mapping is more flexible than
the k-means without requiring the same number of maxima across the ensemble.

### Probabilistic map.

We visualize the probabilistic map using color blending [2] for both k-means
clustering and Morse mapping labeling strategies. Both visualizations in Fig.
7 highlight the positional uncertainty of 2-cell boundaries invisible to the
mean field of Fig. 5d. However, the expected 2-cell boundaries (black
contours) using Morse mapping appear to be more spatially stable than those
obtained via k-means clustering. The expected 2-cell boundaries extract the
expected eddy positions for the ensemble dataset. Figs. 8a-c visualize our
entropy-based exploration of the probabilistic map for lower entropy
thresholds of $1.5$, $1.25$, and $1$, respectively. Figs. 8d-f carve out
regions in the domain, where the ensemble agrees in their gradient
destinations for at least 80%, 70%, and 60% members, respectively. Thus, the
shared features denoting the eddy structures across the ensemble are
discoverable in Figs. 8d-f. In Fig. 9, the probabilistic map is again
visualized for the lower entropy threshold of $0.8$. The gradient flows
originating at the query selections $0-3$ in Fig. 9 have the highest
probability of terminating at the local maxima with green, yellow, gray, and
pink labels, respectively.

##  3 Implementation

In the case of statistical volume visualizations, the renderings are performed
on a machine with Nvidia GPU Quadro P6000, with 24 GB memory. We integrated
the fragment shaders for our statistical frameworks into the Voreen volume
rendering engine (http://voreen.uni-muenster.de) for DVR of ensemble data. In
the case of statistical Morse complex summary maps, we extend the Python code
for topological data analysis available at https://pypi.org/project/topopy/.
We provide the demo of our techniques in action in a supplementary video.

##  4 Conclusion

We demonstrate the effectiveness of statistical visualization techniques for
aggregate analysis of the Red Sea eddy simulation dataset. Specifically, we
illustrate applications of statistical volume rendering [8, 7, 1] and
statistical Morse complex summary maps [2] to extract the likely (expected)
eddy positions as well as their variability. The distribution-based data
representation in the case of statistical volume rendering allows for the
exploration of the large-scale Red Sea eddy simulation ensemble in 3D at
interactive frame rates. Additionally, the distribution-based statistics show
increased robustness to noise compared to the mean statistics and allow for an
uncertainty integration with visualizations using both statistical rendering
techniques.

###### Acknowledgements.

This work was supported in part by the NIH grants P41 GM103545-18 and R24
GM136986; the DOE grant DE-FE0031880; the Intel Graphics and Visualization
Institutes of XeLLENCE; and the NSF grants IIS-1617101, IIS-1910733,
DBI-1661375, and IIS-1513616.

## References

  * [1] T. Athawale, B. Ma, E. Sakhaee, C. R. Johnson, and A. Entezari.  Direct volume rendering with nonparametric models of uncertainty.  IEEE Transactions on Visualization and Computer Graphics, to appear, 2021.
  * [2] T. Athawale, D. Maljovec, C. R. Johnson, V. Pascucci, and B. Wang.  Uncertainty visualization of 2D Morse complex ensembles using statistical summary maps.  2019\.  arXiv:1912.06341.
  * [3] G. David, S. Joseph, and T. Julien.  Mandatory critical points of 2D uncertain scalar fields.  Computer Graphics Forum, 33(3):31–40, 2014.
  * [4] H. Edelsbrunner, J. Harer, and A. Zomorodian.  Hierarchical Morse complexes for piecewise linear 2-manifolds.  In Proceedings of the 17th Annual Symposium on Computational Geometry, pp. 70–79, 2001.
  * [5] H. Edelsbrunner, D. Letscher, and A. J. Zomorodian.  Topological persistence and simplification.  Discrete and Computational Geometry, 28:511–533, 2002.
  * [6] I. Hoteit, T. Hoar, G. Gopalakrishnan, NancyCollins, J. Anderson, B. Cornuelle, A. Köhld, and P. Heimbach.  A MITgcm/DART ocean analysis and prediction system with application to the Gulf of Mexico.  Dynamics of Atmospheres and Oceans, 63:1–23, 2013.
  * [7] S. Liu, J. Levine, P.-T. Bremer, and V. Pascucci.  Gaussian mixture model based volume visualization.  In Proceedings of the IEEE Large-Scale Data Analysis and Visualization Symposium, pp. 73–77, 2012.
  * [8] E. Sakhaee and A. Entezari.  A statistical direct volume rendering framework for visualization of uncertain data.  IEEE Transactions on VIsualization and Computer Graphics, 23(12):2509–2520, 2017.
  * [9] S. Sivareddy, H. Toye, P. Zhan, S. Langodan, G. Krokos, O. Knio, and I. Hoteit.  Impact of atmospheric and model physics perturbations on a high-resolution ensemble data assimilation system of the Red Sea.  Journal of Geophysical Research: Oceans, 2020.
  * [10] J. Wang, S. Hazarika, C. Li, and H.-W. Shen.  Visualization and visual analysis of ensemble data: A survey.  IEEE Transactions on Visualization and Computer Graphics, 25(9), 2018.

Figure 1: Statistical volume rendering of the velocity magnitude ensemble: (a)
arrow glyph visualization for a velocity vector field of a single ensemble
member colored by magnitude, (b) mean statistics, (d-e) parametric noise
models, (f) Gaussian mixture models with ordered Gaussian means, (g) Gaussian
mixture models with Monte Carlo sampling, (h-i) nonparametric density models
with quantile representation, (j) uniform noise model for the interquartile
range (central 50% population). The red, blue, and yellow in the transfer
function (c) indicate relatively high-, moderate-, and low-velocity
magnitudes. $e_{1}$, $e_{2}$, and $e_{3}$ denote the potential eddy positions
observed across different noise models. Figure 2: (a) The quartile view for
uncertainty analysis of an eddy presence in the Red Sea dataset. The solid
boxes enclose the positions that indicate the high likelihood of an eddy
presence, whereas the dotted boxes mark positions with substantial uncertainty
in an eddy presence. (b) Effect of sample size (n) on visualizations. The
dotted boxes illustrate positions with variability in reconstruction. Figure
3: Visualizations of the uncertain velocity magnitude field for a series of
time steps. The mean-field visualization exhibits relatively more fluctuations
in reconstructions compared to the uniform and Gaussian noise models.

Figure 4: Persistence simplification. (a) Persistence graphs. (b) Spaghetti
plots of the simplified Morse complexes. Figure 5: (a-c) Three ensemble
members together with (d) the mean field. Each subfigure visualizes, from left
to right, the negated velocity magnitude field (red means low and blue means
high velocity magnitude), and its corresponding Morse complexes before and
after persistence simplification. The dotted white boxes mark the vortex
features of each member. Figure 6: (a) Labeling with Morse mapping. (b-c)
Labeling with k-means clustering. (d) Mandatory maxima are shown as colored
regions. (e) Labeling with nearest mandatory maxima.  Figure 7: The
probabilistic map is visualized based on (a) k-means clustering and (b) Morse
mapping strategies. Figure 8: (a-c) Entropy-based exploration of uncertain
regions representing entropy greater than or equal to $1.5$ in (a), $1.25$ in
(b), and $1$ in (c), respectively; (d-f) visualizations of the regions that
agree in their gradient destinations for at least 80% members in (d), 70%
members in (e), and 60% members in (f), respectively. Figure 9: The
probabilistic map for the positions with entropy greater than or equal to
$0.8$. The query selections $0-3$ have the highest probability of flowing to
the maxima with green, yellow, gray, and pink labels, respectively.


# Weight Decay and SGD Provably Induce Low-Rank Bias in Neural Networks

Antiquus S. Hippocampus, Natalia Cerebro & Amelie P. Amygdale
Department of Computer Science
Cranberry-Lemon University
Pittsburgh, PA 15213, USA
<EMAIL_ADDRESS>
&Ji Q. Ren & Yevgeny LeNet
Department of Computational Neuroscience
University of the Witwatersrand
Joburg, South Africa
<EMAIL_ADDRESS>
Coauthor
Affiliation
Address
email Use footnote for providing further information about author (webpage,
alternative address)— _not_ for acknowledging funding agencies. Funding
acknowledgements go at the end of the paper.

###### 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.

##  1 Introduction

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.,
[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.,
[DBLP:journals/corr/abs-2110-06914]). The origin of the SGD noise represents a
puzzling and rather neglected question: 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)111When $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$..

redWhy is low-rank/SGD noise important?

###  1.1 Related Work

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 [DBLP:journals/corr/abs-2201-12760],
[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 [DBLP:journals/corr/abs-2012-09839].
[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.
[DBLP:journals/corr/abs-2012-09839] established evidence that the implicit
regularization in matrix factorization is a heuristic for rank minimization.
Beyond factorization problems, [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
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. [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
[DBLP:journals/corr/abs-2201-12760].

###  1.2 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:

  * •

In Thm. 1 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.

  * •

In particular, in Sec. 3.1, we theoretically show that the batch size and
weight decay have a crucial effect on the rank of the learned matrices. In
Sec. 4.3 we provide an extensive empirical study of the relationship between
the rank and certain hyperparameters.

  * •

In Props. 1-2, 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. 4.2.

##  2 Problem 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\mathcal{X}\times\mathcal{Y}$, where
$\mathcal{X}\subset\mathbb{R}^{c_{1}\times h_{1}\times w_{1}}$ is the instance
space (e.g., images), and $\mathcal{Y}\subset\mathbb{R}^{k}$ is a label space.

We consider a parametric model
$\mathcal{F}\subset\\{f^{\prime}:\mathcal{X}\to\mathbb{R}^{k}\\}$, where each
function $f_{W}\in\mathcal{F}$ is specified by a vector of parameters
$W\in\mathbb{R}^{N}$. A function $f_{W}\in\mathcal{F}$ assigns a prediction to
an input point $x\in\mathcal{X}$, and its performance on the distribution $P$
is measured by the expected risk

$L_{P}(f_{W})~{}:=~{}\mathbb{E}_{(x,y)\sim P}[\ell(f_{W}(x),y)],$ (1)

where $\ell:\mathbb{R}^{k}\times\mathcal{Y}\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 regularized
empirical risk

$L^{\lambda}_{S}(f_{W})~{}:=~{}\frac{1}{m}\sum^{m}_{i=1}\ell(f_{W}(x_{i}),y_{i})+\lambda\|W\|^{2}_{2},$
(2)

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.

Optimization.  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
$\tilde{S}=\\{(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}_{\tilde{S}}(f_{W_{t}})$, where
$\eta_{t}>0$ is our learning rate at iteration $t$. The algorithm returns
$W_{T}$.

#### 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
$\mathbb{I}:\\{\textnormal{True},\textnormal{False}\\}\to\\{0,1\\}$ the
indicator function. For a given tensor $x\in\mathbb{R}^{c\times h\times w}$,
we denote
$\textnormal{vec}(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})$.

###  2.1 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.

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}:\mathbb{R}^{c_{1}\times
h_{1}\times w_{1}}\to\mathbb{R}^{c_{i}\times h_{i}\times w_{i}}$ and each
connection holds a transformation $C^{ij}:\mathbb{R}^{c_{j}\times h_{j}\times
w_{j}}\to\mathbb{R}^{c_{i}\times h_{i}\times w_{i}}$. The layers are divided
into three categories: (i) the input layer $v_{1}$, (ii) the output layer
$v_{L}$ and (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 $\tilde{X}$, 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$.

Convolutional layers.  The networks include convolutional layers [6795724,
10.5555/646469.691875] (see also [Goodfellow-et-al-2016])
$C^{ij}:\mathbb{R}^{c_{i}\times h_{i}\times w_{i}}\to\mathbb{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\mathbb{R}^{c_{i}\times c_{j}\times k_{1}\times k_{2}}$ and computes the
following tensor as output

$y_{c,i,j}=\sum^{c_{1}}_{c^{\prime}=1}\textnormal{vec}(Z_{c^{\prime},~{}:~{},~{}:~{}})^{\top}\cdot\textnormal{vec}(\textnormal{Pad}_{p}(x)_{c^{\prime},is~{}:~{}(i+1)s+k_{1},js~{}:~{}(j+1)s+k_{2}}),$
(3)

where $(c,i,j)\in[c_{2}]\times[h_{2}]\times[w_{2}]$. Here,
$\textnormal{Pad}_{p}$ takes a tensor $x\in\mathbb{R}^{c\times h\times w}$ and
returns a new tensor $x^{\prime}\in\mathbb{R}^{c\times(h+2p)\times(w+2p)}$,
where the first and last $p$ rows and columns of each
$x^{\prime}_{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\mathbb{R}^{c_{i}\times h_{i}\times
w_{i}}:~{}W^{ij}\textnormal{vec}(x)=\textnormal{vec}(C^{ij}(x))$. This is
guaranteed since both the padding and convolution operations of the layer are
linear redexplain. We also consider $V^{ij}\in\mathbb{R}^{c_{i}\times
c_{j}k_{1}k_{2}}$ which is a matrix whose $c$’th row is the vectorized filter
$V^{ij}_{c}:=\textnormal{vec}(Z^{ij}_{c,~{}:~{},~{}:~{}})$.

As a special case, the networks may also include fully-connected layers:
$C:\mathbb{R}^{c_{1}\times 1\times 1}\times\mathbb{R}^{c_{2}\times 1\times 1}$
that take an input $x\times\mathbb{R}^{c_{1}\times 1\times 1}$ and apply a
linear transformation $\textnormal{vec}(C(x))=W\textnormal{vec}(x)$, where
$W\in\mathbb{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 $\mathbb{R}^{c_{1}\times 1\times
1}\to\mathbb{R}^{c_{2}\times 1\times 1}$ with $k_{1}=k_{2}$, $p=0$ and $s=1$.

Rearrangement layers.  A rearrangement layer $C:\mathbb{R}^{c_{1}\times
h_{1}\times w_{1}}\to\mathbb{R}^{c_{2}\times h_{2}\times w_{2}}$ takes an
input vector $x\in\mathbb{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.

Pooling Layers.  A pooling layer Zhou1988ComputationOO (see 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

$y_{c,i,j}=\max(\textnormal{Pad}_{p}(x)_{c,is~{}:~{}(i+1)s+k_{1},js~{}:~{}(j+1)s+k_{2}}),$
(4)

##  3 Theoretical Results

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\mathbb{R}^{N}$ starting from a certain
vector $W_{1}\in\mathbb{R}^{N}$ (constant), we say that it converges in mean
if we have

$\exists~{}W^{*}\in\mathbb{R}^{N}:~{}\lim_{t\to\infty}\mathbb{E}[\|W_{t}-W^{*}\|]~{}=~{}0.$
(5)

As a consequence of this definition, we have

$\lim_{t\to\infty}\mathbb{E}[\|W_{t+1}-W_{t}\|]~{}\leq~{}\lim_{t\to\infty}\mathbb{E}[\|W_{t+1}-W^{*}\|]+\lim_{t\to\infty}\mathbb{E}[\|W_{t}-W^{*}\|]~{}=~{}0.$
(6)

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

$\lim_{t\to\infty}\mathbb{E}_{\tilde{S}}\|\nabla
L^{\lambda}_{\tilde{S}}(f_{W^{t}})\|=0,$ (7)

since $\mathbb{E}[\|W^{t+1}-W^{t}\|]=\eta\cdot\mathbb{E}_{\tilde{S}}[\|\nabla
L^{\lambda}_{\tilde{S}}(f_{W^{t}})\|]$.

To analyze the step taken in each step, we write it explicitly:

$\displaystyle\nabla_{V^{pq}}L^{\lambda}_{\tilde{S}}(f_{W})~{}=~{}\frac{1}{B}\sum_{(x,y)\in\tilde{S}}\frac{\partial\ell(f_{W}(x),y)}{\partial
f_{W}(x)}\cdot\frac{\partial f_{W}(x)}{\partial V^{pq}}+2\lambda V^{pq}.$ (8)

In Sec. 3.1 we show that near-convergence, mini-batch SGD implicitly learns
neural networks with low-rank matrices. In Sec. 3.2, we prove that SGD noise
is inevitable when applying weight decay.

###  3.1 Low-Rank Bias in Neural Networks

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}$.

######  Lemma 1.

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,
$\textnormal{rank}\left(\nabla_{V^{ij}}f_{W}(x)\right)\leq N^{ij}$.

We note that a fully-connected layer
$M:\mathbb{R}^{c}\to\mathbb{R}^{c^{\prime}}$ associated with a matrix
$W\in\mathbb{R}^{c^{\prime}\times c}$ can be represented as a convolutional
layer $C:\mathbb{R}^{c\times 1\times 1}\to\mathbb{R}^{c^{\prime}\times 1\times
1}$ whose parameters tensor $Z\in\mathbb{R}^{c^{\prime}\times c\times 1\times
1}$ satisfies $Z_{a,b,1,1}=W_{a,b}$ for all $(a,b)\in[c^{\prime}]\times[c]$.
Since the number of patches in this case is $1$, the rank of the gradient
associated with $V$ is at most $1$.

redSome text

######  Theorem 1.

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,

$\min_{\begin{subarray}{c}V\in R^{d_{p}\times d_{q}}:\\\
\textnormal{rank}(V)\leq
kN^{pq}B^{*}\end{subarray}}\|V^{pq}-V\|~{}\leq~{}\frac{1}{2\lambda}\min_{\begin{subarray}{c}\tilde{S}\subset
S:\\\
|\tilde{S}|=B^{*}\end{subarray}}\|\nabla_{V^{pq}}L^{\lambda}_{\tilde{S}}(f_{W})\|$

###### Proof.

Let $\tilde{S}\subset S$ be a batch of size $B^{*}$. By the chain rule, we can
write the gradient of the loss function as follows

$\displaystyle\frac{\partial L^{\lambda}_{\tilde{S}}(f_{W})}{\partial
V^{ij}}~{}={}$
$\displaystyle\frac{1}{B^{*}}\sum^{B^{*}}_{(x,y)\in\tilde{S}}\frac{\partial\ell(f_{W}(x),y)}{\partial
f_{W}(x)}\cdot\frac{\partial f_{W}(x)}{\partial V^{ij}}+2\lambda V^{ij}.$

According to Lems. 1 and 2, the matrix $E_{\tilde{S}}:=-\frac{1}{2\lambda
B^{*}}\sum^{B^{*}}_{(x,y)\in\tilde{S}}\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

$\displaystyle\min_{V:~{}\textnormal{rank}(V)\leq
B^{*}N^{ij}}\|W^{ij}-V\|~{}\leq~{}\min_{\tilde{S}\subset
S:~{}|\tilde{S}|=B^{*}}\|W^{ij}-E_{\tilde{S}}\|~{}=~{}\frac{1}{2\lambda}\min_{\tilde{S}\subset
S:~{}|\tilde{S}|=B^{*}}\left\|\frac{\partial
L^{\lambda}_{\tilde{S}}(f_{W})}{\partial W^{ij}}\right\|,$

which completes the proof of this theorem. ∎

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.

###  3.2 Degeneracy and the Origin of SGD Noise

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}:\mathbb{R}^{n}\to\mathbb{R}$ and assume that
$\forall~{}i\in[m]:~{}x_{i}\neq 0$.

redWhy 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
$\tilde{S}\subset S$ of size $B$ for which
$\nabla_{W^{pq}}L_{\tilde{S}}(f_{W})\neq 0$.

We note that at convergence, we have $\nabla_{W^{pq}}L_{\tilde{S}}(f_{W})=0$
across all pairs $(v_{p},v_{q})\in E_{T}$ and mini-batches $\tilde{S}\subset
S$ of size $B<m$. Specifically, we we can write,

$\displaystyle 0{}$
$\displaystyle=~{}\nabla_{W^{pq}}L^{\lambda}_{\tilde{S}}(f_{W})~{}=~{}\frac{1}{B}\sum_{(x,y)\in\tilde{S}}\frac{\partial\ell(f_{W}(x),y)}{\partial
f_{W}(x)}\cdot\frac{\partial f_{W}(x)}{\partial W^{pq}}+2\lambda W^{pq}.$ (9)

Suppose we have two batches $\tilde{S}_{1},\tilde{S}_{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,

$\displaystyle 0{}$
$\displaystyle=~{}\nabla_{W^{pq}}L^{\lambda}_{\tilde{S}_{1}}(f_{W})-\nabla_{W^{pq}}L^{\lambda}_{\tilde{S}_{2}}(f_{W})$
(10) $\displaystyle=~{}\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}}.$

Therefore, we conclude that for all $i,j\in[m]$,

$\displaystyle 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}}.$ (11)

Hence, for all $(v_{p},v_{q})\in E_{T}$ and $i\in[m]$,

$\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.$ (12)

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. 11.

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.

######  Proposition 1 ($\lambda>0$).

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. 9). Then, either $f_{W}\equiv 0$
or $\\{x_{i}\\}^{m}_{i=1}$ are collinear vectors.

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.

######  Proposition 2 ($\lambda=0$).

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\mathbb{R}~{}\nexists~{}a^{*}\in\mathbb{R}:~{}\ell(a^{*},b)=\inf_{a}\ell(a,b)$,
then, $\forall~{}i\in[m]:~{}f_{W}(x_{i})=0$.

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\mathbb{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{|}\frac{\partial\ell(f_{W}(x_{i}),y_{i})}{\partial
f_{W}(x_{i})}\big{|}$ 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{|}\frac{\partial\ell(f_{W}(x_{i}),y_{i})}{\partial
f_{W}(x_{i})}\big{|}$ 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}$.

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. 11 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.

##  4 Experiments

|  |  |
---|---|---|---
|  |  |
$\eta=0.001$ | $\eta=0.01$ | $\eta=0.1$ | $\eta=0.5$
Figure 1: Average rank of MLP-BN-10-100 trained on CIFAR10 with varying batch sizes 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. |  |  |
---|---|---|---
|  |  |
$B=16$ | $B=32$ | $B=64$ | $B=128$
Figure 2: Average rank of MLP-BN-10-100 trained on CIFAR10 with varying weight decay blue($\epsilon=0.001$). See Fig. 1 for details. |  |  |
---|---|---|---
|  |  |
$\eta=0.0125$ | $\eta=0.025$ | $\eta=0.05$ | $\eta=0.1$
Figure 3: Average rank of RES-BN-5-500 trained on CIFAR10 with varying batch sizes blue($\epsilon=0.001$). See Fig. 1 for details. |  |  |
---|---|---|---
|  |  |
$\eta=0.0125$ | $\eta=0.025$ | $\eta=0.05$ | $\eta=0.1$
Figure 4: Average ranks and accuracy rates of MLP-5-500 trained on CIFAR10 with varying batch sizes blue($\epsilon=0.001$). See Fig. 1 for details. |  |  |
---|---|---|---
|  |  |
$\eta=0.01$ | $\eta=0.05$ | $\eta=0.1$ | $\eta=0.3$
Figure 5: Average ranks and accuracy rates of blueMLP-BN-10-100 (no-momentum) trained on CIFAR10 with varying batch sizes blue($\epsilon=0.001$). See Fig. 1 for details. |  |  |
---|---|---|---
|  |  |
$B=4$ | $B=8$ | $B=16$ | $B=32$
Figure 6: Average ranks and accuracy rates of blueMLP-BN-10-100 (no-momentum) trained on CIFAR10 with varying weight decay blue($\epsilon=0.001$). See Fig. 1 for details. |  |  |
---|---|---|---
|  |  |
$\eta=0.05$ | $\eta=0.1$ | $\eta=0.2$ | $\eta=0.3$
Figure 7: Average ranks and accuracy rates of blueMLP-5-500 (no-momentum) trained on CIFAR10 with varying batch sizes blue($\epsilon=0.001$). See Fig. 1 for details. |  |  |
---|---|---|---
|  |  |
$B=4$ | $B=8$ | $B=16$ | $B=32$
Figure 8: Average ranks and accuracy rates of blueMLP-5-500 (no-momentum) trained on CIFAR10 with varying weight decay blue($\epsilon=0.001$). See Fig. 1 for details. |  |  |
---|---|---|---
|  |  |
$\eta=0.5$ | $\eta=0.8$ | $\eta=1.0$ | $\eta=2.0$
Figure 9: Average ranks and accuracy rates of blueResNet-18 trained on CIFAR10 with varying batch sizes blue($\epsilon=0.001$). See Fig. 1 for details. |  |  |
---|---|---|---
|  |  |
$B=4$ | $B=8$ | $B=16$ | $B=32$
Figure 10: Average ranks and accuracy rates of blueResNet-18 trained on CIFAR10 with varying weight decay blue($\epsilon=0.001$). See Fig. 1 for details. |  |  |
---|---|---|---
|  |  |
$B=4$ | $B=8$ | $B=16$ | $B=32$
Figure 11: Average ranks and accuracy rates of blueResNet-18 trained on CIFAR10 with varying learning rate blue($\epsilon=0.001$). See Fig. 1 for details. |  |  |
---|---|---|---
|  |  |
$\eta=0.1$ | $\eta=0.2$ | $\eta=0.3$ | $\eta=0.5$
Figure 12: Average ranks and accuracy rates of blueVGG-16 trained on CIFAR10 with varying batch sizes blue($\epsilon=0.001$). See Fig. 1 for details. |  |  |
---|---|---|---
|  |  |
$B=4$ | $B=8$ | $B=16$ | $B=32$
Figure 13: Average ranks and accuracy rates of blueVGG-16 trained on CIFAR10 with varying weight decay blue($\epsilon=0.01$). See Fig. 1 for details. |  |  |
---|---|---|---
|  |  |
$B=4$ | $B=8$ | $B=16$ | $B=32$
Figure 14: Average ranks and accuracy rates of blueVGG-16 trained on CIFAR10 with varying learning rate blue($\epsilon=0.001$). See Fig. 1 for details. |  |  |  |
---|---|---|---|---
|  |  |  |
$\eta=0.004$ | $\eta=0.008$ | $\eta=0.01$ | $\eta=0.02$ | $\eta=0.04$
Figure 15: Average ranks and accuracy rates of blueViT (no-momentum) trained on CIFAR10 with varying batch sizes blue($\epsilon=0.01$). See Fig. 1 for details. |  |  |  |
---|---|---|---|---
|  |  |  |
$B=8$ | $B=16$ | $B=32$ | $B=64$ | $B=128$
Figure 16: Average ranks and accuracy rates of blueViT (no-momentum) trained
on CIFAR10 with varying weight decay blue($\epsilon=0.01$). See Fig. 1 for
details.

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 network222The plots are best
viewed when zooming into the pictures.. For each run we used one Tesla-k80 GPU
for 5-20 hours.

###  4.1 Setup

#### 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}:\mathbb{R}^{n}\to\mathbb{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, blueor 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 blue$\notin[-\epsilon,\epsilon]$, where $\epsilon$
is a small tolerance value. Training was done on the MNIST and CIFAR10
datasets.

#### Architectures.

We consider several network architectures. (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. (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_{1}z))))$, where
$W_{1},W_{2}\in\mathbb{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.
blue(iii) The third, denoted VGG-16, is the convolutional network proposed by
[DBLP:journals/corr/SimonyanZ14a], but with dropout replaced by batch
normalization layers and with only one fully-connected layer at the end. (iv)
The fourth architecture is the residual network proposed in [7780459], denoted
ResNet-18.

###  4.2 Experiments on SGD Noise

In Sec. 3.2 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,{1}\mathrm{e}{-6},{1}\mathrm{e}{-5},{1}\mathrm{e}{-4},{1}\mathrm{e}{-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}{|E_{T}|}\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. 17 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. 1, 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={1}\mathrm{e}{-4},{1}\mathrm{e}{-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. 3.2.
Namely, even though perfect convergence is impossible (see Prop. 2) 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.

###  4.3 Experiments on Rank Minimization

#### Rank, batch size and the learning rate.

blue (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={5}\mathrm{e}{-4}$ and momentum $0.9$. blue (ii) We
train instances of VGG-16 and ResNet-18 with the previous hyperparameters
except with momentum $0$.

As can be seen in Figs. 1, 3, 4, blue9, and 12 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. 3.1 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.

#### Rank minimization and weight decay.

blue (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,{1}\mathrm{e}{-4},{2}\mathrm{e}{-4},{4}\mathrm{e}{-4},{8}\mathrm{e}{-4},{6}\mathrm{e}{-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={5}\mathrm{e}{-4}$. blue (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. 2blue, 10 and 13. As can be seen, by
increasing $\lambda$ we typically impose stronger rank minimization
constraints, as predicted in Sec. 3.1. In addition, regardless of the batch
size, we lose the rank minimization constraint when training with $\lambda=0$
as predicted in Sec. 3.1.

##  5 Conclusions

redneeds 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:

  * •

SGD noise is inevitable when applying weight decay, regardless of the batch
size;

  * •

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;

  * •

Weight decay and training with mini-batches are necessary to obtain rank
minimization.

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 He_2017_ICCV, transformers NIPS2017_3f5ee243 and
convolutional networks 3a6e6cc3ac6e4326b35a6315d7cbddf7) and whether rank
minimization extends beyond the standard settings of supervised learning.

##  Appendix A Additional Experiments

MNIST, CE |  |  |  |
---|---|---|---|---
MNIST, MSE |  |  |  |
CIFAR10, CE |  |  |  |
CIFAR10, MSE |  |  |  |
| (a) | (b) | (c) | (d)
Figure 17: Convergence of MLP-5-2000 trained on MNIST and CIFAR10 with CE/ MSE
loss. In (a) we plot the averaged distance between the weight matrices at
epoch $t$ and epoch $t+1$, captured by, $\frac{1}{|E_{T}|}\sum_{(p,q)\in
E_{T}}\|W^{pq}_{t+1}-W^{pq}_{t}\|$. In (b) we plot the train accuracy rates,
in (c) we plot the averaged train loss and in (d) we plot the average rank
across the trainable matrices.

##  Appendix B Additional Architectures

#### Optimization.

rededit 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
$\tilde{S}=\\{(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}_{\tilde{S}}(f_{W_{t}}(~{}\cdot~{};\tilde{X}))$,
where $\tilde{X}=\\{x_{ij}\\}^{B}_{j=1}$ and $\eta_{t}>0$ is our learning rate
at iteration $t$. Here,
$f_{W}(~{}\cdot~{};\tilde{X}):\mathbb{R}^{n}\to\mathbb{R}$ is the function
given by $f_{W}(x;\tilde{X})$ for a fixed value of $\tilde{X}$. Furthermore,
for a given function $g_{\theta}=g^{L}_{\theta}\circ\dots
g^{1}_{\theta}:\mathbb{R}^{n}\to\mathbb{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 redpseudo gradient of $g_{\theta}$, which is computed using the
chain rule, while treating $\\{g^{i_{j}}_{\theta}\\}^{r}_{j=1}$ as independent
of $\theta$.

#### 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}:\mathbb{R}^{c_{1}\times h_{1}\times
w_{1}}\to\mathbb{R}^{c_{i}\times h_{i}\times w_{i}}$ and each connection holds
a linear transformation $C^{ij}:\mathbb{R}^{c_{j}\times h_{j}\times
w_{j}}\to\mathbb{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 $\tilde{X}$, the output of a given layer $v_{i}$ is evaluated as
follows

$v_{i}(x;\tilde{X})~{}:=~{}\sigma\left(n_{i}\left[\sum_{j\in\textnormal{pred}(i)}C^{ij}(v_{j}(x;\tilde{X}));\tilde{X}\right]\right),$

except for the output layer $v_{L}$ that computes

$f_{W}(x;\tilde{X})~{}:=~{}v_{L}(x;\tilde{X})~{}:=~{}\sum_{j\in\textnormal{pred}(L)}C^{Lj}(v_{j}(x;\tilde{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^{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$.

#### Convolutional layers.

redadd a formal def maybe? Throughout the paper we consider convolutional
layers 6795724,10.5555/646469.691875 (see also GoodBengCour16)
$C:\mathbb{R}^{c_{1}\times h_{1}\times w_{1}}\to\mathbb{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_{2}w_{2}$. For a given
convolutional layer $C^{pq}$ in the network $f_{W}$, we denote
$N^{pq}:=h_{p}w_{p}$. For a given convolutional layer $C^{pq}$ we define a
matrix $W^{pq}$ that computes $\forall~{}x\in\mathbb{R}^{c_{q}\times
h_{q}\times w_{q}}:~{}W^{pq}\textnormal{vec}(x)=\textnormal{vec}(C^{pq}(x))$.
This is guaranteed since both the padding and convolution operations of the
layer are linear. We also consider $V^{pq}\in\mathbb{R}^{c_{p}\times
c_{q}k_{1}k_{2}}$ which is a matrix whose rows are
$\textnormal{vec}(V^{pq}_{c,:,:})$.

#### Pooling Layers.

A pooling layer (see 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.

#### Batch normalization layers.

A batch normalization layer pmlr-v37-ioffe15 is defined in the following
manner,

$\displaystyle n_{l}(v_{l}(x;\tilde{X});\tilde{X}){}$
$\displaystyle=~{}\gamma_{l}\cdot
A_{l}(\tilde{X})\odot(v_{l}(x;\tilde{X})-b_{l}(\tilde{X}))$
$\displaystyle\quad+\mathbbm{1}_{c\times h_{l}\times w_{l}}\cdot\beta_{l}$
$\displaystyle A_{l}(\tilde{X}){}$
$\displaystyle=~{}\left(\left(\sqrt{\mathrm{Var}_{\hat{x}\sim
U[\tilde{X}]}(v_{l}(\hat{x};\tilde{X})_{j})}\right)^{-1}\right)_{j}$
$\displaystyle b_{l}(\tilde{X}){}$ $\displaystyle=~{}\mathbb{E}_{\hat{x}\sim
U[\tilde{X}]}[v_{l}(\hat{x};\tilde{X})],$

where $v_{l}(\hat{x};\tilde{X})_{j}$ is the $j$’th coordinate of
$v_{l}(\hat{x};\tilde{X})$ for $j\in[c_{l}]\times[h_{l}]\times[w_{l}]$. In
addition, $\gamma_{l},\beta_{l}\in\mathbb{R}$ are trainable scale and shift
parameters and $\mathbbm{1}_{a_{1}\times\dots\times
a_{k}}\in\mathbb{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 NEURIPS2020_f3ada80d) the matrices $A_{l}(\tilde{X})$ during
backpropagation at any step during training. Namely, we treat
$\frac{\bar{\partial}A_{l}(\tilde{X})}{\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
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, [NEURIPS2019_2f4fe03d] showed that detaching the variance term of
layer normalization is competitive and even improves the performance on
various tasks.

##  Appendix C Proofs

See 1

###### Proof.

Let $x\in\mathbb{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
$\textnormal{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,

$f_{W}(x)~{}=~{}\sum_{l_{1}\in\textnormal{pred}(l_{0})}C^{l_{0}l_{1}}\circ
v_{l_{1}}(x),$

where $l_{0}=L$ and $C^{l_{0}l_{1}}\circ z:=C^{l_{0}l_{1}}(z)$. In addition,
each layer $v_{l}$ can also be written as

$\textnormal{vec}(v_{l_{1}}(x))~{}=~{}D_{l_{1}}\odot\sum_{l_{2}\in\textnormal{pred}(l_{1})}C^{l_{1}l_{2}}\circ
v_{l_{2}}(x),$

where $D_{l}:=D_{l}(x):=\sigma^{\prime}(v_{l}(x)))\in\mathbb{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

$\displaystyle f_{W}(x){}$ $\displaystyle=~{}\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)$
$\displaystyle\quad\quad+\sum_{\begin{subarray}{c}\pi\textnormal{ from $1$ to
$l_{0}$}\\\ (i,j)\notin\pi\end{subarray}}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,$ $\displaystyle=:~{}A_{W}(x)+B_{W}(x)$

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\mathbb{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.

#### 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\mathbb{R}^{c_{j}\times h_{j}\times w_{j}}$ as a vector
$\textnormal{vec}^{ij}(x)\in\mathbb{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\mathbb{R}^{c_{i}\times(c_{j}k_{1}k_{2})}$, whose $c$’th row is
$\textnormal{vec}(Z^{ij}_{c,~{}:,~{}:~{}})$. We can write the vectorized
output of the convolutional layer as $U^{ij}\cdot\textnormal{vec}^{ij}(x)$,
where

$U^{ij}:=\left(\begin{array}[]{ c c c c c c
}\lx@intercol\hfil\hfil\lx@intercol&&0&0&0&0\\\
\lx@intercol\hfil\raisebox{7.20007pt}[0.0pt][0.0pt]{$V^{ij}$}\hfil\lx@intercol&0&0&0&0\\\
0&0&&&0&0\\\
0&0&\lx@intercol\hfil\raisebox{7.20007pt}[0.0pt][0.0pt]{$\ddots$}\hfil\lx@intercol&0&0\\\
0&0&0&0&&\lx@intercol\hfil\hfil\lx@intercol\\\
0&0&0&0&\lx@intercol\hfil\raisebox{7.20007pt}[0.0pt][0.0pt]{$V^{ij}$}\hfil\lx@intercol\end{array}\right)$
(13)

is a $(N^{ij}c_{i})\times(N^{ij}c_{j}k_{1}k_{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 $\textnormal{vec}(x)\in\mathbb{R}^{c_{j}k_{1}k_{2}}$ of
its input $x$. Since $\textnormal{vec}^{ij}(x)$ consists of the same variables
as in $\textnormal{vec}(x)$ with potentially duplicate items, there is a
linear transformation that translates $\textnormal{vec}(x)$ into
$\textnormal{vec}^{ij}(x)$. Therefore, we can simply write
$W^{ij}\cdot\textnormal{vec}(x)=U^{ij}\cdot\textnormal{vec}^{ij}(x)$.

Representing convolutional layers.  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
$\textnormal{vec}(C^{ld}(x))=W^{ld}\textnormal{vec}(x)$, for any input
$x\in\mathbb{R}^{c_{d}\times h_{d}\times w_{d}}$.

Representing pooling and rearrangement layers.  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
$\textnormal{vec}(C^{ld}(x))=W^{ld}\textnormal{vec}(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.

Computing the rank.  Finally, we note that $\textnormal{vec}(C^{ij}\circ
z)=U^{ij}\cdot\textnormal{vec}^{ij}(z)=W^{ij}\cdot\textnormal{vec}(z)$,
$\textnormal{vec}(D_{l}\odot z)=P_{l}\cdot\textnormal{vec}(z)$ for
$P_{l}:=\textnormal{diag}(\textnormal{vec}(D_{l}))$. Therefore, we can write

$\displaystyle A_{W}(x){}$ $\displaystyle=~{}\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\textnormal{vec}^{ij}(v_{j}(x))$
$\displaystyle=:~{}a(x)^{\top}\cdot U^{ij}\cdot b(x),$

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):=\textnormal{vec}^{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 equation 13 and the definition of
$U^{ij}$, we can write

$\displaystyle 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),$

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,

$\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}.$ (14)

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}$. ∎

######  Lemma 2.

Let $\tilde{f}_{W}$ be a convolutional neural network with batch-
normalization. Let $C^{pq}$ be a given convolutional layer within
$\tilde{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\mathbb{R}^{c_{q}\times(c_{p}k_{1}k_{2})}$ be the parameters matrix
of $C^{pq}$. Let $\tilde{X}=\\{x_{i}\\}^{B}_{i=1}$ be a batch of samples.
Then, the ranks of $\frac{\partial f_{W}(x;\tilde{X})}{\partial V^{pq}}$ and
$\sum^{B}_{i=1}\alpha_{i}\frac{\partial f_{W}(x_{i};\tilde{X})}{\partial
V^{pq}}$ are at most $BN^{pq}$.

###### Proof.

We would like to show that
$\frac{\bar{\partial}f_{W}(x_{i};\tilde{X})}{\partial V^{pq}}$ and
$\sum^{B}_{i=1}\beta_{i}\cdot\frac{\bar{\partial}f_{W}(x_{i};\tilde{X})}{\partial
V^{pq}}$ are matrices of rank $\leq N^{pq}B$. We follow the proof of Lem. 2 in
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. 1. 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\mathbb{R}^{c_{1}\times h_{1}\times
w_{1}}$, the network can be written as follows

$f(x;\tilde{X})~{}=~{}\sum_{l_{1}\in\textnormal{pred}(L)}C^{Ll_{1}}(v_{l_{1}}(x;\tilde{X})).$

In addition, each layer $v_{l_{1}}$ can be written as follows

$\displaystyle v_{l_{1}}(x;\tilde{X}){}$
$\displaystyle=~{}D_{l_{1}}(x;\tilde{X})\odot
n_{l_{1}}(u_{l_{1}}(x;\tilde{X});\tilde{X}),$

where
$D_{l_{1}}(x;\tilde{X}):=\sigma^{\prime}(n_{l_{1}}(u_{l_{1}}(x;\tilde{X});\tilde{X}))$
and
$u_{l_{1}}(x;\tilde{X}):=\sum_{l_{2}\in\textnormal{pred}(l_{1})}C^{l_{1}l_{2}}(v_{l_{2}}(x;\tilde{X}))$.
Therefore, we can write

$n_{l_{1}}(u_{l_{1}}(x;\tilde{X});\tilde{X})~{}=~{}\gamma_{l_{1}}\cdot
A_{l_{1}}(\tilde{X})\odot\left(u_{l_{1}}(x;\tilde{X})-b_{l_{1}}(\tilde{X})\right)+\mathbbm{1}_{c_{l_{1}}\times
h_{l_{1}}\times w_{l_{1}}}\cdot\beta_{l_{1}},$

where $b_{l}(\tilde{X}):=\mathbb{E}_{\hat{x}\sim
U[\tilde{X}]}[u_{l}(\hat{x};\tilde{X})]=\sum_{t\in\textnormal{pred}(l)}C^{lt}(\mu_{t}(\tilde{X}))$
and $\mu_{t}(\tilde{X}):=\mathbb{E}_{x\sim U[\tilde{X}]}[v_{t}(x;\tilde{X})]$.
Hence, we can write

$\displaystyle v_{l_{1}}(x;\tilde{X}){}$
$\displaystyle=~{}\sum_{l_{2}\in\textnormal{pred}(l_{1})}Q_{l_{1}}(x;\tilde{X})\odot\left[C^{l_{1}l_{2}}(v_{l_{2}}(x;\tilde{X}))-C^{l_{1}l_{2}}(\mu_{l_{2}}(\tilde{X}))\right]+D_{l_{1}}(x;\tilde{X})\cdot\beta_{l_{1}}$
$\displaystyle=~{}\sum_{l_{2}\in\textnormal{pred}(l_{1})}Q_{l_{1}}(x;\tilde{X})\odot
C^{l_{1}l_{2}}\left(v_{l_{2}}(x;\tilde{X})-\mu_{l_{2}}(\tilde{X})\right)+D_{l_{1}}(x;\tilde{X})\cdot\beta_{l_{1}},$

where $Q_{l}(x;\tilde{X}):=D_{l}(x;\tilde{X})\odot A_{l}(\tilde{X})$ and the
last equation follows from the linearity of convolutional layers.

We denote $\mu_{l}:=\mu_{l}(\tilde{X})$, $v^{j}_{l}:=v_{l}(x_{j};\tilde{X})$,
$D^{j}_{l}:=D_{l}(x_{j};\tilde{X})$, $A_{l}:=A_{l}(\tilde{X})$,
$Q^{j}_{l}:=Q_{l}(x_{j};\tilde{X})$ and define

$\Delta^{j}_{l}:=\left\\{\begin{array}[]{lll}x_{j}&\mbox{if }l=1\\\
v^{j}_{l}-\mu_{l}&\mbox{if }l\neq 1\\\ \end{array}\right..$

Hence, we can write $f(x_{j_{0}};\tilde{X})$ as a follows

$\displaystyle f(x_{j_{0}};\tilde{X}){}$
$\displaystyle=~{}\sum_{l_{1}\in\textnormal{pred}(l_{0})}C^{l_{0}l_{1}}\cdot
v^{j_{0}}_{l_{1}}$
$\displaystyle=~{}\sum_{l_{1}\in\textnormal{pred}(l_{0})}C^{l_{0}l_{1}}\left(Q^{j_{0}}_{l_{1}}\odot\left(\sum_{l_{2}\in\textnormal{pred}(l_{1})}C^{l_{1}l_{2}}(v^{j_{0}}_{l_{2}}-\mu_{l_{2}})\right)+D^{j_{0}}_{l_{1}}\cdot\beta_{l_{1}}\right)$
$\displaystyle=~{}\sum_{l_{1}\in\textnormal{pred}(l_{0})}C^{l_{0}l_{1}}(D^{j_{0}}_{l_{1}}\cdot\beta_{l_{1}})$
$\displaystyle\quad+\sum_{l_{1}\in\textnormal{pred}(l_{0})}\sum_{l_{2}\in\textnormal{pred}(l_{1})}C^{l_{0}l_{1}}(Q^{j_{0}}_{l_{1}}\odot
C^{l_{1}l_{2}}(v^{j_{0}}_{l_{2}}-\mu_{l_{2}}))$
$\displaystyle=~{}\sum_{l_{1}\in\textnormal{pred}(l_{0})}C^{l_{0}l_{1}}(D^{j_{0}}_{l_{1}}\cdot\beta_{l_{1}})$
$\displaystyle\quad+\sum_{l_{1}\in\textnormal{pred}(l_{0})}\sum_{l_{2}\in\textnormal{pred}(l_{1})}C^{l_{0}l_{1}}(Q^{j_{0}}_{l_{1}}\odot
C^{l_{1}l_{2}}(\Delta^{j_{0}}_{l_{2}}))$

For any $l_{t+2}\neq 1$, we can write

$\displaystyle\Delta^{j_{t}}_{l_{t+2}}{}$
$\displaystyle=~{}v^{j_{t}}_{l_{t+2}}-\mu_{l_{t+2}}$ (15)
$\displaystyle=~{}\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}}$
$\displaystyle\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)$
$\displaystyle\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)$
$\displaystyle=~{}\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}}),$

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}$,

$\displaystyle f_{W}(x_{j_{0}};\tilde{X}){}$
$\displaystyle=~{}\sum_{l_{1}\in\textnormal{pred}(l_{0})}C^{l_{0}l_{1}}\circ
D^{j_{0}}_{l_{1}}\cdot\beta_{l_{1}}$
$\displaystyle\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_{0}l_{1}}\circ
Q^{j_{0}}_{l_{1}}\odot
C^{l_{1}l_{2}}\circ\sum^{B}_{j_{1}=1}\alpha_{j_{0},j_{1}}Q^{j_{1}}_{l_{2}}\odot
C^{l_{2}l_{3}}$ $\displaystyle\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}}$
$\displaystyle\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_{0}l_{1}}\circ
Q^{j_{0}}_{l_{1}}\odot
C^{l_{1}l_{2}}\circ\sum^{B}_{j_{1}=1}\alpha_{j_{0},j_{1}}Q^{j_{1}}_{l_{2}}\odot
C^{l_{2}l_{3}}$ $\displaystyle\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}}.$

Instead of recursively expanding $f_{W}$’s expression $T$ times by rewriting
$\Delta^{j_{t}}_{l_{t+2}}$ using equation 15, 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}};\tilde{X})$ 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$,

$E^{j_{0}}_{1,\pi}~{}:=~{}\left\\{\begin{array}[]{lll}C^{\pi_{2}\pi_{1}}\circ
D^{j_{0}}_{\pi_{1}}\cdot\beta_{\pi_{1}}&\mbox{if }T=2\\\ \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}&\mbox{if }T>2\\\
\end{array}\right.$

and also

$E^{j_{0}}_{2,\pi}~{}:=~{}\left\\{\begin{array}[]{lll}\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}&\mbox{if
}\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}&\mbox{if }\pi_{1}=1\\\
\end{array}\right.$

With this terminology we can write

$\displaystyle f_{W}(x_{j_{0}};\tilde{X}){}$
$\displaystyle=~{}\sum_{\begin{subarray}{c}\textnormal{$\pi$ from some node
$\neq 1$ to $L$}\\\
\textnormal{$(p,q)\notin\pi$}\end{subarray}}E^{j_{0}}_{1,\pi}+\sum_{\begin{subarray}{c}\textnormal{$\pi$
from 1 to $L$}\\\
\textnormal{$(p,q)\notin\pi$}\end{subarray}}E^{j_{0}}_{2,\pi}+\sum_{\begin{subarray}{c}\pi\textnormal{
from $q$ to $L$}\\\
\pi_{2}=p,\pi_{1}=q\end{subarray}}E^{j_{0}}_{2,\pi}~{}=:~{}Z^{j_{0}}_{1}+Z^{j_{0}}_{2}+Z^{j_{0}}_{3}.$

We note that for any $x$ and $\tilde{X}$, with measure $1$ over $W$, the
matrices $D_{l}(x;\tilde{X})$ are constant in the neighborhood of $W$.
Together with the assumption that $A_{l}(x;\tilde{X})$ 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}};\tilde{X})}{\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

$\displaystyle Z^{j_{0}}_{3}{}$
$\displaystyle=~{}\sum_{\begin{subarray}{c}\pi\textnormal{ from $p$ to $L$}\\\
\pi_{2}=p,\pi_{1}=q\end{subarray}}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}}$
$\displaystyle=~{}\sum_{\begin{subarray}{c}\pi\textnormal{ from $p$ to $L$}\\\
\pi_{2}=p,\pi_{1}=q\end{subarray}}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}.$

As a next step we linearize the above expression. We notice that
$\textnormal{vec}(C^{ij}\circ
z)=U^{ij}\cdot\textnormal{vec}^{ij}(z)=W^{ij}\cdot\textnormal{vec}(z)$ and
$\textnormal{vec}(Q^{j}_{l}\odot
z)=\textnormal{diag}(\textnormal{vec}(Q^{j}_{l}))\cdot\textnormal{vec}(z)$.
Therefore, we replace $C^{pq}$ by $U^{pq}$ and $\Delta^{j_{T-3}}_{q}$ with
$\textnormal{vec}^{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}(\textnormal{vec}(Q^{j}_{l}))$. We obtain that

$\displaystyle Z^{j_{0}}_{3}{}$
$\displaystyle=~{}\sum_{\begin{subarray}{c}\pi\textnormal{ from $p$ to $L$}\\\
\pi_{2}=p,\pi_{1}=q\end{subarray}}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\textnormal{vec}^{pq}(\Delta^{j_{T-3}}_{q})$
$\displaystyle=~{}\sum^{B}_{j_{T-3}=1}(z^{1}_{j_{0},j_{T-3}})^{\top}\cdot
U^{pq}\cdot z^{2}_{j_{T-3}},$

where

$\displaystyle(z^{1}_{j_{0},j_{T-3}})^{\top}~{}:$
$\displaystyle=~{}\sum_{\begin{subarray}{c}\pi\textnormal{ from $p$ to $L$}\\\
\pi_{2}=p,\pi_{1}=q\end{subarray}}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}}$
$\displaystyle z^{2}_{j_{T-3}}~{}:$
$\displaystyle=~{}\textnormal{vec}^{pq}(\Delta^{j_{T-3}}_{q}).$

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

$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}},$

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

$\frac{\bar{\partial}f(x_{j_{0}};\tilde{X})}{\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}}$

which is a matrix of rank $\leq BN^{pq}$.

We also notice that $b_{i,j_{T-3}}$ are independent of $j_{0}$. Therefore, we
obtain that

$\displaystyle\sum^{B}_{j_{0}=1}\alpha_{j_{0}}\cdot\frac{\bar{\partial}f_{W}(x_{j_{0}};\tilde{X})}{\partial
V^{pq}}{}$
$\displaystyle=~{}\sum^{B}_{j_{0}=1}\alpha_{j_{0}}\cdot\frac{\bar{\partial}Z^{j_{0}}_{3}}{\partial
V^{pq}}$
$\displaystyle=~{}\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}}$
$\displaystyle=~{}\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}},$

which is also a matrix of rank $\leq BN^{pq}$. ∎