Weak and Strong Gradient Directions: Explaining Memorization, Generalization, and Hardness of Examples at Scale

1 INTRODUCTION . Generalization in over-parameterized neural networks trained using Stochastic Gradient Descent ( SGD ) is not well understood . Such networks typically have sufficient capacity to memorize their training set ( Zhang et al. , 2017 ) which naturally leads to the question : Among all the maps that are consistent with the training set , why does SGD learn one that generalizes well to the test set ? This question has spawned a lot of research in the past few years ( Arora et al. , 2018 ; Arpit et al. , 2017 ; Bartlett et al. , 2017 ; Belkin et al. , 2019 ; Fort et al. , 2020 ; Kawaguchi et al. , 2017 ; Neyshabur et al. , 2018 ; Sankararaman et al. , 2019 ; Rahaman et al. , 2019 ; Zhang et al. , 2017 ) . There have been many attempts to extend classical algorithm-independent techniques for reasoning about generalization ( e.g. , VC-dimension ) to incorporate the “ implicit bias ” of SGD to get tighter bounds ( by limiting the size of the hypothesis space to that reachable through SGD ) . Although this line of work is too large to review here , the recent paper of Nagarajan & Kolter ( 2019 ) provides a nice overview . However , they also point out some fundamental problems with this approach ( particularly , poor asymptotics ) , and come to the conclusion that the underlying proof technique itself ( uniform convergence ) may be inadequate . They argue instead for looking at algorithmic stability ( Bousquet & Elisseeff , 2002 ) . While there has been work on analysing the algorithmic stability of SGD ( Hardt et al. , 2016 ; Kuzborskij & Lampert , 2018 ) , it does not take into account the training data . Since SGD can memorize training data with random labels , and yet generalize on real data ( i.e. , its generalization behavior is data-dependent ( Arpit et al. , 2017 ) ) , any such analysis must lead to vacuous bounds in practical settings ( Zhang et al. , 2017 ) . Thus , in order for an algorithmic stability based argument to work , what is needed is an approach that takes into account both the algorithmic details of SGD as well as the training data . Recently , a new approach , for understanding generalization along these lines has been proposed in Chatterjee ( 2020 ) . Called the Coherent Gradients Hypothesis ( CGH ) , the key observation is that descent directions that are common to multiple examples ( i.e. , similar ) add up in the overall gradient ( i.e. , reinforce each other ) whereas directions that are idiosyncratic to particular examples fail to add up . Thus , the biggest changes to the network parameters are those that benefit multiple examples . In other words , certain directions in the tangent space of the loss function are “ strong ” gradient directions supported by multiple examples whereas other directions are “ weak ” directions supported by only a few examples . Intuitively–and CGH is only a qualitative theory at this point–strong directions are ( algorithmically ) stable ( in the sense of Bousquet & Elisseeff ( 2002 ) , i.e. , altered marginally by the removal of a single example ) whereas weak directions are ( algorithmically ) unstable ( could disappear entirely if the example supporting it is removed ) . Therefore , a change to the parameters along a strong direction should generalize better than one along a weak direction . Since the overall gradient is the mean of per-example gradients , if strong directions exist , the overall gradient has large components along it , and thus the parameter updates are biased towards algorithmic stability . Since CGH is a causal explanation for generalization , Chatterjee ( 2020 ) tested the theory by performing two causal interventions . Although they found good agreement between the qualitative predictions of the theory and experiments , an important limitation of their work is that their experiments were on shallow ( 1-3 hidden layers ) fully connected networks trained on MNIST using SGD with a fixed learning rate . In this work , we test CGH on large convolutional networks such as ResNet , Inception and VGG on ImageNet . While one of the tests of Chatterjee ( 2020 ) ( reducing similarity ) scales to this setting , the more compelling test ( suppressing weak gradients by winsorization ) does not . We propose a new class of scalable techniques for suppressing weak gradients , and also propose an entirely new test of CGH which is not based on causal intervention but on analyzing why some examples are learned earlier in training than others . 2 PRELIMINARY : REDUCING SIMILARITY ON IMAGENET . One test of CGH proposed in Chatterjee ( 2020 ) is to study how dataset similarity impacts training . Since directly studying similarity is difficult because which examples are considered similar may change during training ( in CGH examples are similar if their gradients are similar ) , Chatterjee ( 2020 ) proposed adding label noise to a dataset based on the intuition is that no matter what the notion of similarity , adding label noise is likely to decrease it . Therefore , if CGH is true , we should expect that : • As the label noise increases , the rate at which examples are learned decreases , • Examples whose labels have not been corrupted ( pristine examples ) should be learned faster than the rest ( corrupt examples ) , and , • With increasing noise , since there are fewer pristine examples , the rate at which they are learned should decrease . As preliminary experiment , we ran this test on ImageNet and the results for ResNet-18 are shown in Figure 1 . The results for Inception-V3 and VGG-13 are very similar ( please see Appendix A ) . We note the good agreement with the predictions from CGH thus providing initial evidence that CGH holds at scale . 3 ABLATING SGD TO TEST THE COHERENT GRADIENT HYPOTHESIS : SCALABLE TECHNIQUES TO SUPPRESS WEAK GRADIENT DIRECTIONS . Since weak directions are supported by few examples , CGH holds that overfitting and memorization in SGD is caused by descending down weak directions . The original CGH paper proposed to test CGH by modifying SGD to suppressing weak directions in order to verify that it significantly reduces overfitting ( Chatterjee , 2020 ) , i.e. , improves generalization through greater algorithmic stability ( Bousquet & Elisseeff , 2002 ) . 3.1 REVIEW OF THE WINSORIZATION TECHNIQUE . The test modified SGD by using the ( coordinate-wise ) winsorized mean of the per-example gradients in a mini-batch ( instead of the usual mean ) to update the weights in each step . Everything else , including learning rate was kept the same . The winsorized mean limits the influence of outliers by clipping them to a specified percentile of the data ( called the winsorization level ) . As expected from CGH , they found that as the winsorization level increased , the rate of overfitting decreased . We replicated this study with a ResNet-32 on CIFAR-10 and confirmed the results of the original study on this new dataset and architecture ( please see Appendix B ) . We note that since the original study was only on MNIST which has low generalization gap , the effect of suppressing weak directions only manifested with label noise . On CIFAR-10 even the real label case ( i.e. , 0 % noise ) has significant overfitting which is reduced by winsorization . This provides stronger evidence in support of CGH . However , a big challenge with winsorization is the need to compute and store per-example gradients which makes training much slower . For example , in our CIFAR-10 experiment we had to reduce the mini-batch size to 32 to make training feasible . This is exacerbated with larger models and more complex datasets such as ImageNet , and new techniques are needed to scale up the test . 3.2 TECHNIQUES BASED ON MEDIAN OF MEANS . We start by observing that the problem of suppressing weak gradient directions can be posed as a robust mean estimation problem from the robust statistics literature ( Huber , 1981 ) . Although winsorization is one way of obtaining robust mean estimates , there are others . In particular , the median of means algorithm ( Minsker , 2013 ) is an optimal estimation technique in the sense that deviation from the true mean is bounded above by O ( 1/ √ m ) with high probability ( m is the number of samples ) . The sample mean satisfies this property only if the observations are Gaussian . The main idea of the median of means algorithm is to divide the samples into k groups , computing the sample mean of each group , and then returning the geometric median of these k means . The geometric median of k vectors x1 , x2 , . . . xk ∈ Rd is the vector y∗ such that y∗ = argminy∈Rd ∑k i=1 ‖ y − xi ‖2 . When d = 1 , the geometric median is just the ordinary median of scalars . However , in high dimensions the algorithm to compute the geometric median ( Weiszfeld , 1937 ) is iterative and is expensive to integrate in a traditional training loop . A simpler technique is to apply the median of means algorithm to each coordinate that gives a dimension dependent bound on the performance of the estimator . The M3 Technique . The most obvious way to apply this idea to SGD is to divide a mini-batch into k groups of equal size . We compute the mean gradients of each group as usual , and then take their coordinate-wise median . The median is then used to update the weights of the network.1 Even though the algorithm is straightforward , its most efficient implementation ( i.e. , where the k groups are large and processed in parallel ) on modern hardware accelerators requires low-level changes to the stack to allow for a median-based aggregation instead of the mean . Therefore , in this work , we simply compute the mean gradient of each group as a separate micro-batch and only update the network weights with the median every k micro-batches , i.e. , we process the groups serially . 1Since we are simply replacing the mini-batch gradient with a more robust alternative , this technique may be used to study optimizers other than vanilla SGD ( such as SGD with momentum , ADAM , etc. ) . A systematic exploration of that is outside the scope of this study . In the serial implementation , k = 3 is a sweet spot . We have to remember only 2 previous microbatches , and since median ( x1 , x2 , x3 ) = ∑ i xi −mini xi −maxi xi ( where i ∈ { 1 , 2 , 3 } ) , we can compute the median with simple operations . We call this median-of-3 micro-batches ( M3 ) . Example ( Effectiveness of M3 ) . Fix d and consider a set of m < d training examples . At some point in training , let gi ∈ Rd ( 1 ≤ i ≤ m ) be their gradients . Suppose further that each gradient gi has an idiosyncratic component ui and a common component c , i.e. , gi = ui + c with ui · uj = 0 ( for j 6= i ) and ui · c = 0 . Since M3 involves coordinate-wise operations , let us assume that we are working in a basis where the non-zero coordinates of ui do not overlap with each other or with c. Now , consider a mini-batch of size 3b ≤ m constructed by picking 3b examples ( i.e. , their gradients ) uniformly at random without replacement . The expected value of this update if we take an SGD step ( i.e. , simply take the mean of the mini-batch ) , is gSGD = c+ 1m ∑ i ui . On the other hand , the expected value of the update with M3 is gM3 = c since any non-zero coordinate of any ui can not be the median value for that coordinate across the 3 groups since it can appear at most once . In this extreme case , we see that M3 suppresses the weak gradient directions ( the ui ) while preserving the strong gradient direction c. The RM3 Technique . Now , rather than update the weights every k micro-batches , as we do in M3 , we can update the weights every micro-batch using the median from that micro-batch and the previous two . In this rolling setup , there is no longer a difference between mini-batches and micro-batches , i.e. , this is essentially the same as computing the median over the current mini-batch gradient and the mini-batch gradients from the previous 2 steps . We call this rolling median of 3 mini-batches ( RM3 ) . Two remarks concerning RM3 are in order . First , RM3 may be seen as an approximation to M3 chosen for implementation efficiency . The assumption is that the loss function does not change very much in 1 gradient step , i.e. , it is locally stationary . Second , since RM3 uses recent gradient history , one may be tempted to think of RM3 as a form of momentum , but that would be wrong . We can understand this better by replacing the median operation in RM3 with the mean . We call this RA3 for rolling average of 3 mini-batches . As we shall see , it is the median v/s mean that makes a significant difference in suppressing weak gradient directions , not the rolling v/s non-rolling . Schematically , in their ability to suppress weak directions , we find that SGD ≈ RA3 RM3 < M3 . ( 1 )

Summary: of paper:**  Builds on and tests the "coherent gradients hypothesis" (CGH), which proposes that SGD is able to generalize well because of 'coherence' (similar direction) of gradients. The methodology of CGH involved comparing gradient directions on individual examples, infeasible for large datasets, so this work proposes and compares 2 mean- and median- based methods for testing CGH at scale. It also investigates easy/hard examples and thereby proposes a test of coherence of gradients based on the number of times an example of a certain difficulty is correctly classified.