Title: AdaScale SGD: A Scale-Invariant Algorithm for Distributed Training
Authors: Tyler B. Johnson
Pulkit Agrawal
Haijie Gu
Carlos Guestrin
Venue: ICLR 2020
Keywords: Large-batch SGD
large-scale learning
distributed training

Abstract: When using distributed training to speed up stochastic gradient descent, learning rates must adapt to new scales in order to maintain training effectiveness. Re-tuning these parameters is resource intensive, while fixed scaling rules often degrade model quality. We propose AdaScale SGD, a practical and principled algorithm that is approximately scale invariant. By continually adapting to the gradient’s variance, AdaScale often trains at a wide range of scales with nearly identical results. We describe this invariance formally through AdaScale’s convergence bounds. As the batch size increases, the bounds maintain final objective values, while smoothly transitioning away from linear speed-ups. In empirical comparisons, AdaScale trains well beyond the batch size limits of popular “linear learning rate scaling” rules. This includes large-scale training without model degradation for machine translation, image classification, object detection, and speech recognition tasks. The algorithm introduces negligible computational overhead and no tuning parameters, making AdaScale an attractive choice for large-scale training.


TL;DR: A practical and principled algorithm for distributed SGD, which simplifies the process of scaling up training

Review discussion:
This paper proposes a novel rule for scaling the learning rate, called the gain ratio, for when the effective batch size (induced by synchronous distributed SGD) is increased. This rule for the gain ratio uses estimates of the trace of the covariance and the norm of the true gradient to determine the appropriate steplength. This results in a method with a R-linear convergence guarantee to a neighborhood that is not dependent on $S$ (which is called scale-invariant). The algorithm additionally tracks the sum of the gain ratios in order to determine the "effective" number of iterations taken, and cut the steplength appropriately.
Strengths:
The gain ratio proposed in this paper is intuitive. I particularly like how the algorithm estimates the mean and variance information in an unbiased manner to determine an appropriate scaling for the steplength. The method is able to attain a R-linear rate of convergence and appears to perform well in practice for a wide variety of applications. The gain ratio is simple to estimate within a distributed framework.
Weaknesses:
I found some of the terms in the paper to be unclear or ill-defined. The original use of the term "scale" was unclear to me. Does this refer to the number of nodes in the distributed implementation? What is its relationship to batch size?
I found the definition of scale invariance in this paper to also be unclear on first read. The claim is that the algorithm is scale invariant if its final model does not depend on $S$. What is the "final model"? As currently defined, the current analysis does not guarantee that the algorithm will reach the same final model (assuming that $f(w, x) = \ell(h(w, x), y)$, i.e. a composition of a loss function and model), as the PL condition only ensures that one reaches a global minimum, which may not be unique.
In fact, the analysis only guarantees convergence to a neighborhood. The description within the analysis appears to imply that scale-invariance is a property of the algorithm attached to its convergence property. Is this the case?
The definition of scale invariance is also already used in optimization to mean algorithms that are not modified when the objective is multiplied by a constant or an affine transformation. This adds to the lack of clarity, and I would suggest the authors use a different term for this kind of invariance (batch size invariant, or something like that?).
Is the theoretical comparison between SGD and AdaScale fair? Note that one can prove a stronger convergence result with SGD because one can actually attain a Q-linear rate of convergence to a neighborhood (for a proof, see for example, Bottou, et al. (2018)).
In particular, one should have something like (in the paper's notation):
$$\mathbb{E}[F(w_T) - F^*] \leq (1 - \gamma)^T [F(w_0) - F^*] + \Delta.$$
This means that one can actually guarantee a fixed ratio of decrease in expectation to a neighborhood, whereas AdaScale converges linearly but not with a fixed ratio.
Some other small questions regarding the theory and experiments:
- Is there a reason why batch normalization was not tried for the CIFAR-10 experiments?
- Is it possible for $r_{t - 1} \gamma > 1$?
- Why was it necessary to estimate $\sigma^2 (w_t)$ and $\|\nabla F(w_t)\|^2$ by both aggregating at the current iteration and exponential averaging? What happens if exponential averaging is removed?
- What are the limitations of this method? How large of a batch size can one use with AdaScale before the algorithm breaks down (if at all)?
Additional Comments:
The algorithm is quite reminiscient of the steplength prescribed in Bollapragada, et al. (2018), which consider the steplength:
$$( 1 + \frac{\sigma^2(w_t)}{\|\nabla F(w_t)\|^2})^{-1}.$$
This gain ratio prescribed in this paper is the ratio between this quantity for two different batch sizes. Is there a clear explanation for why the relationship between these two quantities would arise?
This method could also be used for determining an appropriate scaling of the steplength in the sequential setting, when a larger batch size is used. Has this been considered?
Despite the concerns regarding the clarity in writing and the rigor in the theory of the paper, I think that the algorithmic idea proposed in this paper is interesting, novel, and practical. Because of the lack of clarity and rigor, I have given this paper a weak reject, but I would be happy to accept the paper if my concerns above were addressed in the final manuscript.
Thank you for the detailed and constructive review.
We have updated all theorems to achieve the better rates.  Hopefully it is clear now that the theoretical comparisons are fair.
We tried to write for multiple audiences, so that both deep learning practitioners and optimization experts would find the paper useful.  In Section 2, we define the scale S as the number of batches that are processed in parallel during each iteration.
Thus, batch size is proportional to scale.  We have updated the definition of “scale invariance” so that hopefully it is more clear.  AdaScale only approximately achieves scale invariance.  The scale invariant convergence bound (i.e., the bound that does not depend on S) formalizes this approximation, but the bound does not imply an exactly scale invariant algorithm.
We have some experiments in Appendix C that show results with no exponential averaging.  There is more variance in the gain estimate, but AdaScale still performs well.
In extreme settings, it is possible that r \gamma > 1.  In such cases, the bound would be Equation (6) in Appendix A of the updated submission.  In this case, AdaScale can oscillate between small and large r.  Related to this, there is indeed a limit on the scale S that achieves scale-invariant bounds, and we include this in the updated theorems.
In practice, however, we do not think r \gamma > 1 is a significant concern.  System and algorithm scaling inefficiencies limit the practicality of such extreme S.  Users can also decrease \gamma by decreasing the learning rate.
Yes, one could apply AdaScale in the sequential setting, but we think AdaSale is most impactful for distributed training because of the large speed-ups.
We do use batch normalization for CIFAR-10.  We do not train the batch norm scaling parameters because this helped us achieve 94% accuracy with 1 GPU.
The paper presents and evaluate an algorithm, AdaScale SGD, to improve the stochastic gradient decent in distributed training. The proposed algorithm claims to be approximately scale-invariant, by adapting to the gradient's variance. The paper is well-written and generally easy to read, although I didn't check all theory in the paper.
The approach is evaluated using five benchmarks (networks / dataset combinations) from different  domains. The results are promising and seems solid.
The paper is good and I like it, although I think the novelty and contribution is slightly too low for ICLR. The kind of tweaking and minor optimizations to provide some adaptivity (or similar) in existing and established algorithms and approaches that is presented in this is paper is very important from a practical perspective.
However, from a scientific perspective it provides no significant contribution.
We disagree that the novelty and contribution are too low for ICLR.  The ability to scale up training is often critical to the development of state-of-the-art models.  These models tend to be large, and developing them requires fast turnaround times and a lot of data.  AdaScale makes it significantly easier to scale up training, since the need to re-tune learning rate parameters is one of the biggest pain points when scaling to large batch sizes.  AdaScale is an entirely novel algorithm, as no prior algorithm adapts to the gradient’s variance in order to achieve this scaling goal.
Furthermore, while we agree that AdaScale uses a tweak to provide adaptivity, we are surprised this is framed as a criticism.  Highly influential prior algorithms, such as Adam, AdaGrad, etc., could also be considered tweaks.  These tweaks are simple but non-trivial, and the simplicity has made these algorithms extremely successful.
Finally, we do not understand the remark that “this is paper is very important from a practical perspective . . . from a scientific perspective it provides no significant contribution.”  The comment lacks justification and unfairly criticizes our paper’s importance.  Our paper introduces a practical training algorithm, thoroughly demonstrates its usefulness, and supports the algorithm with theoretical understanding.  These contributions fall firmly within ICLR’s established standards, and we believe many ICLR attendees would take interest in the work and find it useful.
Authors use the PL condition to motivate a formula for learning rate selection. (they bound their loss by a quadratic with fixed curvature alpha, and use this to get learning rate).
Their analysis very closely mirrors one presented in "Gradient Diversity" paper, it uses the same assumptions on the loss. IE compare A.1 in the paper to the B.5 in "Gradient Diversity"
Gradient Diversity solves for batch size B after which linear learning rate scaling starts to break down, while this paper instead fixes B and solves for learning rate. Two results are comparable, if you take their learning rate formula in section 3.2 (need formula numbers) and solve for B which gives learning rate halfway between B=1 and B=inf, you get the same expression as "critical batch size" in equation 5 of gradient diversity paper.
It's not immediately obvious how to invert formula in Gradient Diversity paper to solve for learning rate, so I would consider their learning rate formula an interesting result.
I also appreciate the breadth of experiments used for evaluation.
The biggest issue I have with the paper is that I can't tell if it's better of worse than linear learning rate scaling from their experiment section. All of their experiments use more iterations for AS evaluation uses than for LSW evaluation. They demonstrate better training (and test) losses for AS, but because of extra iterations, I can't tell that the improvement in training loss is due to number of iterations, or due to AS scaling.
Thank you for the helpful review.
Regarding the comparison with LSW:
As suggested, we tested a version of LSW that uses more iterations.  We call this scaling strategy “LSW+”.  LSW+ scales the iterations axis of the LSW learning rate schedule so that LSW+ uses the same number of iterations as AdaScale.  The learning rate axis remains the same as LSW.
The behavior of LSW+ is generally similar to that of LSW.  As expected, LSW+ improves upon LSW, but LSW+ still degrades model quality for all benchmarks at the largest scale (CIFAR, Deep Speech, ImageNet, Transformer, and YOLO).  For Deep Speech, Transformer, and YOLO, LSW+ also diverges at the same scales that LSW diverges.  We have included these results in updated Appendix D (Table 5).
We also note that LSW+ is not a practical algorithm, because it requires either (i) first running AdaScale to determine the number of iterations; or (ii) tuning the number of iterations.  Both options are inconvenient in practice.
Moreover, for a fair comparison, it seems we would also need to consider AdaScale with tuning.
Thus, even if LSW+ had matched AdaScale, AdaScale would still be preferable to LSW+.
Regarding the Gradient Diversity Paper:
We agree our analysis shares similarities with Yin et al. (2018)’s analysis — several papers that we cite in Section 5 contain similar analysis.
But the contributions of the papers are different.  Yin et al. show that if we optimize the learning rate for convergence bounds, then large variance implies near-linear speed-ups from data parallelism, while small variance implies small speed-ups.
What our paper shows is that regardless of how you set the learning rate schedule, our adaptive algorithm can achieve similar model quality at many different scales.  This algorithm is useful for many problems, regardless of whether the gradient’s variance is small, large, or constantly changing (as is common in practice).
In contrast, Yin et al. provide thresholds of the variance for which existing algorithms can work well or become less useful.
Overall, we believe AdaScale is an important step toward user-friendly distributed training, and we think many ICLR attendees would find the algorithm interesting and useful.
Our revised version makes the following changes:
* Results of stronger baseline included in Appendix D (Table 5)
* Updated Theorems with tighter bounds
* Updated gain ratio definition with expectations (importantly, this does not change estimation procedure, AdaScale implementation, experiments, etc.) -- it is only necessary for the theory section
* Updated definition of scale invariant algorithm

Meta-review:
role: PC
data: comment: Main summary: Novel rule for scaling learning rate, known as gain ration, for which the effective batch size is increased.

Discussion: 
reviewer 2: main concern is reviewer can't tell if it's better of worse than linear learning rate scaling from their experiment section.
reviewer 3: novlty/contribution is a bit too low for ICLR.
reviewer 1: algorthmic clarity lacking.
Recommendation: all 3 reviewers recommend reject, I agree.

Decision: Reject