Using stochastic gradient descent (SGD) with large batch-sizes to train deep neural networks is an increasingly popular technique. By doing so, one can improve parallelization by scaling to multiple workers (GPUs) and hence leading to significant reductions in training time. Unfortunately, a major drawback is the socalled "generalization gap": large-batch training typically leads to a degradation in generalization performance of the model as compared to small-batch training. In this paper, we propose to correct this generalization gap by adding diagonal Fisher curvature noise to large-batch gradient updates. We provide a theoretical analysis of our method in the convex quadratic setting. Our empirical study with state-of-the-art deep learning models shows that our method not only improves the generalization performance in large-batch training, but furthermore, does so in a way where the training convergence remains desirable and the training duration is not elongated. We additionally connect our method to recent works on loss surface landscape in the experimental section.
It has previously been observed that training deep networks using large batch-sizes leads to a larger generalization gap compared to the gap when training with a relatively small batch-size. This paper proposes to add noise sampled from diagonal "empirical" Fisher matrix to the large batch gradient as a method for closing the generalization gap. The authors motivate the use of empirical Fisher for sampling noise by arguing that the covariance of gradients from small batch-sizes can be seen as approximately equal to a scaled version of the Fisher matrix. It is then pointed out that using the Fisher matrix directly to sample noise could in principle close the generalization gap but would lead to slow converegence similar to SGD with a small batch-size. The authors then claim that the convergence speed is better when noise is sampled from the diagonal Fisher matrix instead of the full Fisher matrix. This claim is proven in theory for a convex quadratic loss surface and experiments are conducted to empirically verify this claim both in the quadratic setting are for realistic deep networks. Finally an efficient method for sampling noise from the diagonal empirical Fisher matrix is proposed.
Comments:
I think the paper is very well written and the results are presented clearly. In terms of novelty, I found the argument about convergence using diagonal Fisher being faster compared with full Fisher quite interesting, and its application for large batch training to be insightful. 
As a minor comment, for motivating theorem 3.1, it is pointed out by the authors that the diagonal Fisher acts as an approximation of the full Fisher and hence their regularization effects should be similar while convergence should be faster for diagonal Fisher. As a caveat, I think the authors should also point out that the convergence rate would be best when C is set to 0 in the result of the theorem. This implies no noise is used during SGD updates. However, this would imply the regularization effect from the noise will also vanish which would lead to poor generalization. 
However, there is a crucial detail that makes the main argument of the paper weak. In the main experiments in section 4.3, for the proposed large batch training method, the authors mention that they use a small batch-size of 128 for the first 50 epochs similar to Smith et al (2017) and then switch to the large batch-size of 4096, at which point, the learning rate is linearly scaled proportional to the large batch-size with a warmup scheme similar to Goyal et al (2017) and ghost batch normalization is used similar to Hoffer et al (2017). The former two tricks have individually been shown on their own to close the generalization gap for large batch-size training on large datasets like ImageNet. This paper combines these tricks and adds noise sampled from the diagonal Fisher matrix on top when switching to large batch-size after epoch 50 and reports experiments on smaller datasets-- MNIST, Fashion MNIST and the CIFAR datasets. Finally, the accuracy numbers for the proposed method is only marginally better than the baseline where isotropic noise is added to the large batch-size gradient. For these reasons, I do not consider the proposed method a significant improvement over existing techniques for closing the generalization gap for large batch training.
There is also a statement in the paper that is problematic but can be fixed by re-writing. In the paper, empirical Fisher matrix, as termed by the authors in the paper, refers to the Fisher matrix where the target values in the dataset is used as the output of the model rather than sampling it from the model itself as done for computing the true Fisher matrix. This empirical (diagonal) Fisher matrix is used to sample noise which is added to the large batch gradient in the proposed method. It is mentioned that the covariance of the noise in small batch SGD is exactly same as the empirical Fisher matrix. This claim is premised on the argument that the expected gradient (over dataset) is unconditionally roughly 0, i.e., throughout the training. This is absolutely false. If this was the case, gradient descent (using full dataset) should not be able to find minima and this is far from the truth. Even if we compare the scale of expected gradient to the mini-batch gradient (for small batch-size), the scale of these two gradients at any point during training (using say small batch-size SGD) is of the same order. I am saying the latter statement from my personal experience. The authors can verify this as well.
Overall, while I found the theoretical argument of the paper to be mostly interesting, I was dissapointed by the experimental details as they make the gains from the proposed method questionable when considered in isolation from the existing methods that close the generalization gap for large batch training.