Disentangling the Mechanisms Behind
Implicit Regularization in SGD

As proposed by Smith et al. [2021], we attempt to understand the generalization behavior of mini-batch SGD by how it implicitly regularizes the norm of the micro-batch gradient, $\|\nabla\mathcal{L}_{M}(\boldsymbol{\theta})\|$ for some micro-batch $M\subseteq\mathcal{B}$. In large-batch SGD, we accomplish this through gradient accumulation (i.e. accumulating the gradients of many small-batches to generate the large-batch update), and thus can add an explicit regularizer (described in Geiping et al. [2022]) that penalizes the average micro-batch norm. Formally, for some large-batch $\mathcal{B}$ and component disjoint micro-batches $M\subseteq\mathcal{B}$, let $\nabla_{\boldsymbol{\theta}}\mathcal{L}_{M}(\boldsymbol{\theta})=\frac{1}{|M|}\sum_{(\mathbf{x},y)\in M}\nabla_{\boldsymbol{\theta}}\ell(\mathbf{x},y;\boldsymbol{\theta})$. The new loss function is:

$$\mathcal{L}_{\mathcal{B}}(\boldsymbol{\theta})+\lambda\frac{|M|}{|\mathcal{B}|}\sum_{M\in\mathcal{B}}\|\nabla_{\boldsymbol{\theta}}\mathcal{L}_{M}(\boldsymbol{\theta})\|^{2}.$$

(1)

While this quantity can be approximated through finite differences, it can also be optimized directly by backpropagation using modern deep learning packages, as we do in this paper.

Note that by definition, we can decompose the regularizer term into the product of the Jacobian of the network and the gradient of the loss with respect to network output. Formally, for some network $f$ with $p$ parameters, if we let $\mathbf{z}=f(\mathbf{x};\boldsymbol{\theta})\in\mathbb{R}^{C}$ be the model output for some input $\mathbf{x}$ and denote its corresponding label as $y$, then:

$${\color[rgb]{0.0,0.0,0.0}\definecolor[named]{pgfstrokecolor}{rgb}{0.0,0.0,0.0}\nabla_{\boldsymbol{\theta}}\ell(\mathbf{x},y;\boldsymbol{\theta})=(\nabla_{\boldsymbol{\theta}}\mathbf{z})(\nabla_{\mathbf{z}}\ell(\mathbf{x},y;\boldsymbol{\theta}))}$$

(2)

Where $\nabla_{\boldsymbol{\theta}}\mathbf{z}\in\mathbb{R}^{p\times C}$ is the Jacobian of the network and the second term is the loss-output gradient. We explicitly show this equivalence for the comparison made in section 2.3.

One noticeable artifact of Equation 1 is its implicit reliance on the true labels $y$ to calculate the regularizer penalty. Jastrzebski et al. [2020] shows that we can derive a similar quantity in the mini-batch SGD setting by penalizing the trace of the Fisher Information Matrix $\mathbf{F}$, which is given by $\text{Tr}(\mathbf{F})=\mathbb{E}_{\mathbf{x}\sim\mathcal{X},\hat{y}\sim p_{\boldsymbol{\theta}}(y\mid\mathbf{x})}[\|\nabla_{\boldsymbol{\theta}}\ell(\mathbf{x},\hat{y};\boldsymbol{\theta})\|^{2}]$, where $p_{\boldsymbol{\theta}}(y\mid\mathbf{x})$ is the predictive distribution of the model at the current iteration and $\mathcal{X}$ is the data distribution. We thus extend their work to the accumulated large-batch regime and penalize an approximation of the average Fisher trace over micro-batches: if we let $\widehat{\mathcal{L}}_{M}(\boldsymbol{\theta})=\frac{1}{|M|}\sum_{\mathbf{x}\in M,\hat{y}\sim p_{\boldsymbol{\theta}}(y\mid\mathbf{x})}\ell(\mathbf{x},\hat{y};\boldsymbol{\theta})$, then our penalized loss is

$$\mathcal{L}_{\mathcal{B}}(\boldsymbol{\theta})+\lambda\frac{|M|}{|\mathcal{B}|}\sum_{M\in\mathcal{B}}\|\nabla_{\boldsymbol{\theta}}\widehat{\mathcal{L}}_{M}(\boldsymbol{\theta})\|^{2}.$$

(3)

The only difference between the expressions in Equation 1 and Equation 3 is that the latter now uses labels sampled from the predictive distribution, rather than the true labels, to calculate the regularizer term. As with Equation 1, we can directly backpropagate using this term in our loss equation.

