---
abstract: 'The mini-batch stochastic gradient descent (SGD) algorithm is widely used in training machine learning models, in particular deep learning models. We study SGD dynamics under linear regression and two-layer linear networks, with an easy extension to deeper linear networks, by focusing on the variance of the gradients, which is the first study of this nature. In the linear regression case, we show that in each iteration the norm of the gradient is a decreasing function of the mini-batch size $b$ and thus the variance of the stochastic gradient estimator is a decreasing function of $b$. For deep neural networks with $L_2$ loss we show that the variance of the gradient is a polynomial in $1/b$. The results back the important intuition that smaller batch sizes yield lower loss function values which is a common believe among the researchers. The proof techniques exhibit a relationship between stochastic gradient estimators and initial weights, which is useful for further research on the dynamics of SGD. We empirically provide further insights to our results on various datasets and commonly used deep network structures.'
bibliography:
- 'ref.bib'
---

Introduction {#sec:1}
============

Deep learning models have achieved great success in a variety of tasks including natural language processing, computer vision, and reinforcement learning [@goodfellow2016deep]. Despite their practical success, there are only limited studies of the theoretical properties of deep learning; see survey papers [@sun2019optimization; @fan2019selective] and references therein. The general problem underlying deep learning models is to optimize (minimize) a loss function, defined by the deviation of model predictions on data samples from the corresponding true labels. The prevailing method to train deep learning models is the mini-batch stochastic gradient descent (SGD) algorithm and its variants [@bottou1998online; @bottou2018optimization]. SGD updates model parameters by calculating a stochastic approximation of the full gradient of the loss function, based on a random selected subset of the training samples called a mini-batch.

It is well-accepted that selecting a large mini-batch size reduces the training time of deep learning models, as computation on large mini-batches can be better parallelized on processing units. For example, Goyal et. al. [@goyal2017accurate] scale ResNet-50 [@he2016deep] from a mini-batch size of 256 images and training time of 29 hours, to a larger mini-batch size of 8,192 images. Their training achieves the same level of accuracy while reducing the training time to one hour. However, noted by many researchers, larger mini-batch sizes suffer from a worse generalization ability [@lecun2012efficient; @keskar2019large]. Therefore, many efforts have been made to develop specialized training procedures that achieve good generalization using large mini-batch sizes [@hoffer2017train; @goyal2017accurate]. Smaller batch sizes have the advantage of allegedly offering better generalization (at the expense of a higher training time).

We hypothesize that smaller sizes lead to lower training loss and, unfortunately, decrease stability of the algorithm. The latter follows from the fact that the smaller is the batch size, more stochasticity and volatility is introduced. After all, if the batch size equals to the number of samples, there is no stochasticity in the algorithm. To this end, we conjecture that the variance of the gradient in each iteration is a decreasing function of the mini-batch size. The conjecture is the focus of the work herein. We are able to prove it in the convex linear regression case and to show significant progress in a two layer neural network setting with samples based on a normal distribution. In this case we show that the variance is a polynomial in the reciprocal of the mini-batch size and that it is decreasing for large enough mini-batch sizes. The increased variance as the mini-batch size decreases should also intuitively imply convergence to lower training loss values and in turn better prediction and generalization ability (these relationships are yet to be confirmed analytically; but we provide empirical evidence to their validity).

Another line of research focuses on how to choose an optimal mini-batch size based on different criteria [@smith2017bayesian; @gower2019sgd]. However, these papers make strong assumptions on the loss function properties (strong or point or quasi convexity, or constant variance near stationary points) or about the formulation of the SGD algorithm (continuous time interpretation by means of differential equations). The statements are approximate in nature and thus not mathematical claims. They also focus on convergence and generalization while our goal is variance. The theoretical results regarding the relationship between the mini-batch size and the performance (variance, loss, generalization ability, etc.) of the SGD algorithm applied to general machine learning models are still missing. The work herein partially addresses this gap by showing the impact of the mini-batch size on the variance of gradients in SGD.

In the linear regression case, we show that in each iteration the norm of any linear combination of sample-wise gradients is a decreasing function of the mini-batch size $b$. As a special case, the variance of the stochastic gradient estimator and the full gradient at the iterate in step $t$ are also decreasing functions of $b$ at any iteration step $t$. In addition, the proof provides a recursive relationship between the norm of gradients and the model parameters at each iteration. This recursive relationship can be used to calculate any quantity related to the stochastic gradient or full gradient at any iteration with respect to the initial weights. We give structural results and not explicit formulas which are impossible to obtain. For the two-layer linear neural network with $L_2$-loss and samples drawn from a normal distribution, we show that in each iteration step $t$ the trace of any product of the stochastic gradient estimators and weight matrices is a polynomial in $1/b$ with coefficients a sum of products of the initial weights. As a special case, the variance of the stochastic gradient estimator is a polynomial in $1/b$ without the constant term and therefore it is a decreasing function of $b$ when $b$ is large enough. The results can be easily extended to general deep linear networks. As a comparison, other papers that study theoretical properties of two-layer networks either fix one layer of the network, or assume the over-parameterized property of the model and they study convergence, while our paper makes no such assumptions on the model and we study variance with respect to the mini-batch size. The proof also reveals the structure of the coefficients of the polynomial, and thus serving as a tool for future work on proving other properties of the stochastic gradient estimators.

The proofs are involved and require several key ideas. The main one is to show a more general result than it is necessary in order to carry out the induction. The induction is not only on time step $t$ but also on the batch size with the latter one being tricky to handle. New concepts and definitions are introduced in order to handle the more general case. Along the way we show a result of general interest establishing expectation of several rank one matrices sampled from a normal distribution intertwined with constant matrices.

In conclusion, we study the dynamics of SGD under linear regression and a two-layer linear network setting by focusing on the decreasing property of the variance of stochastic gradient estimators with respect to the mini-batch size. The proof techniques can also be used to derive other properties of the SGD dynamics in regard to the mini-batch size and initial weights. To the best of authors’ knowledge, the work is the first one to theoretically study the impact of the mini-batch size on the variance of the gradient, under mild assumptions on the network and the loss function. We support our theoretical results by experiments. We further experiment on other state-of-the-art deep learning models and datasets to empirically show the validity of the conjectures about the impact of mini-batch size on average loss, average accuracy and the generalization ability of the model.

The major contributions of this paper are as follows.

-   For linear regression, we show that the norm of any number of linear combinations of the coordinates of the gradient is a decreasing function of the mini-batch size (Theorem \[thm:2\]). As a special case, the variance of the stochastic gradient estimators is also a decreasing function of the mini-batch size, for all iterations and all choices of learning rates (Corollary \[col:1\]) that are independent of the mini-batch size.

-   For a two-layer linear network, we show that any non-negative trace of the product of weight matrices and stochastic gradient estimators is a decreasing function of the mini-batch size for a large enough value. Here samples are drawn from a normal distribution. As a special case, the variance of the stochastic gradient estimators is also a decreasing function for large enough mini-batch size, for all iterations and all choices of learning rates (Theorem \[thm:varianceRepresentation\]) that are independent of the mini-batch size. The proof can be easily extended to more than two layers.

-   In the two-layer network we also show that the variance is a polynomial in $1/b$. In order to establish all of the results we design a new proof technique where the main idea is to show a more general result than only considering variance in order to apply induction in a non-trivial way.

-   We verify the theoretical results on various datasets and provide further understanding. We further empirically show that the results extend to other widely used network structures and hold for all choices of the mini-batch sizes. We also empirically verify that, on average, in each iteration the loss function value and the generalization ability (measured by the gap between accuracy on the training and test sets) are all decreasing functions of the mini-batch size.

The rest of the manuscript is structured as follows. In Section \[sec:2\] we review the literature while in Section \[sec:3\] we present the theoretical results on how mini-batch sizes impact the variance of stochastic gradient estimators, under different models including linear regression and deep linear networks. Section \[sec:4\] introduces the experiments that verify our theorems and provide further insights into the impact of the mini-batch sizes on SGD performance. We defer the proofs of the theorems and other technical details to Appendix A and experimental details to Appendix B.

Literature Review {#sec:2}
=================

Stochastic gradient descent type methods are broadly used in machine learning [@bottou1991stochastic; @lecun1998gradient; @bottou2018optimization]. The performance of SGD highly relies on the choice of the mini-batch size. It has been widely observed that choosing a large mini-batch size to train deep neural networks appears to deteriorate generalization [@lecun2012efficient]. This phenomenon exists even if the models are trained without any budget or limits, until the loss function value ceases to improve [@keskar2019large]. One explanation for this phenomenon is that large mini-batch SGD produces “sharp” minima that generalize worse [@hochreiter1997flat; @keskar2019large]. Specialized training procedures to achieve good performance with large mini-batch sizes have also been proposed [@hoffer2017train; @goyal2017accurate].

It is well-known that SGD has a slow asymptotic rate of convergence due to its inherent variance [@nesterov2013introductory]. Variants of SGD that can reduce the variance of the stochastic gradient estimator, which yield faster convergence, have also been suggested. The use of the information of full gradients to provide variance control for stochastic gradients is addressed in [@johnson2013accelerating; @roux2012stochastic; @shalev2013stochastic]. The works in [@lei2017non; @li2014efficient; @schmidt2017minimizing] further improve the efficiency and complexity of the algorithm by carefully controling the variance.

There is prior work focusing on studying the dynamics of SGD. Neelakantan et. al. [@neelakantan2015adding] propose to add isotropic white noise to the full gradient to study the “structured” variance. The works in [@pmlr-v70-li17f; @mandt2017stochastic; @jastrzkebski2017three] connect SGD with stochastic differential equations to explain the property of converged minima and generalization ability of the model. Smith and Le [@smith2017bayesian] propose an “optimal” mini-batch size which maximizes the test set accuracy by a Bayesian approach. The Stochastic Gradient Langevin Dynamics (SGLD, a variant of SGD) algorithm for non-convex optimization is studied in [@zhang2017hitting; @mou2018generalization].

In most of the prior work about the convergence of SGD, it is assumed that the variance of stochastic gradient estimators is upper-bounded by a linear function of the norm of the full gradient, e.g. Assumption 4.3 in [@bottou2018optimization]. One exception is [@gower2019sgd] which gives more precise bounds of the variance under different sampling methods. These bounds are still dependent on the model parameters at the corresponding iteration. To the best of the authors’ knowledge, there is no existing result connecting the variance of stochastic gradient estimators with the initial weights and the mini-batch size. This paper partially solves this problem.

Analysis {#sec:3}
========

Mini-batch SGD is a lighter-weight version of gradient descent. Suppose that we are given a loss function $L(w)$ where $w$ is the collection (vector, matrix, or tensor) of all model parameters. At each iteration $t$, instead of computing the full gradient $\grad_w L(w_t)$, SGD randomly samples a mini-batch set $\calB_t$ that consists of $b=|\calB_t|$ training instances and sets $$w_{t+1} \gets w_t - \alpha_t \grad_w L_{\calB_t}(w_t),$$ where the positive scalar $\alpha_t$ is the learning rate (or step size) and $\grad_w L_{\calB_t}(w_t)$ denotes the stochastic gradient estimator based on mini-batch $\calB_t$.

An important property of the stochastic gradient estimator $\grad_w L_{\calB_t}(w_t)$ is that it is an unbiased estimator, i.e. $\mathbb{E} \grad_w L_{\calB_t}(w_t) = \grad_w L(w_t)$, where the expectation is taken over all possible choices of mini-batch $\calB_t$. However, it is unclear what is the value of $$\var \pth{\grad_w L_{\calB_t}(w_t)} \triangleq \mathbb{E} \norm{\grad_w L_{\calB_t}(w_t)}^2 - \norm{\mathbb{E} \grad_w L_{\calB_t}(w_t)}^2.$$

Intuitively, we should have $$\label{eq:sgVariancePropto}
        \var \pth{\grad_w L_{\calB_t}(w_t)} \; \propto \; \frac{n^2}{b} \var \pth{\grad_w L(w_t)}$$ where $n$ is the number of training samples and stochasticity on the right-hand side comes from mini-batch samples behind $w_t$. The works in [@smith2017bayesian; @gower2019sgd] also point out this relationship, but a rigorous proof is missing. In addition, even the quantities $\grad_w L(w_t)$ and $\var \pth{\grad_w L(w_t)}$ are still challenging to compute as we do not have direct formulas of their precise values. Besides, as we choose different $b$’s, their values are not comparable as we end up with different $w_t$’s.

A plausible idea to address these issues is to represent $\mathbb{E} \grad_w L_{\calB_t}(w_t)$ and $\var \pth{\grad_w L_{\calB_t}(w_t)}$ using the fixed and known quantities $w_0, b, t$, and $\alpha_t$. In this way, we can further discover the properties, like decreasing with respect to $b$, of $\mathbb{E} \grad_w L_{\calB_t}(w_t)$ and $\var \pth{\grad_w L_{\calB_t}(w_t)}$. The biggest challenge is how to connect the quantities in iteration $t$ with those of iteration $0$. This is similar to discovering the properties of a stochastic differential equation at time $t$ given only the dynamics of the stochastic differential equation and the initial point.

In this section, we address these questions under two settings: linear regression and a deep linear network. In Section \[subsec:linearRegression\] with a linear regression setting, we provide explicit formulas for calculating any norm of the linear combination of sample-wise gradients. We therefore show that the $\var \pth{\grad_w L_{\calB_t}(w_t)}$ is a decreasing function of the mini-batch size $b$. In Section \[subsec:linearNetwork\] with a deep linear network setting and samples drawn from a normal distribution, we show that any trace of the product of weight matrices and stochastic gradient estimators is a polynomial in $1/b$ with finite degree. We further prove that $\var \pth{\grad_w L_{\calB_t}(w_t)}$ is a decreasing function of the mini-batch size $b > b_0$ for some constant $b_0$.

For a random matrix $M$, we define $\var\pth{M} \triangleq \Expect \norm{\text{vec}(M)}^2 - \norm{\Expect\text{vec}(M)}^2$ where $\textrm{vec}(M)$ denotes the vectorization of matrix $M$. We denote $\fromto{m}{n} \triangleq \{m, m+1, \ldots,n\}$ if $m\le n$, and $\emptyset$ otherwise. We use $[n] \triangleq \fromto{1}{n}$ as an abbreviation. For clarity, we use the superscript $b$ to distinguish the variables with different choices of the mini-batch size $b$. In each iteration $t$, we use $\calB_t^b$ to denote the batch of samples (or sample indices) to calculate the stochastic gradient. We denote by $\calF_t^b$ the filtration of information before calculating the stochastic gradient in the $t$-th iteration, i.e. $\calF_t^b \triangleq \sth{w_0,\calB_0^b, \ldots, \calB_{t-1}^b}$.

Linear Regression {#subsec:linearRegression}
-----------------

In this subsection, we discuss the dynamics of SGD applied in linear regression. Given data points $(x_1, y_1), \cdots, (x_n, y_n)$, where $x_i \in \mathbb{R}^{p}$ and $y_i \in \mathbb{R}$, we define the loss function to be $$\label{eq:empiricalLoss}
        L(w) = \frac{1}{n} \sum_{i=1}^n L_i(w) = \frac{1}{n} \sum_{i=1}^n \frac{1}{2} \pth{w^\transp x_i - y_i}^2,$$ where $w \in \mathbb{R}^{p}$ are the model parameters. We consider minimizing by mini-batch SGD. Note that the bias term in the general linear regression models is omitted, however, adding the bias term does not change the result of this section. Formally, we first choose a mini-batch size $b$ and initial weights $w_0$. In each iteration $t$, we sample $\calB_t^b$, a subset of $[n]$ with cardinality $b$, and update the parameters by $$w_{t+1}^b = w_t^b - \alpha_t g_t^b,$$ where $g_t^b = \frac{1}{b} \sum_{i\in \calB_t^b} \grad L_i\pth{w_t^b}.$ We first show the relationship between the variance of stochastic gradient $g_t^b$ and the full gradient $\grad L\pth{w_t^b}$ and sample-wise gradient $\grad L_i\pth{w_t^b}, i\in[n]$, derived by considering all possible choices of the mini-batch $\calB_t^b$. Readers should note that Lemma \[lemma:1\] actually holds for all models with $L_2$-loss, not merely linear regression (since in the proof we do not need to know the explicit form of $L_i(w)$).

\[lemma:1\] Let $c_b \triangleq \frac{n-b}{b(n-1)} \ge 0$. For any matrix $A \in \mathbb{R}^{p\times p}$ we have [$$\begin{aligned}
            \condvar{}{Ag_t^b}{\calF_t^b} &= \condexp{}{\norm{Ag_t^b}^2}{\calF_t^b} - \norm{A\grad L\pth{w_t^b}}^2 \\
            &= c_b \left(\frac{1}{n}\sum_{i=1}^n\norm{A\grad L_i\pth{w_t^b}}^2 - \norm{A\grad L\pth{w_t^b}}^2\right).
        \end{aligned}$$]{}

Lemma \[lemma:1\] provides a bridge to connect the norm and variance of $g_t^b$ with sample-wise gradients $\grad L_i\pth{w_t^b}, i\in [n]$. Therefore, if we can further discover the properties of $\grad L_i\pth{w_t^b}, i\in [n]$, we are able to calculate the variance of $g_t^b$. Lemma \[lemma:3\] addresses this problem by showing the relationship between any linear combination of $\grad L_i\pth{w_t^b}$ and $\grad L_i\pth{w_{t-1}^b}$.

\[lemma:3\] For any set of square matrices $\sth{A_1, \cdots, A_n} \in \mathbb{R}^{p \times p}$, if we denote $A = \sum_{i=1}^n A_i x_i x_i^\transp$, then we have [$$\begin{aligned}
                & \condexp{}{\norm{\sum_{i=1}^{n}  A_i \grad L_i  \pth{w_{t+1}^b}}^2}{\calF_0}  =
                \condexp{}{\norm{\sum_{i=1}^{n} B_i \grad L_i\pth{w_{t}^b}}^2}{\calF_0} \\
                &+ \frac{\alpha_t^2 c_b}{n^2} \sum_{k=1}^n \sum_{l=1}^n\condexp{}{\norm{\sum_{i=1}^{n} B_i^{kl} \grad L_i\pth{w_{t}^b}}^2}{\calF_0}.
            \end{aligned}$$]{} Here $B_i = A_i - \frac{\alpha_t}{n} A $; $B_i^{kl} = A$ if $i=k, i\neq l$, $B_i^{kl} = A$ if $i=l, i\neq k$, and $B_i^{kl}$ equals the zero matrix, otherwise.

Lemma \[lemma:3\] provides the tool to reduce the iteration $t$ by one. Therefore, we can easily use it to recursively calculate the norm of any linear combinations of the sample-wise gradients, for all iterations $t$. Combining the fact that $c_b$ is a decreasing function of $b$, we are able to show Theorem \[thm:1\].

\[thm:1\] For any $t \in \naturals$ and any matrices $A_i \in \mathbb{R}^{p\times p}, i\in[n]$, $\condexp{}{\norm{\sum_{i=1}^{n} A_i \grad L_i\pth{w_{t}^b}}^2}{\calF_0}$ is a decreasing function of $b$ for $b \in [n]$.

Theorem \[thm:1\] states that the norm of any linear combinations of the sample-wise gradients is a decreasing function of $b$. Combining Lemma \[lemma:1\] which connects the variance of $g_t^b$ with the linear combination of $\grad L_i\pth{w_{t}^b}$’s, and the fact that $\grad L\pth{w_{t}^b} = \frac{1}{n}\sum_{i=1}^n\grad L_i\pth{w_{t}^b}$, we have Theorem \[thm:2\].

\[thm:2\] Fixing initial weights $w_0$, both $\var\left(Bg_t^b \, \middle| \calF_0 \right)$ and $\var\left(B\grad L\pth{w_t^b} \, \middle| \calF_0 \right)$ are decreasing functions of mini-batch size $b$ for all $b \in [n]$, $t \in \naturals$, and all square matrices $B \in \mathbb{R}^{p \times p}$.

As a special case, Corollary \[col:1\] guarantees that the variance of the stochastic gradient estimator is a decreasing function of $b$.

\[col:1\] Fixing initial weights $w_0$, both $\var\left(g_t^b \, \middle| \calF_0 \right)$ and $\var\left(\grad L\pth{w_t^b} \, \middle| \calF_0 \right)$ are decreasing functions of mini-batch size $b$ for all $b \in [n]$ and $t \in \naturals$.

In conclusion, we provide a framework for calculating the explicit value of variance of the stochastic gradient estimators and the norm of any linear combination of sample-wise gradients. We further show that the variance of both the full gradient and the stochastic gradient estimator are a decreasing function of the mini-batch size $b$.

Two-layer Linear Network with Online Setting {#subsec:linearNetwork}
--------------------------------------------

In this section, we study the dynamics of SGD on deep linear networks. We consider the two-layer linear network while the results and proofs can be easily extended to deep linear network with any depth. We consider the population loss $$\label{eq:populationLoss} 
        \calL (w) = \Expect_{x\sim \calN(0, I_p)} \left[ \frac{1}{2} \norm{W_2 W_1 x - W_2^* W_1^* x}^2 \right]$$ under the teacher-student learning framework [@hinton2015distilling] with $w=(W_1,W_2)$ a tuple of two matrices. Here $W_1 \in \mathbb{R}^{p_1 \times p}$ and $W_2 \in \mathbb{R}^{p_2 \times p_1}$ are parameter matrices of the student network and $W_1^*$ and $W_2^*$ are the fixed ground-truth parameters of the teacher network. We use online SGD to minimize the population loss $\calL(w)$. Formally, we first choose a mini-batch size $b$ and initial weight matrices $\sth{W_{0,1}, W_{0,2}}$. In each iteration $t$, we draw $b$ independent and identically distributed samples $x_{t,i}, i\in[b]$ from $\calN(0, I_p)$ to form the mini-batch $\calB_t^b$ and update the weight matrices by $W_{t+1,1}^b = W_{t,1}^b - \alpha_t g_{t,1}^b$ and $W_{t+1, 2}^b = W_{t,2}^b - \alpha_t g_{t,2}^b$, where $$\begin{aligned}
        g_{t, 1}^b &= \frac{1}{b} \sum_{i=1}^b \nabla_{W_{t,1}^b} \pth{\frac{1}{2} \norm{W_{t,2}^b W_{t,1}^b x_{t,i} - W_2^* W_1^* x_{t,i}}^2} \notag \\
        &= \frac{1}{b}\sum_{i=1}^b {W_{t,2}^b}^\transp \pth{W_{t,2}^b W_{t,1}^b - W_{2}^* W_{1}^*} x_{t,i} x_{t,i}^\transp, \label{eq:gt1}\\
        g_{t, 2}^b &= \frac{1}{b} \sum_{i=1}^b \nabla_{W_{t,2}^b} \pth{\frac{1}{2} \norm{W_{t,2}^b W_{t,1}^b x_{t,i} - W_2^* W_1^* x_{t,i}}^2} \notag \\
        &= \frac{1}{b}\sum_{i=1}^b \pth{W_{t,2}^b W_{t,1}^b - W_{2}^* W_{1}^*} x_{t,i} x_{t,i}^\transp{W_{t,1}^b}^\transp. \label{eq:gt2}
    \end{aligned}$$

The derivation follows from the formulas in [@IMM2012-03274]. In the following, we use $\calW_t^b = W_{t,2}^b W_{t,1}^b - W_{2}^* W_{1}^* $ to denote the gap between the product of model weights and ground-truth weights.

For ease of developing our proofs, we first introduce the definition of a *multiplicative term* in Definition \[def:mul\]. Intuitively, a multiplicative term is a matrix which equals to the product of its parameter matrices and constant matrices (and their transpose). The degree of a matrix $A$ in a multiplicative term $M$ is the number of appearance of $A$ and $A^\transp$ in $M$. The degree of $M$ is exactly the number of appearances of all weight matrices in $M$.

\[def:mul\] For any set of matrices $\calS$, we denote $\widebar{\calS} = \calS \cup \{M^\transp: M \in \calS\}$. Given a set of parameter matrices $\calX = \{X_1, X_2, \cdots, X_{n_v}\}$ and constant matrices $\calC = \{C_1, C_2, \cdots, C_{n_c}\}$, we say that a matrix $M$ is *a multiplicative term of parameter matrices $\calX$ and constant matrices $\calC$* if it can be written in the form of $$M = M(\calX, \calC) = \prod_{i=1}^k A_i,$$ where $A_i \in \widebar{\calX} \cup \widebar{\calC}$. We write $\deg(X_j; M) = \sum_{i=1}^k \pth{\mathbbm{1}\sth{X_j = A_i} + \mathbbm{1}\sth{X_j = A_i^\transp}}, j \in [n_v]$ as the degree of parameter matrix $X_j$ in $M$, $\deg(C_j; M) = \sum_{i=1}^k \pth{\mathbbm{1}\sth{C_j = A_i} + \mathbbm{1}\sth{C_j = A_i^\transp}}, j \in [n_c]$ as the degree of constant matrix $C_j$ in $M$, and $\deg(M) = \sum_{i=1}^k \mathbbm{1}\sth{A_i \in \widebar{\calX}} = \sum_{j=1}^{n_v} \deg(X_j; M)$ as the total degree of the parameter matrices of $M$.

As pointed out in the Section \[sec:1\], the difficulty of studying the dynamics of SGD is how to connect the quantities in iteration $t$ with fixed variables, like initial weights $W_{0,1}, W_{0,2}$ and mini-batch size $b$. We overcome this challenge by the following two lemmas. Lemma \[lemma:GtoW\] provides the relationship between $g_{t,i}^b, i=1,2$ and $W_{t,i}^b, i=1,2$ by taking expectation over the distribution of random samples in $\calB_t^b$. Lemma \[lemma:WtoG\] shows the relationship between $W_{t,i}^b, i=1,2$ and $g_{t-1,i}^b, i=1,2$ using and .

\[lemma:GtoW\] For multiplicative terms $M_i, i \in \fromto{0}{m} $ of parameter matrices $\left\{ g_{t,1}^b, g_{t,2}^b \right\}$ and constant matrices $\sth{W_{t,1}^b, W_{t,2}^b, W_1^*, W_2^*}$ with degree $d_i$, respectively, we denote $M = \prod_{i=1}^m \trace{M_i} M_0$ and $d = \sum_{i=0}^m d_i$. There exists a set of multiplicative terms $\sth{M_{ij}^k, i\in [m_k], j \in \fromto{0}{m_{ki}}, k \in \fromto{0}{q}}$ of parameter matrices $\sth{W_{t,1}^b, W_{t,2}^b}$ and constant matrices $\sth{W_1^*, W_2^*}$ such that $$\begin{aligned}
                \condexp{}{M}{\calF_t^b} = N_0 + N_1 \frac{1}{b} + \cdots + N_d \frac{1}{b^d},
                            \end{aligned}$$ where $N_k = \sum_{i=1}^{m_k} \prod_{j=1}^{m_{ki}} \trace{M_{ij}^k} M_{i0}^k, k\in \fromto{0}{d}  $. Here $m_k, m_{ki}$ are constants independent of $b$, and $\sum_{j=0}^{m_{ki}} \deg\pth{M_{ij}^k} \le 3d + \sum_{i=0}^m \pth{\deg\pth{W_{t,1}^b; M_i} + \deg(W_{t,2}^b; M_i)}$.

\[lemma:WtoG\] For multiplicative term $M_i, i \in \fromto{0}{m}$ of parameter matrices $\sth{W_{t,1}^b, W_{t,2}^b}$ and constant matrices $\sth{W_1^*, W_2^*}$ of degree $d_i$, let $d = 2^{d_0+\cdots +d_m}$. There exists a set of multiplicative terms $\{M_{ik}, i \in \fromto{0}{m}, k \in [d]\}$ of parameter matrices $\sth{g_{t,1}^b, g_{t,2}^b}$ and constant matrices $\sth{W_{t,1}^b, W_{t,2}^b, W_1^*, W_2^*}$ such that $$\begin{aligned}
                \prod_{i=1}^m \trace{M_i} M_0 = \sum_{k=1}^{d} \prod_{i=1}^{m} \trace{M_{ik}} M_{0k},
            \end{aligned}$$ where $\sum_{i=0}^m \deg\pth{M_{ik}} \le d$.

With the help of Lemmas \[lemma:GtoW\] and \[lemma:WtoG\], we can represent $g_{t,i}^b, i=1,2$ using multiplicative terms of $g_{t-1,i}^b, i=1,2$ and some other constant matrices. Furthermore, by iteratively reducing the value of $t$, we are able to represent $g_{t,i}^b, i=1,2$ by the variables in $t = 0$. Theorem \[thm:representation\] precisely gives the representation in the form of a polynomial of $\frac{1}{b}$ and the coefficients as the sum of multiplicative terms of parameter matrices $\sth{W_{0,1}^b, W_{0,2}^b}$ and constant matrices $\sth{W_1^*, W_2^*}$.

\[thm:representation\] Given $t \ge 0$, for any multiplicative terms $M_i, i \in \fromto{0}{m} $ of parameter matrices $\left\{ g_{t,1}^b, g_{t,2}^b \right\}$ and constant matrices $\sth{W_{t,1}^b, W_{t,2}^b, W_1^*, W_2^*}$ with degree $d_i$, respectively, we denote $M = \prod_{i=1}^m \trace{M_i} M_0$, $d = \sum_{i=0}^m d_i$ and $d' = \sum_{i=0}^m \pth{\deg\pth{W_{t,1}^b; M_i} + \deg(W_{t,2}^b; M_i)}$. There exists a set of multiplicative terms $\sth{M_{ij}^k, i\in [m_k], j \in \fromto{0}{m_{ki}}, k \in \fromto{0}{q}}$ of parameter matrices $\sth{W_{0,1}^b, W_{0,2}^b}$ and constant matrices $\sth{W_1^*, W_2^*}$ such that $$\begin{aligned}
                \condexp{}{M}{\calF_0} = N_0 + N_1 \frac{1}{b} + \cdots + N_q \frac{1}{b^q},
            \end{aligned}$$ where $N_k = \sum_{i=1}^{m_k} \prod_{j=1}^{m_{ki}} \trace{M_{ij}^k} M_{i0}^k, k\in \fromto{0}{q}  $. Here $m_k, m_{ki}$ and $q \le \frac{1}{2}(3^{t+1}-1)d + \frac{1}{2}(3^{t}-1)d'$ are constants independent of $b$, and $\sum_{j=0}^{m_{ki}} \deg\pth{M_{ij}^k} \le 3^{t} (3d + d')$.

By changing the role of parameter and constant matrices we obtain the following corollary.

\[thm:representationoCorollary\] Given $t \ge 0$, for any multiplicative terms $M_i, i \in \fromto{0}{m} $ of parameter matrices $\sth{W_{t,1}^b, W_{t,2}^b, \calW_t^b}$ and constant matrices $\sth{W_1^*, W_2^*}$ such that $\sum_{i=1}^2 \deg \pth{W_{t,i}^b; M} = d$ and $\deg \pth{\calW_t^b; M} = d'$, we denote $M = \prod_{i=1}^m \trace{M_i} M_0$. There exists a set of multiplicative terms $\sth{M_{ij}^k, i\in [m_k], j \in \fromto{0}{m_{ki}}, k \in \fromto{0}{q}}$ of parameter matrices $\sth{W_{0,1}^b, W_{0,2}^b}$ and constant matrices $\sth{W_1^*, W_2^*}$ such that $$\begin{aligned}
                \condexp{}{M}{\calF_0} = N_0 + N_1 \frac{1}{b} + \cdots + N_q \frac{1}{b^q},
            \end{aligned}$$ where $N_k = \sum_{i=1}^{m_k} \prod_{j=1}^{m_{ki}} \trace{M_{ij}^k} M_{i0}^k, k\in \fromto{0}{q}  $. Here $m_k, m_{ki}$ and $q \le 3^{t}\pth{d + 2d'}$ are constants independent of $b$, and $\sum_{j=0}^{m_{ki}} \deg\pth{M_{ij}^k} \le 3^{t} \pth{d + 2d'}$.

As a special case of Theorem \[thm:representation\], Theorem \[thm:varianceRepresentation\] shows that the variance of the stochastic gradient estimators is also a polynomial of $\frac{1}{b}$ but with no constant term. This backs the important intuition that the variance is approximately inversely proportional to the mini-batch size $b$. Besides, note that if we consider $b \goto \infty$, intuitively we should have $\condvar{}{g_{t,i}^b}{\calF_0} \goto 0, i=1,2$. This observation aligns with the statement of Theorem \[thm:varianceRepresentation\].

\[thm:varianceRepresentation\] Given $t \ge 0$, value $\condvar{}{g_{t,i}^b}{\calF_0}, i=1,2$ can be written as a polynomial of $\frac{1}{b}$ with degree at most $2\cdot 3^t$ with no constant term. Formally, we have $$\condvar{}{g_{t,i}^b}{\calF_0} = \beta_1 \frac{1}{b} + \cdots + \beta_r \frac{1}{b^r}, \label{eq:2.7}$$ where $r \le 2\cdot 3^{t+1}$ and each $\beta_i$ is a constant independent of $b$.

Finally, to show the that the variance is a decreasing function of $b$ for large enough $b$, we only need to show that the leading coefficient $\beta_1$ is non-negative. This is guaranteed by the fact that variance is always non-negative. We therefore have Theorem \[thm:varianceDecrease\].

\[thm:varianceDecrease\] Given $t \in \mathbb{N}$, there exists a constant $b_0$ such that for all $b \ge b_0 $ function $\condvar{}{g_{t,i}^b}{\calF_0}, i=1,2$ is a decreasing function of $b$.

In conclusion, we present the relationship between any multiplicative terms of parameter matrices $\sth{g_{t,i}^b, W_{t,i}^b, i=1,2}$ and constant matrices $\sth{W_1^*, W_2^*}$ and the initial weights $W_{0,1}, W_{0,2}$ and the mini-batch size $b$. Unlike the linear regression setting, the closed form expressions for the variance are unknown. However, Theorem \[thm:varianceRepresentation\] conquers this issue by iteratively deducing $t$ one by one and it provides a polynomial representation. We are also able to show the decreasing property of the variance of stochastic gradient estimators with respect to $b$, based on this polynomial representation.

  ---------------------------------------------------------- ------------------------------------------------------------ -- --
      ![image](figure/regression.png){width="50.00000%"}      ![image](figure/regression_fitpoly.png){width="50.00000%"}     
   \(a) Variance of stochastic gradients and full gradients        \(b) Fitting polynomials of mini-batch size $b$           
  ---------------------------------------------------------- ------------------------------------------------------------ -- --

  ----------------------------------------------------- ----------------------------------------------------- -- --
   ![image](figure/twolayer_W1.png){width="50.00000%"}   ![image](figure/twolayer_W2.png){width="50.00000%"}     
    \(a) Variance of gradients with respect to $W_1$      \(b) Variance of gradients with respect to $W_2$       
  ----------------------------------------------------- ----------------------------------------------------- -- --

  ------------------------------------------------------- ---------------------------------------------------------- -- -- -- --
   ![image](figure/MNIST_initial.png){width="50.00000%"}   ![image](figure/MNIST_gen_zoomed.png){width="50.00000%"}           
              \(a) Different initial weights                           \(b) Gap of accuracy (zoomed-in)                       
  ------------------------------------------------------- ---------------------------------------------------------- -- -- -- --

Experiments {#sec:4}
===========

In this section, we present numerical results to support the theorems in Section \[sec:3\] and provide further insights into the impact of the mini-batch size on the dynamics of SGD. The experiments are conducted on four datasets and models that are relatively small due to the computational cost of using large models and datasets. The goal of these experiments is to support the theorems in Section \[sec:3\], to backup the hypotheses discussed in the introduction, and to provide further insights.

For all experiments, we perform mini-batch SGD multiple times starting from the same initial weights and following the same choice of the learning rates and other hyper-parameters, if applicable. This enables us to calculate the variance of the gradient estimators and other statistics in each iteration, where the randomness comes only from different samples of SGD. The learning rate $\alpha_t$ is selected to be inversely proportional to iteration $t$, or fixed, depending on the task at hand.

All models are implemented using PyTorch version 1.4 [@paszke2019pytorch] and trained on NVIDIA 2080Ti/1080 GPUs. We report the details about the hyperparameters and training procedures in Appendix \[appendix:exp\].

Datasets and Settings
---------------------

The Graduate Admission dataset[^1] [@acharya2019comparison] is to predict the chance of a graduate admission using linear regression. The dataset contains 500 samples with 6 features. This is a popular regression dataset with clean data. We build a linear regression model to predict the chance of acceptance (we include the intercept term in the model) and minimize the empirical $L_2$ loss using mini-batch SGD, as stated in Section \[subsec:linearRegression\]. The purpose of this experiment is to empirically study the rate of decrease of the variance. The theoretical study exhibited in Section \[subsec:linearRegression\] establishes the non-increasing property but it does not state anything about the rate of decrease.

We build a synthetic dataset of standard normal samples to study the setting in Section \[subsec:linearNetwork\]. We fix the teacher network with 64 input neurons, 256 hidden neurons and 128 output neurons. We optimize the population $L_2$ loss by updating the two parameter matrices of the student network using online SGD, as stated in Section \[subsec:linearNetwork\]. In this case we have proved the functional form of the variance as a function of $b$ and show the decreasing property of the variance of the stochastic gradient estimators for large mini-batch sizes. However, we do not show the decreasing property for every $b$. With this experiment we confirm that the conjecture likely holds.

The MNIST dataset is to recognize digits in handwritten images of digits. We use all 60,000 training samples and 10,000 validation samples of MNIST. We build a three-layer fully connected neural network with 1024, 512 and 10 neurons in each layer. For the two hidden layers, we use the ReLU activation function. The last layer is the softmax layer which gives the prediction probabilities for the 10 digits. We use mini-batch SGD to optimize the cross-entropy loss of the model. The model deviates from our analytical setting since it has non-linear activations, it has the cross-entropy loss function (instead of $L_2$), and empirical loss (as opposed to population). MNIST is selected due to its fast training and popularity in deep learning experiments. The goal is to verify the results in this different setting and to back up our hypotheses.

The Yelp Review dataset from the Yelp Dataset Challenge 2015 [@zhang2015character] contains 1,569,264 samples of customer reviews with positive/negative sentiment labels. We use 10,000 samples as our training set and 1,000 samples as the validation set. We use XLNet [@yang2019xlnet] to perform sentiment classification on this dataset. Our XLNet has 6 layers, the hidden size of 384, and 12 attention heads. There are in total 35,493,122 parameters. We intentionally reduce the number of layers and hidden size of XLNet and select a relatively small size of the training and validation sets since training of XLNet is very time-consuming ([@yang2019xlnet] train on 512 TPU v3 chips for 5.5 days) and we need to train the model for multiple runs. This setting allows us to train our model in several hours on a single GPU card. We train the model using the Adam weight decay optimizer, and some other techniques, as suggested in Table 8 of [@yang2019xlnet]. This dataset represents sequential data where we further consider the hypotheses.

Discussion {#sec:4.2}
----------

As observed in Figure \[fig:linear\](a), under the linear regression setting with the Graduate Admission dataset, the variance of the stochastic gradient estimators and full gradients are all strictly decreasing functions of $b$ for all iterations. This result verifies the theorems in Section \[subsec:linearRegression\]. Figure \[fig:linear\](b) further studies the rate of decrease of the variance. From the proofs in Section \[subsec:linearRegression\] we see that $\condvar{}{g_t^b}{\calF_0}$ is a polynomial of $\frac{1}{b}$ with degree $t+1$. Therefore, for every $t$, we can approximate this polynomial by sampling many different $b$’s and calculate the corresponding variances. We pick $b$ to cover all numbers that are either a power of 2 or multiple of 40 in $[2, 500]$ (there are a total of 21 such values) and fit a polynomial with degree 6 (an estimate from the analyses) at $t=10, 20, 30, 40$. Figure \[fig:linear\](b) shows the fitted polynomials. As we observe, the value $\condvar{}{g_t^b}{\calF_0}$ (approximated by the value of the polynomial) is both decreasing with respect to the mini-batch size $b$ and iteration $t$. Further, the rate of decrease in $b$ is slower as the $b$ increasing. This provides a further insight into the dynamics of training a linear regression problem with SGD.

Under the two-layer linear network setting with the synthetic dataset, Figure \[fig:twolayer\] verifies that the variance of the stochastic gradient estimators and full gradients are all strictly decreasing functions of $b$ for all iterations. This figure also empirically shows that the constant $b_0$ in Theorem \[thm:varianceDecrease\] could be as small as $b_0 = 4$. In fact, we also experiment with the mini-batch size of 1 and 2, and the decreasing property remains to hold. We also test this on multiple choices of initial weights and learning rates and this pattern remains clear.

  ---------------------------------------------------- ---------------------------------------------------- -- --
   ![image](figure/MNIST_loss.png){width="50.00000%"}   ![image](figure/MNIST_acc.png){width="50.00000%"}      
   \(a) Log of loss for training and validation sets    \(b) Log of error for training and validation sets     
  ---------------------------------------------------- ---------------------------------------------------- -- --

  -------------------------------------------------------- ---------------------------------------------------- ----------------------------------------------------- -- -- --
   ![image](figure/xlnet_variance.png){width="35.00000%"}   ![image](figure/xlnet_loss.png){width="35.00000%"}   ![image](figure/xlnet_error.png){width="35.50000%"}        
           \(a) Variance of stochastic gradients                    \(b) Training and validation loss                  \(c) Training minus validation accuracy              
  -------------------------------------------------------- ---------------------------------------------------- ----------------------------------------------------- -- -- --

In aforementioned two experiments we use SGD in its original form by randomly sampling mini-batches. In deep learning with large-scale training data such a strategy is computationally prohibitive and thus samples are scanned in a cyclic order which implies fixed mini-batches are processed many times. Therefore, in the next two datasets we perform standard “epoch” based training to empirically study the remaining two hypotheses discussed in the introduction (decreasing loss and error as a function of $b$) and sensitivity with respect to the initial weights. Note that we are using cross-entropy loss in the MNIST dataset and the Adam optimizer in the Yelp dataset and thus these experiments do not meet all of the assumptions of the analysis in Section \[sec:3\].

As shown in Figure \[fig:MNIST2\](a), we run SGD with two batch sizes 64 and 128 on five different initial weights. This plot shows that, even the smallest value of the variance among the five different initial weights with a mini-batch size of 64, is still larger than the largest variance of mini-batch size 128. We observe that the sensitivity to the initial weights is not large. This plot also empirically verifies our conjecture in the introduction that the variance of the stochastic gradient estimators is a decreasing function of the mini-batch size, for all iterations of SGD in a general deep learning model.

In addition, we also conjecture that there exists the decreasing property for the expected loss, error and the generalization ability with respect to the mini-batch size. Figure \[fig:MNIST1\](a) shows that the expected loss (again, randomness comes from different runs of SGD through the different mini-batches with the same initial weights and learning rates) on the training set is a decreasing function of $b$. However, this decreasing property does not hold on the validation set when the loss tends to be stable or increasing, in other words, the model starts to be over-fitting. We hypothesize that this is because the learned weights start to bounce around a local minimum when the model is over-fitting. As the larger mini-batch size brings smaller variance, the weights are closer to the local minimum found by SGD, and therefore yield a smaller loss function value. Figure \[fig:MNIST1\](b) shows that both the expected error on training and validation sets are decreasing functions of $b$.

Figure \[fig:MNIST2\](b) exhibits a relationship between the model’s generalization ability and the mini-batch size. As suggested by [@1227801], we build a test set by distorting the 10,000 images of the validation set. The prediction accuracy is obtained on both training and test sets and we calculate the gap between these two accuracies every 100 epochs. We use this gap to measure the model generalization ability (the smaller the better). Figure \[fig:MNIST2\](b) shows that the gap is an increasing function of $b$ starting at epoch 500, which partially aligns with our conjecture regarding the relationship between the generalization ability and the mini-batch size. We also test this on multiple choices of the hyper-parameters which control the degree of distortion in the test set and this pattern remains clear.

Figure \[fig:YELP\] shows the similar phenomenon that the variance of stochastic estimators and the expected loss and error on both training and validation sets are decreasing functions of $b$ even if we train XLNet using Adam. This example gives us confidence that the decreasing properties are not merely restricted on shallow neural networks or vanilla SGD algorithms. They actually appear in many advanced models and optimization methods.

Summary and Future Work
=======================

We examine the impact of the mini-batch size on the dynamics of SGD. Our focus is on the variance of stochastic gradient estimators. For linear regression and a two-layer linear network, we are able to theoretically prove that the variance conjecture holds. We further experiment on multiple models and datasets to verify our claims and their applicability to practical settings. Besides, we also empirically address the conjectures about the expected loss and the generalization ability.

There are several possible directions for future work. One obvious extension of this work is to show the decreasing property of variance to more general machine learning models, like fully connected networks with activation functions and residual connections. Another challenging research direction is to theoretically investigate the impact of the mini-batch size on the expected loss and the generalization ability of machine learning models (the conjectures we mentioned in Section \[sec:1\]). The extensions of this work to other optimization algorithms, like Adam and Gradient Boosting Machines, are also very attractive. We hope our proof techniques can serve as a tool for future research.

Proofs
======

Proofs of Results in Section \[subsec:linearRegression\]
--------------------------------------------------------

For two matrices $A, B$ with the same dimension, we define the inner product $\iprod{A}{B} \triangleq \trace{A^\transp B} $.

\[lemma:product\] Suppose that $f(x)$ and $g(x)$ are both smooth, non-negative and decreasing functions of $x \in \reals$. Then $h(x) = f(x) g(x)$ is also a non-negative and decreasing function of $x$.

It is obvious that $h(x)$ is non-negative for all $x$. The first-order derivative of $h$ is $$h'(x) = f'(x) g(x) + f(x) g'(x) \le 0,$$ and thus $h(x)$ is also a decreasing function of $x$.

\[Proof of Lemma \[lemma:1\]\] Note that $$\begin{aligned}
            \condexp{}{g_t^b \pth{g_t^b}^\transp}{\calF_t^b} &= \frac{1}{b^2}\condexp{}{\sum_{i\in \calB_t^b} \grad L_i\pth{w_t^b} \sum_{i\in \calB_t^b} \grad L_i\pth{w_t^b}^\transp}{\calF_t^b} \\
            &=\frac{1}{b^2}\left( \frac{C_{n-1}^{b-1}}{C_n^b} \sum_{i=1}^n \grad L_i\pth{w_t^b}\grad L_i\pth{w_t^b}^\transp + \frac{C_{n-2}^{b-2}}{C_n^b} \sum_{i\neq j} \grad L_i\pth{w_t^b}\grad L_j\pth{w_t^b}^\transp \right) \\
            &=\frac{1}{b^2}\left( \frac{b}{n} \sum_{i=1}^n \grad L_i\pth{w_t^b}\grad L_i\pth{w_t^b}^\transp + \frac{b(b-1)}{n(n-1)} \sum_{i\neq j} \grad L_i\pth{w_t^b}\grad L_j\pth{w_t^b}^\transp \right) \\
            &= \frac{1}{b^2}\left( \frac{b(n-b)}{n(n-1)}\sum_{i=1}^n \grad L_i\pth{w_t^b}\grad L_i\pth{w_t^b}^\transp + \frac{b(b-1)}{n(n-1)} \sum_{i=1}^n \grad L_i\pth{w_t^b}\sum_{i=1}^n \grad L_i\pth{w_t^b}^\transp \right)\\
            &= \frac{n-b}{bn(n-1)} \sum_{i=1}^n \grad L_i\pth{w_t^b}\grad L_i\pth{w_t^b}^\transp + \frac{(b-1)n}{b(n-1)}\grad L\pth{w_t^b}\grad L\pth{w_t^b}^\transp.
        \end{aligned}$$ For any $A  \in \mathbb{R}^{p\times p}$, we have $$\begin{aligned}
            \condexp{}{\norm{A g_t^b}^2}{\calF_t^b} &= \condexp{}{\pth{g_t^b}^\transp A^\transp A g_t^b}{\calF_t^b}  =  \condexp{}{\trace{\pth{g_t^b}^\transp A^\transp A g_t^b}}{\calF_t^b} \\
            &= \condexp{}{\trace{ A^\transp A g_t^b \pth{g_t^b}^\transp}}{\calF_t^b} \\
            &= \trace{ A^\transp A \condexp{}{g_t^b \pth{g_t^b}^\transp}{\calF_t^b}} \\
            &= \trace{\frac{n-b}{bn(n-1)} \sum_{i=1}^n A^\transp A\grad L_i\pth{w_t^b}\grad L_i\pth{w_t^b}^\transp + \frac{(b-1)n}{b(n-1)}A^\transp A\grad L\pth{w_t^b}\grad L\pth{w_t^b}^\transp} \\
            &= \frac{n-b}{bn(n-1)} \sum_{i=1}^n \norm{A\grad L_i\pth{w_t^b}}^2 + \frac{(b-1)n}{b(n-1)} \norm{A\grad L\pth{w_t^b}}^2\\
            &= c_b \left(\frac{1}{n}\sum_{i=1}^n\norm{A\grad L_i\pth{w_t^b}}^2 - \norm{A\grad L\pth{w_t^b}}^2\right) + \norm{A\grad L\pth{w_t^b}}^2.
        \end{aligned}$$

Therefore, we have $$\begin{aligned}
            \condvar{}{A g_t^b}{\calF_t^b} &= \condexp{}{\norm{A g_t^b}^2}{\calF_t^b} - \norm{\condexp{}{A g_t^b}{\calF_t^b}}^2 \\
            &= \condexp{}{\norm{A g_t^b}^2}{\calF_t^b} - \norm{A\grad L\pth{w_t^b}}^2 \\ 
            &= c_b \left(\frac{1}{n}\sum_{i=1}^n\norm{A\grad L_i\pth{w_t^b}}^2 - \norm{A\grad L\pth{w_t^b}}^2\right).
        \end{aligned}$$

Let $C_i = x_i x_i^\transp$ and $C = \frac{1}{n} \sum_{i=1}^n C_i $. For the given $A_1, \ldots, A_n$, we denote $A = \sum_{i=1}^n A_i C_i $. Then we have $$\begin{aligned}
            &\condexp{}{\norm{\sum_{i=1}^{n} A_i \grad L_i\pth{w_{t+1}^b}}^2}{\calF_0} = \condexp{}{\condexp{}{\norm{\sum_{i=1}^{n} A_i \grad L_i\pth{w_{t+1}^b}}^2}{\calF_t^b}}{\calF_0}  \\
            =\ & \condexp{}{\condexp{}{\norm{\sum_{i=1}^{n} A_i \pth{x_i^\transp w_{t+1}^b - y_i}x_i}^2}{\calF_t^b}}{\calF_0} \\
            =\ & \condexp{}{\condexp{}{\norm{\sum_{i=1}^{n} A_i \pth{x_i^\transp \pth{w_{t}^b - \alpha_t g_t^b} - y_i}x_i}^2}{\calF_t^b}}{\calF_0} \\
            =\ & \condexp{}{\condexp{}{\norm{\sum_{i=1}^{n} A_i \grad L_i\pth{w_{t}^b} - \alpha_t A g_t^b}^2}{\calF_t^b}}{\calF_0} \\
            =\ & \condexp{}{\norm{\sum_{i=1}^{n} A_i \grad L_i\pth{w_{t}^b}}^2}{\calF_0} - 2\alpha_t \condexp{}{\condexp{}{\iprod{\sum_{i=1}^{n} A_i \grad L_i\pth{w_{t}^b}}{A g_t^b}}{\calF_t^b}}{\calF_0}\\
            &+\alpha_t^2 \condexp{}{\condexp{}{\norm{A g_t^b}^2}{\calF_t^b}}{\calF_0} \\
            =\ & \condexp{}{\norm{\sum_{i=1}^{n} A_i \grad L_i\pth{w_{t}^b}}^2}{\calF_0} - 2\alpha_t \condexp{}{\iprod{\sum_{i=1}^{n} A_i \grad L_i\pth{w_{t}^b}}{A \grad L\pth{w_t^b}}}{\calF_0}\\
            &+\alpha_t^2 \condexp{}{c_b \pth{\frac{1}{n} \sum_{i=1}^n \norm{A\grad L_i(w_t^{b})}^2 - \norm{A\grad L(w_t^{b})}^2} + \norm{A \grad L(w_t^{b})}^2 }{\calF_0} \\
            =\ & \condexp{}{\norm{\sum_{i=1}^{n} A_i \grad L_i\pth{w_{t}^b} - \alpha_t A \grad L(w_t^{b})}^2}{\calF_0} +\alpha_t^2 c_b \condexp{}{\frac{1}{n} \sum_{i=1}^n \norm{A\grad L_i(w_t^{b})}^2 - \norm{A\grad L(w_t^{b})}^2}{\calF_0} \\
            =\ & \condexp{}{\norm{\sum_{i=1}^{n} A_i \grad L_i\pth{w_{t}^b} - \alpha_t A \grad L(w_t^{b})}^2}{\calF_0} + \frac{\alpha_t^2 c_b}{n^2} \sum_{i\neq j}\condexp{}{\norm{A \grad L_i\pth{w_{t}^b} - A\grad L_j \pth{w_{t}^b}}^2}{\calF_0}\\
            =\ & \condexp{}{\norm{\sum_{i=1}^{n} \pth{A_i - \frac{\alpha_t}{n} A} \grad L_i\pth{w_{t}^b}}^2}{\calF_0} + \frac{\alpha_t^2 c_b}{n^2} \sum_{i=1}^n \sum_{j=1}^n\condexp{}{\norm{A \grad L_i\pth{w_{t}^b} - A\grad L_j \pth{w_{t}^b}}^2}{\calF_0}.
        \end{aligned}$$ Therefore, if we set $B_i = A_i - \frac{\alpha_t}{n} A $ and $$B_i^{kl}=\left\{
            \begin{aligned}
                A &  & i=k, i\neq l, \\
                -A &  & i=l, i\neq k, \\
                0 &  & \textrm{otherwise},
            \end{aligned}
            \right.$$ we have $$\begin{aligned}
                \condexp{}{\norm{\sum_{i=1}^{n} A_i \grad L_i\pth{w_{t+1}^b}}^2}{\calF_0} = \condexp{}{\norm{\sum_{i=1}^{n} B_i \grad L_i\pth{w_{t}^b}}^2}{\calF_0} + \frac{\alpha_t^2 c_b}{n^2} \sum_{k=1}^n \sum_{l=1}^n\condexp{}{\norm{\sum_{i=1}^{n} B_i^{kl} \grad L_i\pth{w_{t}^b}}^2}{\calF_0}.
            \end{aligned}$$

We use induction to show this statement.

When $t = 0$, $\condexp{}{\norm{\sum_{i=1}^{n} A_i \grad L_i\pth{w_{t}^b}}^2}{\calF_0} = \norm{\sum_{i=1}^{n} A_i \grad L_i\pth{w_0}}^2 $ which is invariant of $b$. Therefore, it is a decreasing function of $b$.

Suppose the statement holds for $t$. For any set of matrices $\left\{A_1, \ldots, A_n \right\}$ in $\mathbb{R}^{p\times p}$, by Lemma \[lemma:3\] we know that there exist matrices $\left\{B_1, \cdots, B_n\right\}$ and $\left\{ B_i^{kl}: i,k,l \in [n] \right\}$ such that $$\begin{aligned}
                \condexp{}{\norm{\sum_{i=1}^{n} A_i \grad L_i\pth{w_{t+1}^b}}^2}{\calF_0}= \condexp{}{\norm{\sum_{i=1}^{n} B_i \grad L_i\pth{w_{t}^b}}^2}{\calF_0} + \frac{\alpha_t^2 c_b}{n^2} \sum_{k=1}^n \sum_{l=1}^n\condexp{}{\norm{\sum_{i=1}^{n} B_i^{kl} \grad L_i\pth{w_{t}^b}}^2}{\calF_0}.
            \end{aligned}$$

By induction, we know that $\condexp{}{\norm{\sum_{i=1}^{n} B_i \grad L_i\pth{w_{t}^b}}^2}{\calF_0}$ and all $\condexp{}{\norm{\sum_{i=1}^{n} B_i^{kl} \grad L_i\pth{w_{t}^b}}^2}{\calF_0}$ are non-negative and decreasing functions of $b$. Besides, clearly $\frac{\alpha_t^2 c_b}{n^2} = \frac{\alpha_t^2 (n-b)}{bn^3(n-1)}$ is a non-negative and decreasing function of $b$. By Lemma \[lemma:product\], we know that $ \frac{\alpha_t^2 c_b}{n^2} \condexp{}{\norm{\sum_{i=1}^{n} B_i^{kl} \grad L_i\pth{w_{t}^b}}^2}{\calF_0} $ is also a non-negative and decreasing function of $b$. Finally, $\condexp{}{\norm{\sum_{i=1}^{n} A_i \grad L_i\pth{w_{t+1}^b}}^2}{\calF_0}$, as the sum of non-negative and decreasing functions in $b$, is a non-negative and decreasing function of $b$.

In order to prove Theorem \[thm:2\], we split the task to two separate theorems about the full gradient and the stochastic gradient and prove them one by one.

\[theorem:varianceFullGradient\] Fixing initial weights $w_0$, $\var\left(B\grad L\pth{w_t^b} \, \middle| \calF_0 \right)$ is a decreasing function of mini-batch size $b$ for all $b \in [n]$, $t \in \naturals$, and all square matrices $B \in \mathbb{R}^{p \times p}$.

\[thm:decLinear\] Fixing initial weights $w_0$, $\var\left(Bg_t^b \, \middle| \calF_0 \right)$ is a decreasing function of mini-batch size $b$ for all $b \in [n]$, $t \in \naturals$, and all square matrices $B \in \mathbb{R}^{p \times p}$.

We induct on $t$ to show that the statement holds. For $t = 0$, we have $\var\left(B\grad L\pth{w_t^b} \, \middle| \calF_0 \right) = 0 $ for any matrix $B$. Suppose the statement holds for $t - 1 \ge 0$. Note that from $$\begin{aligned}
            \grad L\pth{w_t^b} &= \frac{1}{n}\sum_{i=1}^n x_i\pth{x_i^\transp w_t^b - y_i}\\
            &= \frac{1}{n}\sum_{i=1}^n x_i\pth{x_i^\transp \pth{w_{t-1}^b - \alpha_t g_{t-1}^b} - y_i}\\
            &= \frac{1}{n}\sum_{i=1}^n x_i\pth{x_i^\transp w_{t-1}^b - y_i} - \frac{\alpha_t}{n}\sum_{i=1}^n x_i x_i^\transp g_{t-1}^b \\
            &= \grad L\pth{w_{t-1}^b} - \alpha_t C g_{t-1}^b,
        \end{aligned}$$

we have $$\begin{aligned}
            &\var\left(B\grad L\pth{w_t^b} \, \middle| \calF_0 \right) \nonumber\\
            =\  & \var\left(B\grad L\pth{w_{t-1}^b} - \alpha_t BC g_{t-1}^b \, \middle| \calF_0 \right) \nonumber\\
            =\  & \condexp{}{\norm{B\grad L\pth{w_{t-1}^b} - \alpha_t BC g_{t-1}^b}^2}{\calF_0^b} - \norm{\condexp{}{B\grad L\pth{w_{t-1}^b} - \alpha_t BC g_{t-1}^b}{\calF_0^b}}^2 \nonumber\\
            =\  & \condexp{}{\norm{B\grad L\pth{w_{t-1}^b}}^2 - 2\alpha_t \iprod{B\grad L\pth{w_{t-1}^b}}{BC g_{t-1}^b} + \alpha_t^2 \norm{BC g_{t-1}^b}^2}{\calF_0^b} - \norm{\condexp{}{B\grad L\pth{w_{t-1}^b} - \alpha_t BC g_{t-1}^b}{\calF_0^b}}^2 \nonumber\\
            =\  & \condexp{}{\norm{B\grad L\pth{w_{t-1}^b}}^2}{\calF_0} + \alpha_t^2 \condexp{}{\condexp{}{\norm{BC g_{t-1}^b}^2}{\calF_{t-1}^b}}{\calF_0^b} - 2\alpha_t \condexp{}{ \condexp{}{\iprod{B\grad L\pth{w_{t-1}^b}}{BC g_{t-1}^b}}{\calF_{t-1}^b} }{\calF_0} \nonumber\\
            \  & - \norm{\condexp{}{ \condexp{}{B\grad L\pth{w_{t-1}^b} - \alpha_t BC g_{t-1}^b}{\calF_{t-1}^b} }{\calF_0^b}}^2 \nonumber\\
            =\  & \condexp{}{\norm{B\grad L\pth{w_{t-1}^b}}^2}{\calF_0} + \alpha_t^2 \condexp{}{ c_b \left(\frac{1}{n}\sum_{i=1}^n\norm{BC\grad L_i\pth{w_{t-1}^b}}^2 - \norm{BC\grad L\pth{w_{t-1}^b}}^2\right) + \norm{BC\grad L\pth{w_{t-1}^b}}^2 }{\calF_0} \nonumber\\
            \  &  - 2\alpha_t \condexp{}{ \iprod{B\grad L\pth{w_{t-1}^b}}{BC\grad L\pth{w_{t-1}^b}} }{\calF_0} - \norm{\condexp{}{ B\grad L\pth{w_{t-1}^b} - \alpha_t BC \grad L\pth{w_{t-1}^b} }{\calF_0^b}}^2 \label{eq:1.5}\\
            =\  & \condexp{}{\norm{B\pth{I - \alpha_t C}\grad L\pth{w_{t-1}^b}}^2}{\calF_0^b} + \alpha_t^2c_b \condexp{}{  \left(\frac{1}{n}\sum_{i=1}^n\norm{BC\grad L_i\pth{w_{t-1}^b}}^2 - \norm{BC\grad L\pth{w_{t-1}^b}}^2\right) }{\calF_0} \nonumber\\
            \  & - \norm{\condexp{}{ B\pth{I - \alpha_t C}\grad L\pth{w_{t-1}^b} }{\calF_0^b}}^2 \nonumber\\
            =\  & \var\left( B\pth{I - \alpha_t C}\grad L\pth{w_{t-1}^b} \, \middle| \calF_0 \right) + \alpha_t^2c_b  \left(\frac{1}{n}\sum_{i=1}^n\condexp{}{\norm{BC\grad L_i\pth{w_{t-1}^b}}^2}{\calF_0} - \condexp{}{\norm{BC\grad L\pth{w_{t-1}^b}}^2 }{\calF_0} \right) \nonumber\\
            =\  & \var\left( B\pth{I - \alpha_t C}\grad L\pth{w_{t-1}^b} \, \middle| \calF_0 \right) + \frac{\alpha_t^2c_b}{n^2}\sum_{i\neq j}\condexp{}{\norm{BC\grad L_i\pth{w_{t-1}^b} - BC\grad L_j\pth{w_{t-1}^b}}^2}{\calF_0}, \label{eq:1.6}
        \end{aligned}$$ where is by Lemma \[lemma:1\]. By induction, we know that the first term of is a decreasing function of $b$. Taking $A_i = BC, A_j = -BC, A_k = 0, k \in [n] \backslash \{i, j\}$ in Theorem \[thm:1\], we know that $$\condexp{}{\norm{BC\grad L_i\pth{w_{t-1}^b} - BC\grad L_j\pth{w_{t-1}^b}}^2}{\calF_0}$$ is also a decreasing function of $b$. Note that $\frac{\alpha_t^2c_b}{n^2}$ decreases as $b$ increases. By Lemma \[lemma:product\] we learn that is a decreasing function of $b$ and hence we have completed the induction.

We have $$\begin{aligned}
            \var\left( Bg_t^b \, \middle| \calF_0 \right) &= \condexp{}{\norm{Bg_t^b}^2}{\calF_0} - \norm{\condexp{}{Bg_t^b}{\calF_0}}^2 \\
            &= \condexp{}{\condexp{}{\norm{Bg_t^b}^2}{\calF_t^b}}{\calF_0} - \norm{\condexp{}{\condexp{}{Bg_t^b}{\calF_t^b}}{\calF_0}}^2 \\
            &= c_b \left(\frac{1}{n}\sum_{i=1}^n\condexp{}{\norm{B\grad L_i\pth{w_t^b}}^2}{\calF_0} - \condexp{}{\norm{B\grad L\pth{w_t^b}}^2}{\calF_0}\right) \\
            & \quad + \condexp{}{\norm{B\grad L\pth{w_t^b}}^2}{\calF_0} - \norm{\condexp{}{B\grad L\pth{w_t^b}}{\calF_0}}^2 \\
            &= \frac{c_b}{n^2} \sum_{i\neq j}\condexp{}{\norm{B \grad L_i\pth{w_{t}^b} - B\grad L_j \pth{w_{t}^b}}^2}{\calF_0} + \var\left(B\grad L\pth{w_t^b} \, \middle| \calF_0 \right).
        \end{aligned}$$

Taking $A_i = B, A_j = -B, A_k = 0, k \in [n] \backslash \{i, j\}$ in Theorem \[thm:1\], we know that $$\condexp{}{\norm{B \grad L_i\pth{w_{t}^b} - B\grad L_j \pth{w_{t}^b}}^2}{\calF_0}$$ is a decreasing and non-negative function of $b$ for all $i, j \in [n]$. By Theorem \[theorem:varianceFullGradient\], we know that $\var\left(B\grad L\pth{w_t^b} \, \middle| \calF_0 \right)$ is also a decreasing function of $b$. Therefore, $ \var\left(Bg_t^b \, \middle|\calF_0 \right)$, as the sum of two decreasing functions of $b$, is also a decreasing function of $b$.

Simply taking $B = I_p$ in Theorem \[thm:1\] yields the proof.

Proofs for Results in \[subsec:linearNetwork\]
----------------------------------------------

We often rely on the trivial facts that $x_1x_2^\transp = x_1 I_p x_2^\transp$ and $x_1x_2^\transp x_3x_4^\transp = x_1x_2^\transp I_p x_3x_4^\transp$.

\[lemma:rewriteTrace\] Given a multiplicative term of parameter matrices $\sth{u_i v_i^\transp: u_i, v_i \in \reals^{p}, i \in [n_1]} \cup \{A_j: A_j \in \reals^{p\times p}, j \in [n_2]\} $ and constant matrix $\sth{I_p}$ such that $\deg (u_1 v_1^\transp; M) \ge 1$, we have $$\trace{M} = v_{1}^\transp M' u_{1},$$ where $M'$ is a multiplicative term of parameter matrices $\sth{u_i v_i^\transp: u_i, v_i \in \reals^{p}, i \in [n_1]} \cup \sth{A_j: A_j \in \reals^{p\times p}, j \in [n_2]} $ and constant matrix $\sth{I_p}$ such that $\deg(M) = \deg(M') + 1, \deg(A_j; M) = \deg (A_j; M'), j\in [n_2], \deg(u_i v_i^\transp; M) = \deg (u_i v_i^\transp; M'), i \in \fromto{2}{n_1}$ and $\deg(u_{1} v_{1}^\transp; M) = \deg (u_{1} v_{1}^\transp; M') + 1$.

By the definition of multiplicative terms, we know that there exist two multiplicative terms $M_1, M_2$ of parameter matrices $\sth{u_i v_i^\transp: u_i, v_i \in \reals^{p}, i \in [n_1]} \cup \{A_j: A_j \in \reals^{p\times p}, j \in [n_2]\} $ and constant matrix $\sth{I_p}$ such that $$M = M_1 u_1 v_1^\transp M_2,$$ where $\deg(M) = \deg(M_1) + \deg(M_2) + 1, \deg(A_j; M) = \deg (A_j; M_1) + \deg (A_j; M_2), j\in [n_2], \deg(u_i v_i^\transp; M) = \deg (u_i v_i^\transp; M_1) + \deg (u_i v_i^\transp; M_2), i \in \fromto{2}{n_1}$ and $\deg (u_1 v_1^\transp; M) = \deg (u_1 v_1^\transp; M_1) + \deg (u_1 v_1^\transp; M_2) + 1$. Therefore we have $$\trace{M} = \trace{M_1 u_1 v_1^\transp M_2} = \trace{v_1^\transp M_2 M_1 u_1} = v_1^\transp M_2 M_1 u_1.$$ Note that $M' = M_2 M_1$ satisfies that $\deg(M') = \deg(M_1) + \deg(M_2), \deg(A_j, M') = \deg (A_j; M_1) + \deg (A_j; M_2), j\in [n_2], \deg(u_i v_i^\transp; M) = \deg (u_i v_i^\transp; M_1) + \deg (u_i v_i^\transp; M_2), i \in \fromto{2}{n_1}$ and $\deg (u_1 v_1^\transp; M') = \deg (u_1 v_1^\transp; M_1) + \deg (u_1 v_1^\transp; M_2) + 1$. We have finished the proof.

The following two lemmas focus on the expectation of the product of quadratic forms of the standard normal samples. Lemma \[lemma:4\] focuses on single sample while \[lemma:normalProperty\] focuses on the same form with $b$ i.i.d. samples drawn from the standard normal distribution.

\[lemma:4\] Given matrices $A_j \in \reals^{p \times p}, j\in [m-1]$, we have $$\mathbb{E}_{x\sim\calN(0,I_p)} \qth{x x^\transp A_1 x x^\transp A_2 \cdots A_{m-1} x x^\transp} = \sum_{i=1}^{N_m} \prod_{k=1}^{n_i} \trace{M_{ik}}M_{i0},$$ where $N_m$ and $ n_i,i\in [N_m]$ are constants depending on $m$ and $\sth{M_{ik}, k\in\fromto{0}{n_i}, i\in [N_m]}$ are multiplicative terms of parameter matrices $\sth{A_j,j\in [m-1]}$ and constant matrix $\sth{I_p}$. Furthermore, for every $i\in [N_m]$, we have $\sum_{k=0}^{n_i} \deg(A_j; M_{ik}) = 1, j\in [m-1]$ and therefore $\sum_{k=0}^{n_i} \deg \pth{M_{ik}} = m - 1$.

See [@magnus1978moments].

\[lemma:normalProperty\] We are given matrices $A_j \in \reals^{p \times p}, j\in [m-1] $ and random vectors $x_i, i\in [b]$ independently and identically drawn from $\calN(0, I_p)$. We assume that the multi-set $\calS = \sth{i_j, i_j': j \in [m]}$ satisfies that for every $i \in \calS$, $i$ is an element of $[b]$ and the number of appearance of $i$ in $\calS$ is even. Then $$\mathbb{E}_{x_i\sim\calN(0,I_p)} \qth{ x_{i_1} x_{i_1'}^\transp A_1 x_{i_2} x_{i_2'}^\transp A_2 \cdots A_{m-1} x_{i_m} x_{i_m'}^\transp} = \sum_{i=1}^{N_m} \prod_{k=1}^{n_i} \trace{M_{ik}}M_{i0}, \label{eq:normalProduct}$$ where $N_m$ and $n_i$ are constants depending on $m$ (and independent of $b$) and $M_{ik}, k \in \fromto{0}{n_i}, i\in [N_m]$ are multiplicative terms of parameter matrices $\sth{A_j, j\in[m-1] }$ and constant matrix $\{I_p\}$. Furthermore, for every $i\in [N_m]$, we have $\sum_{k=0}^{n_i} \deg(A_j; M_{ik}) = 1, j\in [m-1] $ and therefore $\sum_{k=0}^{n_i} \deg \pth{M_{ik}} = m-1$.

Let $\beta_i, i \in [b]$ be the number of appearances of $i$ in $\calS$, which are even by assumption. We induct on the quantity $N = \sum_{i=1}^b \mathbbm{1}\sth{\beta_i \neq 0}$.

For the base case of $N = 1$, all elements in the multi-set $\calS$ have the same value. Without loss of generality, we assume $i_j = i_j' = 1, j\in [m]$. Then $$\mathbb{E}_{x_i\sim\calN(0,I_p)} \qth{x_{i_1} x_{i_1'}^\transp A_1 x_{i_2} x_{i_2'}^\transp \cdots A_{m-1} x_{i_m} x_{i_m'}^\transp} = \mathbb{E}_{x_1\sim\calN(0,I_p)} \qth{x_1 x_1^\transp A_1 x_1 x_1^\transp \cdots A_{m-1} x_1 x_1^\transp},$$ which is the statement of Lemma \[lemma:4\].

Suppose the statement holds for $N \ge 1$, and we consider the case of $N + 1$. Note that $x_{i_j'}^\transp A_j x_{i_{j+1}} = x_{i_{j+1}}^\transp A_j x_{i_{j}'}$ is a scalar so that we can move it around without changing the value of the expression[^2]. We distinguish two cases.

-   Let $i_1 \neq i_m'$. Without loss of generality, we assume $i_1 = 1$. We can always change the order of $x_{i_j'}^\transp A_j x_{i_{j+1}}, j \in [m-1]$ (and flip it to be $x_{i_{j+1}}^\transp A_j x_{i_j'}$ if necessary) such that all $x_1$’s appear in the form of $x_1 x_1^\transp$: $$\begin{aligned}
     \label{eq:3.10}
                     x_{i_1} x_{i_1'}^\transp A_1 x_{i_2} x_{i_2'}^\transp A_2 \cdots A_{m-1} x_{i_m} x_{i_m'}^\transp &= x_1 \pth{x_{i_1'}^\transp A_1 x_{i_2} x_{i_2'}^\transp A_2 \cdots A_{m-1} x_{i_m}} x_{i_m'}^\transp \notag \\
                     &= x_1 x_1^\transp \tilde{A}_1 x_1 x_1^\transp \tilde{A}_2 \cdots \tilde{A}_{\frac{\beta_1}{2}-1} x_1 x_1^\transp \tilde{A}_{\frac{\beta_1}{2}} \tilde{x} x_{i_m'}^\transp \notag
                \end{aligned}$$ where $\tilde{x} \in \sth{x_i, i\in[b]}, \tilde{x} \neq x_1$ and $\tilde{A}_i$’s are multiplicative terms of parameter matrices $\{x_u x_v^\transp: u,v \in \fromto{2}{b}\} \cup \sth{A_j: j \in [m-1]}$ and constant matrix $\sth{I_p}$ such that $\sum_{u,v\in \fromto{2}{b}} \sum_{k=1}^{\frac{\beta_1}{2}} \deg(x_u x_v^\transp; \tilde{A}_k) = m - \frac{\beta_1}{2} - 1$ and $\sum_{k=1}^{\frac{\beta_1}{2}} \deg(A_j; \tilde{A}_k) = 1, j \in [m-1]$[^3].

    Applying Lemma \[lemma:4\] and the law of iterative expectations, we have $$\begin{aligned}
                    \mathbb{E}_{x_i\sim\calN(0,I_p)} \qth{x_{i_1} x_{i_1'}^\transp A_1 x_{i_2} x_{i_2'}^\transp \cdots A_{m-1} x_{i_m} x_{i_m'}^\transp} &= \mathbb{E}_{x_1,\cdots, x_b} \qth{x_1 x_1^\transp \tilde{A}_1 x_1 x_1^\transp \tilde{A}_2 \cdots \tilde{A}_{\frac{\beta_1}{2}-1} x_1 x_1^\transp \tilde{A}_{\frac{\beta_1}{2}} \tilde{x} x_{i_m'}^\transp} \\
                    &= \mathbb{E}_{x_2, \cdots, x_b} \qth{\pth{\sum_{i=1}^{N_m} \prod_{k=1}^{n_i} \trace{M_{ik}}M_{i0}} \tilde{A}_{\frac{\beta_1}{2}} \tilde{x} x_{i_m'}^\transp} \\
                    &= \sum_{i=1}^{N_m} \mathbb{E}_{x_2, \cdots, x_b} \qth{\pth{\prod_{k=1}^{n_i} \trace{M_{ik}}M_{i0}} \tilde{A}_{\frac{\beta_1}{2}} \tilde{x} x_{i_m'}^\transp},
                \end{aligned}$$ where $N_m$ and $n_i$ are constant depending on $m$ (and independent of $b$) and $M_{ik}, k \in \fromto{0}{n_i}, i\in [N_m]$ are multiplicative terms of parameter matrices $\sth{\tilde{A}_j, j\in[\frac{\beta_1}{2}-1] }$ and constant matrix $\{I_p\}$. Furthermore, for every $i\in [N_m]$, we have $\sum_{k=0}^{n_i} \deg(\tilde{A}_j; M_{ik}) = 1, j\in [\frac{\beta_1}{2}-1] $ and therefore $\sum_{k=0}^{n_i} \deg \pth{M_{ik}} = \frac{\beta_1}{2}-1$.

    Combining the definition of $\tilde{A}_j$’s, we know that $M_{ik}, k \in \fromto{0}{n_i}, i\in [N_m]$ are multiplicative terms of parameter matrices $\{x_u x_v^\transp: u,v \in \fromto{2}{b}\} \cup \sth{A_j: j \in [m-1]}$ and constant matrix $\sth{I_p}$ such that for every $i\in [N_m]$, we have $\sum_{u,v\in \fromto{2}{b}} \sum_{k=0}^{n_i} \deg(x_u x_v^\transp; M_{ik}) = m - \frac{\beta_1}{2} - 1$ and $\sum_{k=0}^{n_i} \deg(A_j; M_{ik}) = 1, j \in [m-1]$.

    Applying Lemma \[lemma:rewriteTrace\], for every $ k \in \fromto{0}{n_i}$ and every $ i\in [N_m]$, there exists $u_{ik}, v_{ik} \in \sth{x_j: j \in \fromto{2}{b}}$ and multiplicative term $M_{ik}'$ of parameter matrices $\{x_u x_v^\transp: u,v \in \fromto{2}{b}\} \cup \sth{A_j: j \in [m-1]}$ and constant matrix $\sth{I_p}$ such that $$\trace{M_{ik}} = u_{ik}^\transp M_{ik}' v_{ik}.$$

    Therefore, we have $$\begin{aligned}
                    \pth{\prod_{k=1}^{n_i} \trace{M_{ik}}M_{i0}} \tilde{A}_{\frac{\beta_1}{2}} \tilde{x} x_{i_m'}^\transp = \prod_{k=1}^{n_i} \pth{u_{ik}^\transp M_{ik}' v_{ik}} M_{i0} \tilde{A}_{\frac{\beta_1}{2}} \tilde{x} x_{i_m'}^\transp = M_{i0} \tilde{A}_{\frac{\beta_1}{2}} \tilde{x} \prod_{k=1}^{n_i} \pth{u_{ik}^\transp M_{ik}' v_{ik}} x_{i_m'}^\transp \triangleq U_i.
                \end{aligned}$$

    Note that for every $i \in [N_m]$, we have $$\begin{aligned}
                    \sum_{j=1}^{m-1} \deg(x_i; A_j) = \sum_{k=1}^{n_i} \deg(x_i; M_{ik}') + \deg(x_i; M_{i0}) + \deg\pth{x_i; \tilde{A}_{\frac{\beta_1}{2}}} + \deg(x_i; \tilde{x}) + \deg\pth{x_i; x_{i_m'}^\transp},
                \end{aligned}$$ and for every $j\in [m-1]$, we have $$\begin{aligned}
                    \sum_{k=1}^{n_i} \deg(A_j; M_{ik}') + \deg(A_j; M_{i0}) + \deg\pth{A_j; \tilde{A}_{\frac{\beta_1}{2}}} = 1.
                \end{aligned}$$

    In other words, for every $i \in [N_m]$, $U_i$ has the form of $\hat{A}_0 x_{\hat{i}_1} x_{\hat{i}_1'}^\transp \hat{A}_1 x_{\hat{i}_2} x_{\hat{i}_2'}^\transp \cdots \hat{A}_{m-1} x_{\hat{i}_{m'}} x_{i_{m}'}^\transp \hat{A}_{m'}$ but there is no appearance of $x_1$. Here $x_{\hat{i}_j}, x_{\hat{i}_j} \in \sth{x_j, j \in \fromto{2}{b}}$, and $\hat{A}_i, i\in \fromto{0}{m}$ are multiplicative terms of parameter matrices $\sth{A_j, j\in[m-1] }$ and constant matrix $\{I_p\}$. Furthermore, for every $j \in [m-1]$, we have $\sum_{k=0}^{n_i} \deg(A_j; \hat{A}_i) = 1$. Note that here we use the liberty of adding identity matrices if more than two consecutive $x$’s appear. Since we have reduced $N+1$ by one, we can use induction on $x_{\hat{i}_1} x_{\hat{i}_1'}^\transp \hat{A}_1 x_{\hat{i}_2} x_{\hat{i}_2'}^\transp \cdots \hat{A}_{m-1} x_{\hat{i}_{m'}} x_{i_{m}'}^\transp$ and finish the proof. The two constant matrices $\hat{A}_0$ and $\hat{A}_m$ do not change the result of expectation since $\mathbb{E}\pth{\hat{A}_0 X \hat{A}_{m'}}=\hat{A}_0 \Expect(X) \hat{A}_{m'}$.

-   If $i_1 = i_m'$, without loss of generality we assume, $i_1' = 1$ and $i_1' \neq i_1$ (note that all $x_{i_j'}^\transp A_j x_{i_{j+1}}, j \in [m-1]$ are inter-changeable and there is at least one element in $\calS$ that is not equal to $i_1$). We change the orders of $x_{i_j'}^\transp A_j x_{i_{j+1}}, j \in [m-1]$ (and flip it to be $x_{i_{j+1}}^\transp A_j x_{i_j'}$ if necessary) such that all $x_1$’s appear in a consecutive form of $x_1 x_1^\transp$: $$\begin{aligned}
                    x_{i_1} x_{i_1'}^\transp A_1 x_{i_2} x_{i_2'}^\transp A_2 \cdots A_{m-1} x_{i_m} x_{i_m'}^\transp &= x_{i_1} \pth{x_{i_1'}^\transp A_1 x_{i_2} x_{i_2'}^\transp A_2 \cdots A_{m-1} x_{i_m}} x_{i_m'}^\transp \\
                    &= x_{i_1} \pth{\tilde{x}_1^\transp \tilde{A}_0 \qth{x_1 x_1^\transp \tilde{A}_1 \cdots \tilde{A}_{\frac{\beta_1}{2} - 1} x_1 x_1^\transp} \tilde{A}_{\frac{\beta_1}{2}} \tilde{x}_2} x_{i_m'}^\transp,
                \end{aligned}$$ where $ \tilde{x}_1, \tilde{x}_2 \in \sth{x_i, i\in[b]}, \tilde{x}_1, \tilde{x}_2 \neq x_1$ and $\tilde{A}_i$’s are multiplicative terms of parameter matrices $\{x_u x_v^\transp: u,v \in \fromto{2}{b}\} \cup \sth{A_j: j \in [m-1]}$ and constant matrix $\sth{I_p}$ such that $$\sum_{u,v\in \fromto{2}{b}} \sum_{k=0}^{\frac{\beta_1}{2}} \deg(x_u x_v^\transp; \tilde{A}_k) = m - \frac{\beta_1}{2} - 2$$ and $\sum_{k=0}^{\frac{\beta_1}{2}} \deg(A_j; \tilde{A}_k) = 1, j \in [m-1]$. The remaining reasoning is the same as the previous case.

If one of the $\beta_i$ numbers of appearance of $x_j, j\in[b]$ is odd, then it is easy to see that the result in is the zero matrix.

By and we have $$\begin{aligned}
            M = \prod_{i=1}^m \trace{M_i} M_0 = \frac{1}{b^d}\sum_{k=1}^{b^d} \prod_{i=1}^m \trace{M_{ki}} M_{k0},
        \end{aligned}$$ where each $M_{ki}, k\in [b^d], i \in \fromto{0}{m}$ is a multiplicative term of parameter matrices $\sth{x_{t,i}x_{t,i}^\transp, i\in [b]}$ and constant matrices $\sth{W_{t,1}^b, W_{t,2}^b, \calW_t^b}$. Let $\widetilde{M}_k = \prod_{i=1}^m \trace{M_{ki}} M_{k0}, k\in \qth{b^d}$. We split set $\sth{\widetilde{M}_k: k\in \qth{b^d}}$ into disjoint and non-empty sets (equivalent classes) $S_1, \ldots, S_{n_M}$ such that

1.  for every $i \in [n_M]$ and every $M_1, M_2 \in S_i$, we have $ \condexp{}{M_1}{\calF_t^b} = \condexp{}{M_2}{\calF_t^b} $,

2.  for every $i, j \in [n_M], i \neq j$ and every $M_1 \in S_i \textrm{ and } M_2 \in S_j$, we have $\condexp{}{M_1}{\calF_t^b} \neq \condexp{}{M_2}{\calF_t^b} $.

Note that $\cup_{i=1}^{n_M} S_i = \sth{\widetilde{M}_k: k\in \qth{b^d}}$. Let $\hat{M}_k \in S_k$ represent the equivalent class $S_k$ (it can be any member of $S_k$). For every $i \in [n_M]$, we can always write $|S_i| = e_{i,0} + e_{i,1} b + \cdots + e_{i,d} b^{d}$ such that $e_{i,j} \in \mathbb{N}, e_{i,j} < b, j\in\fromto{0}{d}$ (actually $e_{i,j}$’s are the digits of the base-$b$ representation of $|S_i|$). Then we have $$\begin{aligned}
            \condexp{}{M}{\calF_t^b} &= \condexp{}{\frac{1}{b^d}\sum_{k=1}^{b^d} \widetilde{M}_k}{\calF_t^b} = \frac{1}{b^d} \condexp{}{\sum_{i=1}^{n_M} \pth{e_{i,0} + e_{i,1} b + \cdots + e_{i,d} b^{d}} \hat{M}_i }{\calF_t^b}\notag\\ 
            &= \frac{1}{b^d}\sum_{i=1}^{n_M} \pth{e_{i,0} + e_{i,1} b + \cdots + e_{i,d} b^{d}} \condexp{}{\hat{M}_i }{\calF_t^b} \label{eq:2.5} \\
            &= \sum_{i=1}^{n_M} \pth{e_{i,d} + e_{i,d-1} \frac{1}{b} + \cdots + e_{i,0} \frac{1}{b^d} } \condexp{}{\hat{M}_i }{\calF_t^b}. \notag
        \end{aligned}$$ It is important to note that $n_M$, the number of different equivalent classes, is independent of $b$. This follows from the fact that each $\condexp{}{\tilde{M}_k}{\calF_t^b}$ (and so as $\condexp{}{\hat{M}_k}{\calF_t^b}$) includes a finite number of weight matrices $W_{t,1}^b$ and $W_{t,2}^b$ with degree less than or equal to $3d + \sum_{i=0}^m \pth{\deg\pth{W_{t,1}^b; M_i} + \deg(W_{t,2}^b; M_i)}$ (see Lemma \[lemma:normalProperty\]). Thus the number of partition sets is bounded by a quantity independent of $b$.

Note that each $M_{ki}$ can be represented as $$M_{ki} = A_0^{ki} x_{t,i_1}^{ki} {x_{t,i_1}^{ki}}^\transp A_1^{ki} \cdots A_{d_i-1}^{ki} x_{t,i_{d_i}}^{ki} {x_{t,i_{d_i}}^{ki}}^\transp A_{d_i}^{ki}$$ for some matrices $A_0^{ki}, \ldots, A_{d_i}^{ki} $ that are multiplicative term of parameter matrices $\sth{W_{t,1}^b, W_{t,2}^b and \calW_t^b}$ constant matrix $\sth{I_p}$ (we stress again that some $A$ matrices can be identities, based on the definition of multiplicative terms), and $x_{t,i_1}^{ki}, \ldots, x_{t,i_{d_i}}^{ki} \in \sth{x_{t,1}, \ldots, x_{t,b}}$. We have $$\begin{aligned}
            \trace{M_{ki}} &= \trace{ A_0^{ki} x_{t,i_1}^{ki} {x_{t,i_1}^{ki}}^\transp A_1^{ki} \cdots A_{d_i-1}^{ki} x_{t,i_{d_i}}^{ki} {x_{t,i_{d_i}}^{ki}}^\transp A_{d_i}^{ki} }\\
            &= {x_{t,i_{d_i}}^{ki}}^\transp A_{d_i}^{ki} A_0^{ki} x_{t,i_1}^{ki} {x_{t,i_1}^{ki}}^\transp A_1^{ki} \cdots A_{d_i-1}^{ki} x_{t,i_{d_i}}^{ki}. 
        \end{aligned}$$ For every $k \in \qth{b^d}$, we have $$\begin{aligned}
                \prod_{i=1}^m \trace{M_{ki}} M_{k0} &= \qth{\prod_{i=1}^m {x_{t,i_{d_i}}^{ki}}^\transp A_{d_i}^{ki} A_0^{ki} x_{t,i_1}^{ki} {x_{t,i_1}^{ki}}^\transp A_1^{ki} \cdots A_{d_i-1}^{ki} x_{t,i_{d_i}}^{ki}} A_0^{k0} x_{t,i_1}^{k0} {x_{t,i_1}^{k0}}^\transp A_1^{k0} \cdots A_{d_0-1}^{k0} x_{t,i_{d_0}}^{k0} {x_{t,i_{d_0}}^{k0}}^\transp A_{d_0}^{k0}\\
                &= \qth{\prod_{i=1}^m {x_{t,i_{d_i}}^{ki}}^\transp A_{d_i}^{ki} A_0^{ki} x_{t,i_1}^{ki} {x_{t,i_1}^{ki}}^\transp A_1^{ki} \cdots A_{d_i-1}^{ki} x_{t,i_{d_i}}^{ki}} \qth{{x_{t,i_1}^{k0}}^\transp A_1^{k0} \cdots A_{d_0-1}^{k0} x_{t,i_{d_0}}^{k0} } A_0^{k0} x_{t,i_1}^{k0} {x_{t,i_{d_0}}^{k0}}^\transp A_{d_0}^{k0},
        \end{aligned}$$ which can be rewritten as $$\widetilde{M}_k =  \prod_{i=1}^m \trace{M_{ki}} M_{k0} = \pth{\prod_{j=1}^d x_{t, \bar{i}_j}^\transp A_j^k x_{t, \bar{i}_j'}} A_0^{k0} x_{t,i_1}^{k0} {x_{t,i_{d_0}}^{k0}}^\transp A_{d_0}^{k0}.$$

Note that the randomness of each $\tilde{M}_k$ given $\calF_t^b$ only comes from the randomness of $x_{t,j}$’s, i.e. for all $k \in \qth{b^d} $ we have $$\begin{aligned}
            \condexp{}{\tilde{M}_k}{\calF_t^b} &= \mathbb{E}_{x_{t,j} \sim \calN(0, I)} \qth{\pth{\prod_{j=1}^d x_{t, i_j}^\transp A_j^k x_{t, i_j'}} A_0^k x_{t, i_0'}x_{t, i_0}^\transp {A_0^k}'}\notag \\
            &=\mathbb{E}_{x_{t,j} \sim \calN(0, I)} \qth{ A_0^k x_{t, i_0'} \pth{\prod_{j=1}^d x_{t, i_j}^\transp A_j^k x_{t, i_j'}} x_{t, i_0}^\transp {A_0^k}'} \label{eq:2.6}\\
            &= \sum_{i=1}^{n_M^k} \prod_{j=1}^{n_i^k} \trace{\tilde{M}_{ij}^k} \tilde{M}_{i0}^k, \notag
        \end{aligned}$$ where the last equation comes from Lemma \[lemma:normalProperty\]. Here $n_M^k, n_i^k, i \in \qth{n_M^k}, k \in \qth{b^d}$ are constants independent of $b$, $M_{ij}^k$’s are multiplicative terms of parameter matrices $\sth{W_{t,1}^b, W_{t,2}^b, \calW_t^b}$ and constant matrix $\sth{I_p}$ such that for every $i \in \qth{n_M^k}$, we have $$\label{eq:3.15}
            \sum_{j=0}^{n_i^k} \deg\pth{\calW_t^b; \tilde{M}_{ij}^k} = d$$ and $$\label{eq:3.16}
            \sum_{j=0}^{n_i^k} \pth{\deg\pth{W_{t,1}^b; \tilde{M}_{ij}^k} + \deg\pth{W_{t,2}^b; \tilde{M}_{ij}^k}} = d + \sum_{r=0}^m \pth{\deg\pth{W_{t,1}^b; M_r} + \deg(W_{t,2}^b; M_r)}.$$

These degree relationships can be observed from , , and the fact that each $g_{t,1}^b$ or $g_{t,1}^b$ contributes one $\calW_t^b$ and one of $W_{t,1}^b$ or $W_{t,2}^b$ in $\prod_{j=1}^{n_i^k} \trace{\tilde{M}_{ij}^k} \tilde{M}_{i0}^k$. Note that $\calW_t = W_{t,2}^bW_{t,2}^b - W_2^*W_1^*$. For every $i \in \qth{n_M^k}$, if we replace all appearances of $\calW_t^b$ in $\prod_{j=1}^{n_i^k} \trace{\tilde{M}_{ij}^k}\tilde{M}_{i0}^k$ and expand all parentheses of $\pth{W_{t,2}^bW_{t,2}^b - W_2^*W_1^*}$, we have $$\begin{aligned}
 \label{eq:3.17}
            \prod_{j=1}^{n_i^k} \trace{\tilde{M}_{ij}^k} \tilde{M}_{i0}^k = \sum_{l=1}^{2^d} \prod_{j=1}^{n_i^k} \trace{\tilde{M}_{ij}^{kl}} \tilde{M}_{i0}^{kl},
        \end{aligned}$$ where $\tilde{M}_{ij}^{kl}$’s are multiplicative terms of parameter matrices $\sth{W_{t,1}^b, W_{t,2}^b}$ and constant matrices $\sth{W_1^*, W_2^*}$ such that $$\label{eq:3.18}
            \sum_{j=0}^{n_i^k} \pth{\deg\pth{W_{t,1}^b; \tilde{M}_{ij}^{kl}} + \deg\pth{W_{t,2}^b; \tilde{M}_{ij}^{kl}}} \le 3d + \sum_{r=0}^m \pth{\deg\pth{W_{t,1}^b; M_r} + \deg(W_{t,2}^b; M_r)},$$ where the inequality comes from and and the fact that each $g_{t,1}^b$ or $g_{t,2}^b$ contributes 2 or 0 degrees in the form of $W_{t,2}^bW_{t,1}^b$ or $W_2^*W_1^*$, respectively.

Combining , and , we have $$\begin{aligned}
            \condexp{}{M}{\calF_t^b} &= \sum_{k=1}^{n_M} \pth{e_{k,d} + e_{k,d-1} \frac{1}{b} + \cdots + e_{k,0} \frac{1}{b^d} } \condexp{}{\hat{M}_k }{\calF_t^b} \\
            &= \sum_{k=1}^{n_M} \pth{e_{k,d} + e_{k,d-1} \frac{1}{b} + \cdots + e_{k,0} \frac{1}{b^d} } \sum_{i=1}^{n_M^{s_k}} \sum_{l=1}^{2^d} \prod_{j=1}^{n_i^k} \trace{\tilde{M}_{ij}^{kl}} \tilde{M}_{i0}^{kl} \\
            &= N_0 + N_1\frac{1}{b} + \cdots + N_d \frac{1}{b^d},
        \end{aligned}$$ where $$\label{eq:3.19}
            N_r = \sum_{k=1}^{n_M} e_{k,d-r} \pth{\sum_{i=1}^{n_M^{s_k}} \sum_{l=1}^{2^d} \prod_{j=1}^{n_i^k} \trace{\tilde{M}_{ij}^{kl}} \tilde{M}_{i0}^{kl}}.$$

Note that all constants in are independent of $b$ and combining with , we have finished the proof.

Simply using the fact that $W_{t, i}^b = W_{t-1, i}^b - \alpha_t g_{t-1, i}^b, i = 1, 2$, if we replace each $W_{t,i}^b$ in the left-hand-side of (\[eq:3.19\]) by $W_{t-1, i}^b - \alpha_t g_{t-1, i}^b$ and expand all the parentheses, then each $M_i, i \in \fromto{0}{m}$ becomes the sum of $2^{d_i}$ multiplicative terms of parameter matrices $\sth{g_{t,1}^b, g_{t,2}^b}$ and constant matrices $\sth{W_{t,1}^b, W_{t,2}^b, W_1^*, W_2^*}$ with degree at most $d_i$. As a result, $\prod_{i=1}^m \trace{M_i} M_0$ becomes the sum of $2^d$ terms in the form of $\prod_{i=1}^{m} \trace{M_{ik}} M_{0k}$ where $\deg\pth{M_{ik}} \le 2^{d_i}$, and therefore $\sum_{i=0}^m \deg\pth{M_{ik}} \le \prod_{i=0}^m 2^{d_i} = d$.

We use induction on $t$ to show this result. The base case of $t = 0$ it is the same as the statement in Lemma \[lemma:GtoW\].

Suppose that the statement holds for $t \ge 0$, and we consider the case of $t+1$. By Lemma \[lemma:GtoW\], there exists a set of multiplicative terms $\sth{M_{t+1,i,j}^k, i\in [m_{t+1,k}], j \in \fromto{0}{m_{t+1,k,i}}, k \in \fromto{0}{d}}$ of parameter matrices $\sth{W_{t+1,1}^b, W_{t+1,2}^b}$ and constant matrices $\sth{W_1^*, W_2^*}$ such that $$\begin{aligned}
 \label{eq:3.20}
                \condexp{}{M}{\calF_{t+1}^b} = N_{t+1,0} + N_{t+1,1} \frac{1}{b} + \cdots + N_{t+1,d} \frac{1}{b^d},
            \end{aligned}$$ where $N_{t+1,k} = \sum_{i=1}^{m_{t+1,k}} \prod_{j=1}^{m_{t+1,k,i}} \trace{M_{t+1,i,j}^k} M_{t+1,i,0}^k, k\in \fromto{0}{d}$. Here $m_{t+1,k}, m_{t+1,k,i}$ are constants independent of $b$, and $\sum_{j=0}^{m_{t+1,k,i}} \deg\pth{M_{t+1,i,j}^k} \le 3d + d'$.

For each $i \in \qth{m_{t+1, k}}$ and each $k \in \fromto{0}{d}$, by Lemma \[lemma:WtoG\], there exists a set of multiplicative terms $\{M_{t,i,j,k,l}, j\in\qth{m_{t+1,i,k}}, l \in [d_{t,i,k}]\}$ of parameter matrices $\sth{g_{t,1}^b, g_{t,2}^b}$ and constant matrices $\sth{W_{t,1}^b, W_{t,2}^b, W_1^*, W_2^*}$ such that $$\begin{aligned}
 \label{eq:3.21}
                \prod_{j=1}^{m_{t+1,k,i}} \trace{M_{t+1,i,j}^k} M_{t+1,i,0}^k = \sum_{l=1}^{d_{t,i,k}} \prod_{j=1}^{m_{t+1,k,i}} \trace{M_{t,i,j,k,l}} M_{t,i,0,k,l},
            \end{aligned}$$ where $d_{t,i,k} = 2^{\sum_{j=0}^{m_{t+1,k,i}} \pth{\deg\pth{W_{t,1}^b; M_{t,i,j,k,l}} + \deg(W_{t,2}^b; M_{t,i,j,k,l})}}$ is a constant independent of $b$ and $$\label{eq:3.23}
                \sum_{j=0}^{m_{t+1,k,i}} \deg \pth{M_{t,i,j,k,l}} \le 3d + d',$$ and $$\label{eq:3.24}
                \sum_{j=0}^{m_{t+1,k,i}} \pth{\deg \pth{W_{t,1}; M_{t,i,j,k,l}} + \deg \pth{W_{t,2}; M_{t,i,j,k,l}}} \le 3d + d'.$$

Combining and , we have for every $k \in \fromto{0}{d}$ $$N_{t+1, k} = \sum_{i=1}^{m_{t+1,k}} \sum_{l=1}^{d_{t,i,k}} \prod_{j=1}^{m_{t+1,k,i}} \trace{M_{t,i,j,k,l}} M_{t,i,0,k,l}.$$

Note that $$\begin{aligned}
            \condexp{}{M}{\calF_0} =&\; \condexp{}{\condexp{}{M}{\calF_{t+1}^b}}{\calF_0} =  \condexp{}{N_{t+1,0}}{\calF_0} + \condexp{}{N_{t+1,1}}{\calF_0} \frac{1}{b} + \cdots + \condexp{}{N_{t+1,d}}{\calF_0} \frac{1}{b^d} \notag \\
            =&\; \sum_{i=1}^{m_{t+1,0}} \sum_{l=1}^{d_{t,i,0}} \condexp{}{\prod_{j=1}^{m_{t+1,0,i}} \trace{M_{t,i,j,0,l}} M_{t,i,0,0,l}}{\calF_0} +\notag \\
            +& \; \sum_{i=1}^{m_{t+1,1}} \sum_{l=1}^{d_{t,i,1}} \condexp{}{\prod_{j=1}^{m_{t+1,1,i}} \trace{M_{t,i,j,1,l}} M_{t,i,0,1,l}}{\calF_0} \frac{1}{b} + \cdots + \notag\\
            +& \; \sum_{i=1}^{m_{t+1,d}} \sum_{l=1}^{d_{t,i,d}} \condexp{}{\prod_{j=1}^{m_{t+1,d,i}} \trace{M_{t,i,j,d,l}} M_{t,i,0,d,l}}{\calF_0} \frac{1}{b^d}, \label{eq:3.25}
        \end{aligned}$$ and each $M_{t,i,j,k,l}$ is a multiplicative term of parameter matrices $\left\{ g_{t,1}^b, g_{t,2}^b \right\}$ and constant matrices $\{W_{t,1}^b, W_{t,2}^b, W_1^*, W_2^*\}$ such that the degree is at most $1$. Therefore, by induction, for every $i,k,l$, we have $$\label{eq:3.26}
            \condexp{}{\prod_{j=1}^{m_{t+1,k,i}} \trace{M_{t,i,j,k,l}} M_{t,i,0,k,l}}{\calF_0} = N_{t,i,k,l,0} + N_{t,i,k,l,1} \frac{1}{b} + \cdots N_{t,i,k,l,q_t}\frac{1}{b^{q_t}},$$ where $q_t \le d' + \frac{1}{2}(3^t-1)(3d+d')$ and $N_{t,i,k,l,0}, \cdots, N_{t,i,k,l,q_t}$ are sum of multiplicative terms of parameter matrices $\sth{W_{0,1}^b, W_{0,2}^b}$ and constant matrices $\sth{W_1^*, W_2^*}$ with degree at most $d\cdot 3^t$.

Combining and , we can rewrite $$\condexp{}{M}{\calF_0} = N_0 + N_1 \frac{1}{b} + \cdots + N_q \frac{1}{b^q},$$ in the same form as in the statement. Here $q\le d + 3q_t \le \frac{1}{2}(3^{t+2}-1)d + \frac{1}{2}(3^{t+1}-1)d'$ and $\sum_{j=0}^{m_{ki}} \deg\pth{M_{ij}^k} \le 3 \times 3^{t} (3d + d') = 3^{t+1} (3d + d')$ follow from and .

In conclusion, we have shown that the statement holds for $t+1$, and therefore finishes the proof.

We simply note that $M$ can be written as the sum of at most $2^d$ multiplicative terms of parameter matrices $\sth{W_{t,1}^b, W_{t,2}^b, W_1^*, W_2^*}$ and constant matrix $\sth{I_0}$. Then we apply Lemmas \[lemma:GtoW\] and \[lemma:WtoG\] iteratively in the same way as in the proof of Theorem \[thm:representation\] to finish the proof.

We only show the case for $g_{t,1}$ since the proof for $g_{t,2}$ can be tackled similarly. Note that $$\begin{aligned}
            \condvar{}{g_{t,1}^b}{\calF_0} &= \condvar{}{\frac{1}{b} \sum_{i=1}^b {W_{t,2}^b}^\transp \calW_t^b x_{t,i} x_{t,i}^\transp }{\calF_0} = \frac{1}{b^2} \sum_{i=1}^b\condvar{}{{W_{t,2}^b}^\transp \calW_t^b x_{t,i} x_{t,i}^\transp }{\calF_0}\\
            &= \frac{1}{b} \condvar{}{{W_{t,2}^b}^\transp \calW_t^b x_{t,1} x_{t,1}^\transp }{\calF_0} \\
            &= \frac{1}{b}\pth{\condexp{}{\norm{{W_{t,2}^b}^\transp \calW_t^b x_{t,1} x_{t,1}^\transp}^2}{\calF_0} - \norm{\condexp{}{{W_{t,2}^b}^\transp \calW_t^b x_{t,1} x_{t,1}^\transp}{\calF_0}}^2} \\
            &= \frac{1}{b} \pth{\condexp{}{\trace{   x_{t,1} x_{t,1}^\transp {\calW_t^b}^\transp W_{t,2}^b {W_{t,2}^b}^\transp \calW_t^b x_{t,1} x_{t,1}^\transp }}{\calF_0} - \norm{\condexp{}{{W_{t,2}^b}^\transp \calW_t^b x_{t,1} x_{t,1}^\transp}{\calF_0}}^2} \\
            &= \frac{1}{b} \pth{\condexp{}{\condexp{}{\trace{x_{t,1} x_{t,1}^\transp {\calW_t^b}^\transp W_{t,2}^b {W_{t,2}^b}^\transp \calW_t^b x_{t,1} x_{t,1}^\transp }}{\calF_t^b}}{\calF_0} - \norm{\condexp{}{\condexp{}{{W_{t,2}^b}^\transp \calW_t^b x_{t,1} x_{t,1}^\transp}{\calF_t^b}}{\calF_0}}^2} \\
            &= \frac{1}{b} \pth{\condexp{}{ (p+2)\trace{{\calW_t^b}^\transp W_{t,2}^b {W_{t,2}^b}^\transp \calW_t^b} }{\calF_0} - \norm{\condexp{}{ {W_{t,2}^b}^\transp \calW_t^b }{\calF_0}}^2} \\
            &= \frac{1}{b} \pth{(p+2)\trace{\condexp{}{{\calW_t^b}^\transp W_{t,2}^b {W_{t,2}^b}^\transp \calW_t^b }{\calF_0}} - \norm{\condexp{}{ {W_{t,2}^b}^\transp \calW_t^b }{\calF_0}}^2}. \\
            &= \frac{1}{b} \pth{(p+2)\trace{\condexp{}{{\calW_t^b}^\transp W_{t,2}^b {W_{t,2}^b}^\transp \calW_t^b }{\calF_0}} - \norm{\condexp{}{ {W_{t,2}^b}^\transp \calW_t^b }{\calF_0}}^2}. \\
        \end{aligned}$$ Here we have used the fact that $\Expect_{x\sim \calN(0, I_p)}\trace{x x^\transp A x x^\transp} = (p+2) \trace{A}$. By Corollary \[thm:representationoCorollary\] we know that there exists a set of multiplicative terms $\sth{M_{ij}^k, i\in [m_k], j \in \fromto{0}{m_{ki}}, k \in \fromto{0}{q}}$ of parameter matrices $\sth{W_{0,1}^b, W_{0,2}^b}$ and constant matrices $\sth{W_1^*, W_2^*}$ such that $$\begin{aligned}
            \trace{\condexp{}{{\calW_t^b}^\transp W_{t,2}^b {W_{t,2}^b}^\transp \calW_t^b}{\calF_0}} = \gamma_0 + \gamma_1 \frac{1}{b} + \cdots + \gamma_q \frac{1}{b^q}, \label{eq:gamma1}
        \end{aligned}$$ where $\gamma_k = \sum_{i=1}^{m_k} \prod_{j=0}^{m_{ki}} \trace{M_{ij}^k}, k\in \fromto{0}{q}$. Here $m_k, m_{ki}$ and $q \le 6\cdot 3^{t}$ are constants independent of $b$, and $\sum_{j=0}^{m_{ki}} \deg\pth{M_{ij}^k} \le 6\cdot 3^{t}$. Note that $W_{0,1}^b, W_{0,2}^b$ are fixed, and we have $\gamma_k, k\in\fromto{0}{q}$ are constants independent of $b$.

Similarly we observe that there exist constants $q' \le 2\cdot 3^{t+1}$ and $\gamma_k', k\in \fromto{0}{q'}$ such that $$\begin{aligned}
            \norm{\condexp{}{ {W_{t,2}^b}^\transp \calW_t^b }{\calF_0}}^2 = \gamma_0' + \gamma_1' \frac{1}{b} + \cdots + \gamma_q' \frac{1}{b^{q'}}. \label{eq:gamma2}
        \end{aligned}$$ By defining $\gamma_i = 0, i > q$ and $\gamma_i' = 0, i > q'$, and combining and we have $$\begin{aligned}
        \condvar{}{g_{t,1}^b}{\calF_0} &= \frac{1}{b} \pth{(p+2)\trace{\condexp{}{{\calW_t^b}^\transp W_{t,2}^b {W_{t,2}^b}^\transp \calW_t^b }{\calF_0}} - \norm{\condexp{}{ {W_{t,2}^b}^\transp \calW_t^b }{\calF_0}}^2} \\
        &= \frac{p+2}{b} \pth{\gamma_0 + \gamma_1 \frac{1}{b} + \cdots + \gamma_q \frac{1}{b^q}} - \frac{1}{b} \pth{\gamma_0' + \gamma_1' \frac{1}{b} + \cdots + \gamma_q' \frac{1}{b^{q'}}} \\
        &= \sum_{k=1}^{\max\sth{q,q'}} \pth{(p+1) \gamma_k - \gamma_k'} \frac{1}{b^k}.
    \end{aligned}$$

Note that $\gamma_k$’s and $\gamma_k'$’s are all constants independent of $b$, and $\max\sth{q, q'} \le 2\cdot 3^{t+1}$. This completes the proof.

We first show that in we have $\beta_1 \ge 0$. If $r=1$, the statement obviously holds. Let us assume that the statement does not hold for $r > 1$, i.e. $\beta_1 < 0$. Taking $b$ large enough such that $\beta_1 b^{r-1} + \beta_2 b^{r-2} + \cdots + \beta_r < 0$ yields $$\condvar{}{g_{t,i}^b}{\calF_0} = \frac{1}{b^r} \pth{\beta_1 b^{r-1} + \beta_2 b^{r-2} + \cdots + \beta_r} < 0,$$ which contradicts the fact that $\condvar{}{g_{t,i}^b}{\calF_0} \ge 0$. Therefore, we have $\beta_1 \ge 0$.

Let $b_0$ be large enough such that for all $b \ge b_0$, we have $\beta_1 b^{r-1} + 2\beta_2 b^{r-2} + \cdots + r \beta_r \ge 0$. We denote $f(b) = \beta_1 \frac{1}{b} + \beta_2 \frac{1}{b^2} + \cdots + \beta_r \frac{1}{b^r} \ge 0$. For all $b > b_0$ we have $$\begin{aligned}
        f'(b) = - \frac{1}{b^{r+1}} \pth{\beta_1 b^{r-1} + 2\beta_2 b^{r-2} + \cdots + r \beta_r} \le 0.
    \end{aligned}$$ Therefore, for all $b > b_0$ we have $\pth{\condvar{}{g_{t,i}^b}{\calF_0}}' = -\frac{r}{b^{r+1}} f(b) + \frac{1}{b^r} f(b) \le 0 $, and thus $\condvar{}{g_{t,i}^b}{\calF_0}$ is a decreasing function of $b$ for all $b > b_0$.

Experimental Details {#appendix:exp}
====================

In many experiments we fix the initial and ground-truth weights (in the case of Section \[subsec:linearNetwork\]), and the learning rate. We have also tested several other random initial weights and ground-truth weights, and learning rates, and the results and conclusions are similar and not presented.

Graduate Admission Dataset with Linear Regression
-------------------------------------------------

The dataset is normalized by mean and variance of each feature. For the experiment in Figure \[fig:linear\](a), we randomly select an initial weight vectors $w_0$ and run SGD for 2,000 iterations where it appears to converge. We record all statistics at every iteration. There are in total 1,000 runs behind each observation which yields a p-value lower than 0.05. As for Figure \[fig:linear\](b), we select 20 different $b$’s and run SGD from the same initial point for 40 iterations. There are in total of 200,000 runs to make sure the p-value of all statistics are lower than 0.05. In all experiments, the learning rate is chosen to be $\alpha_t = \frac{1}{2t}, t\in[2000]$ because this rate yields a theoretical convergence guaranteed (factor 1/2 has been fine tuned).

Synthetic Dataset with Two-layer Linear Network
-----------------------------------------------

In Figure \[fig:twolayer\], we randomly select two initial weight matrices $W_{0,1}, W_{0,2}$ and the ground-truth weight matrices $W_1^*, W_2^*$. We run SGD for 1,000 iterations which appears to be a good number for convergence while there are 1,000 runs of SGD in total to again give a p-value below 0.05. We record all statistics at every iteration. The learning rate is chosen to be $\alpha_t = \frac{1}{10t}, t\in [1000]$ for the same reason as in the regression experiment.

MNIST with Fully Connected Neural Network {#appendix:MNIST}
-----------------------------------------

The images are normalized by mapping each entry to $[-1, 1]$. We run SGD for 1,000 epochs on the training set which is enough for convergence. The learning rate is a constant set to $3\cdot 10^{-3}$ (which has been tuned). For the experiment in Figure \[fig:MNIST1\], there are in total 100 runs to give us the p-value below 0.05. For the experiment in Figure \[fig:MNIST2\](a), we randomly select five different initial points and we have 50 runs for each initial point.

For the experiment corresponding to Figure \[fig:MNIST2\](b), we choose $\alpha=8$ and $\sigma=2$ as in [@1227801]. The initial weights and other hyper-parameters are chosen to be the same as in Figure \[fig:MNIST1\].

Yelp with XLNet {#appendix:yelp}
---------------

We randomly select a set of initial parameters and run Adam with two different mini-batch sizes of 32 and 64. For computational tractability reasons, for each mini-batch size there are in total of 100 runs and each run corresponds to 20 epochs. We record the variance of the stochastic gradient, loss and accuracy in every step of Adam. The statistics reported in Figure \[fig:YELP\] are averaged through each epoch. In all experiments, the learning rate is set to be $4 \cdot 10^{-5}$ and the $\epsilon$ parameter of Adam is set to be $10^{-8}$ (these two have been tuned). The stochastic gradients of all parameter matrices are clipped with threshold 1 in each iteration. We use the same setup for the learning rate warm-up strategy as suggested in [@yang2019xlnet]. The maximum sequence length is set to be 128 and we pad the sequences with length smaller than 128 with zeros.

[^1]: <https://www.kaggle.com/mohansacharya/graduate-admissions>

[^2]: For example, we can rewrite $$\begin{aligned}
                x_{i_1} x_{i_1'}^\transp A_1 x_{i_2} x_{i_2'}^\transp A_2 x_{i_3} x_{i_3'}^\transp &= x_{i_1} \pth{x_{i_1'}^\transp A_1 x_{i_2}} \qth{x_{i_2'}^\transp A_2 x_{i_3}} x_{i_3'}^\transp = x_{i_1} \qth{x_{i_2'}^\transp A_2 x_{i_3}} \pth{x_{i_1'}^\transp A_1 x_{i_2}} x_{i_3'}^\transp \\
                &= x_{i_1} \qth{x_{i_2'}^\transp \pth{x_{i_1'}^\transp A_1 x_{i_2}} A_2 x_{i_3}} x_{i_3'}^\transp = x_{i_1} \qth{x_{i_2'}^\transp A_2 \pth{x_{i_1'}^\transp A_1 x_{i_2}} x_{i_3}} x_{i_3'}^\transp.
            \end{aligned}$$

[^3]: For example, we can rewrite $$\begin{aligned}
                    &x_1 x_2^\transp A_1 x_1 x_1^\transp A_2 x_3 x_3^\transp A_3 x_1 x_2 = x_1 \pth{x_2^\transp A_1 x_1} \qth{x_1^\transp A_2 x_3} \sth{x_3^\transp A_3 x_1} x_2 = x_1 \pth{x_1^\transp A_1 x_2} \qth{x_3^\transp A_2 x_1} \sth{x_1^\transp A_3 x_3} x_2 \\
                    =& x_1 x_1^\transp A_1 x_2 x_3^\transp A_2 x_1 x_1^\transp A_3 x_3 x_2  = x_1 x_1^\transp \tilde{A}_1 x_1 x_1^\transp \tilde{A}_2 \tilde{x} x_2,
                \end{aligned}$$ where $\tilde{A}_1 = A_1 x_2 x_3^\transp A_2, \tilde{A}_2 = A_3$ and $\tilde{x} = x_3$. Besides, $m = 4, \beta_1 = 4$, thus the degree of $x_u x_v^\transp$ in all $\tilde{A}_k$ sum up to $m - \frac{\beta_1}{2} - 1 = 1$
