Notes on "SGD Implicitly Regularizes Generalization Error"

These are my notes on the following paper: https://arxiv.org/abs/2104.04874 which analyses the inductive bias of SGD relative to full-batch GD using Taylor expansions. In order to make this paper easier to understand, I present the key ideas using modern machine learning notation.

The expected effect of a GD step on the generalisation gap

The generalisation gap, for the purposes of this paper, is the difference between a model's test loss and training loss. In this work we look at how this generalisation gap changes in a single step of gradient descent on the training data, starting from a fixed parameter

$\theta$. Since the training and test data are random, the change in change in generalisation gap is also random, we will be interested in the average of this change, when we average over realisations of the training and test data.

Let

${\mathcal{D}}_{train}$ and

${\mathcal{D}}_{test}$ be a random training and test set, drawn i.i.d. from the same underlying distribution over some data

$x$.

let

$\mathcal{L}(\theta ,\mathcal{D})=\frac{1}{|\mathcal{D}|}\sum _{x\in \mathcal{D}}\ell (x,\theta )$ be the empirical loss evaluated on a set of data

$\mathcal{D}$.

Let

$g(\theta ,\mathcal{D})={\mathrm{\nabla }}_{\theta }\mathcal{L}(\theta ,\mathcal{D})$ be the gradient evaluated on a set of data

$\mathcal{D}$.

Let's say we take a gradient step on

${\mathcal{D}}_{train}$. We can approximate how the training error changes via Taylor expansion of

$\mathcal{L}(\theta ,{\mathcal{D}}_{train})$:

$$\begin{array}{rl}\mathrm{\Delta }{\mathcal{L}}_{train}& =\mathcal{L}({\theta }_0-\eta g({\theta }_0,{\mathcal{D}}_{train}),{\mathcal{D}}_{train})-\mathcal{L}({\theta }_0,{\mathcal{D}}_{train})\\ & =\eta g({\theta }_0,{\mathcal{D}}_{train}{)}^{T}g({\theta }_0,{\mathcal{D}}_{train})+o({\eta }^2)\end{array}$$

Similarly, we can express the change in the test loss as:

$$\begin{array}{rl}\mathrm{\Delta }{\mathcal{L}}_{test}& =\mathcal{L}({\theta }_0-\eta g({\theta }_0,{\mathcal{D}}_{train}),{\mathcal{D}}_{test})-\mathcal{L}({\theta }_0,{\mathcal{D}}_{test})\\ & =\eta g({\theta }_0,{\mathcal{D}}_{test}{)}^{T}g({\theta }_0,{\mathcal{D}}_{train})+o({\eta }^2)\end{array}$$

To move forward, we will want to average the above quantities over random realization of the training and test dataset. Let's consider a dataset of size

$N$,

$\mathcal{D}=\{x_1,\dots ,x_N\}$, where each datapoint

$x_i$ is drawn independently from the same distribution

$P$. Let's say we have two index sets

${\mathcal{I}}_1,{\mathcal{I}}_2\subseteq \{1\dots N\}$, and corresponding subsets

${\mathcal{D}}_1=\{x_i:i\in {\mathcal{I}}_1\}\subseteq \mathcal{D}$ and

${\mathcal{D}}_2=\{x_i:i\in {\mathcal{I}}_2\}\subseteq \mathcal{D}$. We now consider the following expectation:

$$\begin{array}{rl}|{\mathcal{I}}_1||{\mathcal{I}}_2|{\mathbb{E}}_{\mathcal{D}}& [g({\theta }_0,{\mathcal{D}}_1{)}^{T}g({\theta }_0,{\mathcal{D}}_2)]=\sum _{n\in {\mathcal{I}}_1}\sum _{m\in {\mathcal{I}}_2}{\mathbb{E}}_{\mathcal{D}}[g({\theta }_0,x_n{)}^{T}g({\theta }_0,x_m)]\\ & =\sum _{n\in {\mathcal{I}}_1\cap {\mathcal{I}}_2}{\mathbb{E}}_{\mathcal{D}}[g({\theta }_0,x_n{)}^{T}g({\theta }_0,x_n)]+\sum _{n\in {\mathcal{I}}_1}\sum _{n\in {\mathcal{I}}_2\setminus {\mathcal{I}}_1}{\mathbb{E}}_{\mathcal{D}}[g({\theta }_0,x_n{)}^{T}g({\theta }_0,x_m)]\\ & =|{\mathcal{I}}_1\cap {\mathcal{I}}_2|{\mathbb{E}}_{X\sim P}\Vert g({\theta }_0,X){\Vert }^2+(|{\mathcal{I}}_1||{\mathcal{I}}_2|-|{\mathcal{I}}_1\cap {\mathcal{I}}_2|){\Vert {\mathbb{E}}_{X\sim P}g({\theta }_0,X)\Vert }^2\\ & =|{\mathcal{I}}_1\cap {\mathcal{I}}_2|\mathrm{tr}({\mathbb{Cov}}_{X\sim P}g({\theta }_0,X))+|{\mathcal{I}}_1||{\mathcal{I}}_2|{\Vert {\mathbb{E}}_{X\sim P}g({\theta }_0,X)\Vert }^2,\end{array}$$