Like, in equation 2, we can decompose the regularizer term as:

$${\color[rgb]{0.0,0.0,0.0}\definecolor[named]{pgfstrokecolor}{rgb}{0.0,0.0,0.0}\ell(\mathbf{x},\hat{y};\boldsymbol{\theta})=(\nabla_{\boldsymbol{\theta}}\mathbf{z})(\nabla_{\mathbf{z}}\ell(\mathbf{x},\hat{y};\boldsymbol{\theta}))}$$

(4)

Where the second term is another loss-output gradient.

Jastrzebski et al. [2020] observes that models with poor generalization typically show a large spike in the Fisher Trace during the early phases of training, which they coin as Catastrophic Fisher Explosion. In Figure 1, we show that this behavior also occurs when looking at the average Micro-Batch gradient norm.

Given the decompositions shown in equations 2 and 4, it is unclear in either case whether the Jacobian term is the sole source of possible generalization benefit, or if the loss-output gradient is also needed. To disentangle this effect, we borrow from Lee et al. [2022] and use the average and unit Jacobian regularization losses:

$\displaystyle\mathcal{L}_{\mathcal{B}}(\boldsymbol{\theta})+\lambda\frac{|M|}{|\mathcal{B}|}\sum_{M\subseteq\mathcal{B}}\|J_{\text{avg}}(M)\|^{2},\quad$

$\displaystyle\quad J_{\text{avg}}(M)=\frac{1}{|M|}\sum_{(\mathbf{x},y)\in M}(\nabla_{\boldsymbol{\theta}}\mathbf{z})(\frac{1}{C}\mathbbm{1}),$

(5)

$\displaystyle\mathcal{L}_{\mathcal{B}}(\boldsymbol{\theta})+\lambda\frac{|M|}{|\mathcal{B}|}\sum_{M\subseteq\mathcal{B}}\|J_{\text{unit}}(M)\|^{2},\quad$

$\displaystyle\quad J_{\text{unit}}(M)=\frac{1}{|M|}\sum_{(\mathbf{x},y)\in M}(\nabla_{\boldsymbol{\theta}}\mathbf{z})(\mathbf{u}),$

(6)

where $\mathbf{u}\in\mathbb{R}^{C}$ is randomly sampled from the unit hypersphere (i.e. $\|\mathbf{u}\|_{2}=1$), and $\mathbf{u}$ is sampled once per micro-batch. In words, the average Jacobian case penalizes the Jacobian with equal weighting on every class and every example, while the unit Jacobian case penalizes the Jacobian with different but random weighting on each class and example. Note that the unit Jacobian penalty is an unbiased estimator of the Frobenius norm of the Jacobian $\|\nabla_{\boldsymbol{\theta}}\mathbf{z}\|_{F}^{2}$, which is an upper bound on its spectral norm $\|\nabla_{\boldsymbol{\theta}}\mathbf{z}\|_{2}^{2}$ (see Lee et al. [2022] for a more detailed theoretical analysis).

We first focus our experiments on the case of using a ResNet-18 [He et al., 2015], with standard initialization and batch normalization, on the CIFAR10, CIFAR100, Tiny-ImageNet, and SVHN image classification benchmarks [Krizhevsky, 2009, Le and Yang, 2015, Netzer et al., 2011]. Additional experiments on different architectures are detailed in Appendix A.1. Besides our small-batch ($\mathcal{B}=128$) and large-batch ($\mathcal{B}=5120$) SGD baselines, we also train the networks in the large-batch regime using (i) average Micro-batch Gradient Norm Regularization (GN); (ii) average Micro-batch Fisher Trace Regularization (FT); and (iii) average and unit Micro-batch Jacobian Regularizations (AJ and UJ). Note that for all the regularized experiments, we use a component micro-batch size equal to the small-batch size (i.e. 128). In order to compare the best possible performance within each experimental regime, we separately tune the optimal learning rate $\eta$ and optimal regularization parameter $\lambda$ independently for each regime. Additional experimental details can be found in Appendix A.5.

(i) Average micro-batch Gradient Norm and average Fisher trace Regularization recover SGD generalization   For CIFAR100, Tiny-ImageNet, and SVHN we find that we can fully recover small-batch SGD generalization performance by penalizing the average micro-batch Fisher trace and nearly recover performance by penalizing the average micro-batch gradient norm (with an optimally tuned regularization parameter $\lambda$, see Figure 1 and Table 1). In CIFAR10, neither penalizing the gradient norm nor the Fisher trace completely recovers small-batch SGD performance, but rather come within $\approx 0.3\%$ and $\approx 0.4\%$ (respectively) the small-batch SGD performance and significantly improves over large-batch SGD.

We additionally find that using the micro-batch gradient norm leads to slightly faster per-iteration convergence but less stable training (as noted by the tendency for the model to exhibit random drops in performance), while using the Fisher trace leads to slightly slower per-iteration convergence but much more stable training (see Figure 1). This behavior may be due to the Fisher trace’s ability to more reliably mimic the small-batch SGD micro-batch gradient norm behavior throughout training, whereas penalizing the gradient norm effectively curbs the initial explosion but collapses to much smaller norm values as training progresses.

(ii) Average and Unit Jacobian regularizations do not recover SGD generalization   Observe in Table 1 that we are unable to match SGD generalization performance with either Jacobian regularization. In Section 2 we showed that each regularization method can be viewed as penalizing the norm of the Jacobian matrix-vector product with some $C$-dimensional vector. Crucially, both the gradient norm and Fisher trace regularizers use some form of loss-output gradient, which is data-dependent and has no constraint on the weighting of each class, while both Jacobian regularizers use data-independent and comparatively simpler vectors. Given the noticeable difference in generalization performance between the regularization methods that weight the Jacobian with the loss-output gradient and those that do not, we indicate that the loss-output gradient may be crucial to either applying the beneficial regularization effect itself and/or stabilizing the training procedure.

In both successful regularization regimes, namely the average micro-batch gradient norm and average fisher trace regularizers, there is an implicit hyperparameter: the size of the micro-batch used to calculate the regularization term. Note that this hyperparameter is a practical artifact of modern GPU memory limits, as efficiently calculating higher-order derivatives for large batch sizes is not feasible in standard autodifferentiation packages. Consequently, gradient accumulation (and the use of the average micro-batch regularizer, rather than taking the norm over the entire batch) must be used on most standard GPUs (more detailed hadware specifications can be found in appendix A.5).

This restriction, however, may actually be beneficial, as Geiping et al. [2022], Barrett and Dherin [2020], Smith et al. [2021] have noted that they expect the benefits of gradient norm regularizers to break down when the micro-batch size becomes too large. To test this hypothesis, we return to the ResNet-18 in CIFAR100 and Tiny-ImageNet settings and increase the micro-batch size to as large as we could reasonably fit on a single GPU at $|M|=2560$ in both the gradient norm and Fisher trace experiments. Additionally, we run experiments using a VGG11 [Simonyan and Zisserman, 2015] on CIFAR10, interpolating the micro-batch size from the small to large regimes. In both settings, we separately tune the learning rate $\eta$ and regularization coefficient $\lambda$ in each experiment to find the best possible generalization performance in the large micro-batch regimes.

Results   We successfully show that such hypotheses mentioned in Geiping et al. [2022], Smith et al. [2021], Barrett and Dherin [2020] hold true: as the micro-batch size approaches the mini-batch size, both regularization mechanisms lose the ability to recover small-batch SGD performance (see Tables 2 and 3). Additionally, we note that using large micro-batch sizes no longer effectively mimics the average micro-batch gradient norm behavior of small-batch SGD, thus supporting our claim that matching this quantity throughout training is of key importance to recovering generalization performance (Figure 2).

One potential practical drawback of these gradient-based regularization terms is the relatively high computation cost needed to calculate the second-order gradients for every component micro-batch. Instead of penalizing the average micro-batch gradient norm, we can penalize one micro-batch gradient norm. For some large batch $\mathcal{B}$ and fixed sample micro-batch $S$ from batch $\mathcal{B}$, we define the modified loss function

$$\mathcal{L}_{\mathcal{B}}(\boldsymbol{\theta})+\lambda\|\nabla_{\boldsymbol{\theta}}\mathcal{L}_{S}(\boldsymbol{\theta})\|^{2}.$$

(7)

Results   In Figure 3, we plot the final test accuracy (left column) and the average gradient norm (right column) as a function of $\lambda$. We observe that both a larger $\lambda$ and a smaller micro-batch size $|S|$ boost test accuracy. Furthermore, we find that with the “optimal” $\lambda$ and micro-batch size $|S|$, the final test accuracy for sample micro-batch gradient norm regularization is close to (and sometimes better) than the final test accuracy for SGD. Just as we observed with the average Micro-batch Gradient Norm regularization, generalization benefits diminish as the sample micro-batch size approaches the mini-batch size.