Where we used the fact that the expectation of the product of independent random variables is the product of their expetations on the second term in the second line, and the following identity that for

$\mathrm{tr}\mathbb{Cov}[G]=\mathbb{E}[\Vert G{\Vert }^2]-\Vert \mathbb{E}[G]{\Vert }^2$ for vector valued

$X$ which we prove below using the cyclic property of trace:

$$\begin{array}{rl}\mathrm{tr}\mathbb{Cov}[G]& =c\mathbb{E}[G{G}^T]-\mathrm{tr}\mathbb{E}[G]\mathbb{E}[G{]}^T\\ & =\mathbb{E}\mathrm{tr}{G}^{T}G-\mathrm{tr}\mathbb{E}[G{]}^T\mathbb{E}[G]\\ & =\mathbb{E}\Vert G{\Vert }^2-\Vert \mathbb{E}[G]{\Vert }^2\end{array}$$

Using this, we can now express the expected change in training and test errors, under the assumption that

${\mathcal{I}}_{test}\cap {\mathcal{I}}_{t}rain=\mathrm{\varnothing }$:

$$\begin{array}{rl}\mathbb{E}\mathrm{\Delta }{\mathcal{L}}_{test}& =-\eta {\Vert {\mathbb{E}}_{X\sim P}g({\theta }_0,X)\Vert }^2\\ \mathbb{E}\mathrm{\Delta }{\mathcal{L}}_{train}& =-\eta {\Vert {\mathbb{E}}_{X\sim P}g({\theta }_0,X)\Vert }^2-\frac{\eta }{N_{train}}\mathrm{tr}({\mathbb{Cov}}_{X\sim P}g({\theta }_0,X))\end{array}$$

Thus, the training error is reduced more than the test error, in expectation. This is what the authors mean when they write that 'GD develops a “bias” toward overfitting', because in expectation, one gradient step increases the gap between training and test losses.

This is the first main finding of the paper, equation (1), just written differently:

$$\mathbb{E}[\mathrm{\Delta }{\mathcal{L}}_{test}-\mathrm{\Delta }{\mathcal{L}}_{train}]=\frac{\eta }{N_{train}}\mathrm{tr}\mathrm{\Sigma },$$

where

$\mathrm{\Sigma }$ is the covariance matrix of gradients on a single datapoint.

This also provides a nice, intuitive explanation for how the local total variance of gradients

$\mathrm{tr}\mathrm{\Sigma }$ is related to generalisation. Ideally, if the variance minibatch gradients is negligible - or small, it means that running gradient descent is a 'safe way' to learn, because in this case the gradient one would get on the training set is the same as one would get on the test set.

A key limitation of this analysis is that it assummes

${\theta }_0$ is chosen independently of

${\mathcal{D}}_{train}$. So long as we never recycle data in subsequent training steps this would be true in each step of gradient descent. In practice, however, we cycle through training data multiple times. If

${\theta }_0$ in our equations is some intermediate state, and is itself found by running gradient descent on data that overlapped with the training data we currently evaluate, the above analysis is flawed (because it considers

${\theta }_0$ fixed and not a function of

${\mathcal{D}}_{train}$). In this situation it can no longer be guaranteed that the distribution of gradients is going to be the same on the training and test set and the whole argument fails.

Characterising the difference between SGD and GD

The second finding of this paper contrasts stochastic and full-batch gradient descent.