Inspired by the work of Agarwal et al. [2020], we proposed to use gradient grafting in order to control the loss gradient norm behavior during training. Formally, for any two different optimization algorithms $\mathcal{M},\mathcal{D}$, the grafted updated rule is arbitrarily:

	$\displaystyle g_{\mathcal{M}}$	$\displaystyle=\mathcal{M}(\boldsymbol{\theta}_{k}),\quad\quad g_{\mathcal{D}}=\mathcal{D}(\boldsymbol{\theta}_{k})$
	$\displaystyle\boldsymbol{\theta}_{k+1}$	$\displaystyle=\boldsymbol{\theta}_{k}-\|g_{\mathcal{M}}\|\frac{g_{\mathcal{D}}}{\|g_{\mathcal{D}}\|}$		(8)

In this sense, $\mathcal{M}$ controls the magnitude of the update step and $\mathcal{D}$ controls the direction. We first propose Iterative Grafting, wherein $\mathcal{M}(\boldsymbol{\theta}_{k})=\eta\nabla\mathcal{L}_{M}(\boldsymbol{\theta}_{k})$ and $\mathcal{D}(\boldsymbol{\theta}_{k})=\nabla\mathcal{L}_{\mathcal{B}}(\boldsymbol{\theta}_{k})$, where $M\in\mathcal{B}$ is sampled uniformly from the component micro-batches at every update. In words, at every update step we take the large batch gradient, normalize it, and then rescale the gradient by the norm of one of the component micro-batch gradients.

Additionally, we propose External Grafting, where $\mathcal{M}(\boldsymbol{\theta}_{k})=\eta\nabla\mathcal{L}_{M}(\boldsymbol{\theta}_{k^{\prime}})$ and $\mathcal{D}(\boldsymbol{\theta}_{k})=\nabla\mathcal{L}_{\mathcal{B}}(\boldsymbol{\theta}_{k})$. Here, we use $\nabla\mathcal{L}_{B}(\boldsymbol{\theta}_{k^{\prime}})$ to denote the gradient norm at step $k$ from a separate small-batch SGD training run. We propose this experiment to make a comparison with the Iterative Grafting case, since here the implicit step length schedule is independent of the current run, while with Iterative grafting the schedule depends upon the current training dynamics.

Aside from grafting algorithms, which define the implicit step length schedule at every step, we also consider the situation where the step length is fixed throughout training through normalized gradient descent (NGD) [Hazan et al., 2015], wherein $\mathcal{M}(\boldsymbol{\theta}_{k})=\eta$ and $\mathcal{D}(\boldsymbol{\theta}_{k})=\eta\nabla\mathcal{L}_{\mathcal{B}}(\boldsymbol{\theta}_{k})$.

Results   We find that both forms of grafting and NGD can recover the generalization performance of SGD in some model-dataset combinations (see Table 4). Namely, though grafting / NGD seems to work quite well in CIFAR10, no amount of hyperparameter tuning was able to recover the SGD performance for either model in CIFAR100. That being said, in the CIFAR10 case we see (in Figure 4) that the grafting experiments (and NGD, not pictured) do not replicate the same average mini-batch gradient norm behavior of small-batch SGD despite sometimes replicating its performance. This thus gives us solid empirical evidence that while controlling the average mini-batch gradient norm behavior through explicit regularization may aide generalization, it is not the only mechanism in which large-batch SGD can recover performance.

The stark disparity in performance between the CIFAR10 and CIFAR100 benchmarks are of key importance. These differences may be explained by the much larger disparity between the mid-stage average micro-batch gradient norm behavior in the CIFAR100 case than in the CIFAR10 case (see Figure 4). This situation highlights a possible cultural issue within the deep learning community: there is a concerning trend of papers in the deep learning field that cite desired performance on CIFAR10, and no harder datasets, as empirical justification for any posed theoretical results [Orvieto et al., 2022, Geiping et al., 2022, Agarwal et al., 2020, Smith et al., 2021, Barrett and Dherin, 2020, Cheng et al., 2020]. Given the continued advancement of state-of-the-art deep learning models, we argue that it is imperative that baselines like CIFAR100 and ImageNet are adopted as the main standard for empirical verification, so that possibly non-generalizable results (as the grafting / NGD results would have been had we stopped at CIFAR10) do not fall through the cracks in the larger community (see Appendix A.3 for more information).