Venue: ICLR 2020 Conference

Reviews:

final_score: rating: 8: Accept
review_content: ============= Update after rebuttal

Thanks for the clarifications; I have updated by score and recommend acceptance.

Please do not forget to implement the promise changes in the camera ready version (e.g. about L-BFGS tradeoff work; as well as gap between OCO and non-convex framework).


==============

Motivated from NLP applications where models with billions of parameters are used, this paper proposes a *memory efficient* variant of Adagrad by maintaining a rank-one tensor approximation of the "second-moment" accumulator normally used in Adagrad. From the theoretical side, a simple regret bound is proved which provides an intuitive quantity quantifying the convergence guarantee loss for memory-efficient version (vs. Adagrad), and which is empirically evaluated to be small on a large-scale language modeling task (Section 5.3). From the empirical side, the memory vs. generalization performance tradeoff is evaluated on this language modeling task, showing that similar performance to Adagrad can be obtained with much a smaller (optimization overhead) memory footprint. A convex toy task also indicate a similar result.

I like this paper, I am leaning towards acceptance. The write-up can be improved in a few places (see detailed comment below); but overall, I find the idea refreshing for optimization and the proof is simple and elegant.

The main motivation for the paper is that these large models used in NLP are often making us hit the memory limitation of the hardware just to store the model, and so one then has to tradeoff the size of the model with the memory requirement of the *optimization* algorithm (e.g. one step-size accumulator per dimension for Adagrad). Trying to get the gains of an adaptive preconditioned gradient method but with lower memory footprint thus seems valuable (and Section 5.2 which compares doubling the size of the model + using their memory efficient method vs. original model + Adagrad highlights the gains one can get).

On the other hand, I wish the paper provided a bit more intuition on when the rank-one tensor approximation structure will give good results (e.g. Section 5.3 gives an empirical measure; but what about some simple theoretical examples which give low values, to provide more intuitions?). In particular, the main inequality relating the Adagrad step-size with the "extreme indexing" step-size is arising from the first equation in the proof of Lemma 4.3 on p.5 -- basically one replace one entry of the squared gradient to store with the whole sum over all entries of the squared gradient except a slice. This inequality can be very loose (O(d) in the worst-case), so it is surprising that one could get good approximation ratios much better than the O(sqrt(d)) worst-case (a square-root is taken afterwards), and I wonder what structure yields this.

I also note that Anil et al. (2019) tackles a similar problem with a different approach, and so ideally a more detailed comparison (theoretically and empirically) would be provided (rather than just one sentence as in the current submission), especially since this appeared on arXiv 9 months before the ICLR deadline. But the authors presented it as concurrent work (and it does not appear published anywhere yet) [and a Google search shows the authors had a first version on arXiv only a few weeks after this one, so this seems righs], and so I decided to not hold it against this submission. Disclaimer: I made my general opinion about the paper before doing these Google searches and finding the paper on arXiv.

I am not very familiar with the modern transformer architecture experiments, so cannot evaluate their quality, but they seem sensible from an outsider's perspective.

== Detailed comments ==

- p.1 "the first empirical study of the tradeoff between training convergence and memory in the optimizer." This should be properly qualified by "in the Adagrad setup" or something like that; I am pretty sure it is false *for general optimizers*. For example, I guess there are several papers which empirically studied the memory - performance tradeoff for L-BFGS optimizers...

- Important: I think the paper should be more transparent about the *big gap* between the online *convex* optimization framework and stochastic optimization for a *non-convex* loss. There is only one sentence at the end of Section 2.1 mentioning vaguely that Agarwal et al. (2018) provide a reduction from stochastic non-convex setting to the online convex setting, but this is buried in the appendix of the said paper, with several caveats, and for a modified algorithm (i.e. as far as I understand, no convergence guarantee would be given for Algorithm 1 to find a stationary point on a non-convex loss; but one could apply Algorithm 1 on a series of convex problems to obtain overall guarantees on the non-convex loss [but this is a different algorithm!]). I suggest that either the argument of Agarwal et al. (2018) is summarized in a few sentences at the end of Section 2.1 to clarify the real link; or give more caveats [also just before 4.1] (e.g., my current guess is that we use the convex analysis setup to gain insights on the behavior in a controlled setup; and then just hope that some of it applies in the non-convex setup, even though no guarantees is provided whatsoever). 

- Some undefined notation: line 8 of Algorithm 1, it seems they use the dot for the Hadamard product between two vectors. The big dot in Lemma 4.5 is most likely the dot product between matrices (but please mention explicitly for clarity). 

- Typos in Section 5.1 & 5.2: several references seem wrong. E.g. Figure 5 is probably Figure 2; Table 4 is Table 1, etc... Please correct!

- Clarity: I suggest to put much more description in the captions of Figure 2, Table 2, etc. (for example it was much clearer in their arXiv version). I was very confused by the extra light blue dot in Figure 2 which is only explained in 5.2; I suggest that it is already mentioned (with forward pointer) in the caption. Use more of p.9 for the clarity sake...

- Appendix B.1: please provide hardware information when you give wall-clock comparison.
review_title: Official Blind Review #2

final_score: rating: 6: Weak Accept
review_content: This paper considers the problem of the need for memory-efficient optimizers given the increase in model complexity. They study memory-efficient adaptively preconditioned gradient methods, and see the trade-offs among expressivity and preconditioner quality. They show results on a large-scale NLP model, and show that the memory overhead can be reduced by 3 orders of magnitude, without sacrificing performance. 

+ show tradeoff among training convergence and memory in the optimizer.
+ introduce extreme tensoring -- a modification that can be applied to *any* 2nd moment-based adaptive optimizer. It uses a compressed preconditioner. 
+ extend diagonal Shampoo to tensor factorization 
+ nicely written Related Work.
+ applied their idea to widely used optimizers, such as AdaGrad, Adam, etc.
+ the derived regret bound is only a multiplicative constant from the regret bound of AdaGrad.
+ interesting experiments showing tradeoff among lack of memory (SGD) and full memory (AdaGrad), applied on a real-world machine learning setting of large-scale natural language modeling with Transformers. They show how their extreme tensoring modification can achieve an intermediate ground among lack of memory and full memory.
- I am missing a plot similar to Figure 4 but instead of being on the synthetic data, being on the application of NLP.

Overall, although I do not have a lot of experience in this area, it appears to be that the introduced family of algorithms can be promising as a nice interpolation among SGD-type algorithms and AdaGrad ones. I would personally have preferred to see a wider range of experiments, considering other application scenarios besides NLP, just to see the wide success of the proposed approach. Also, I would urge the authors to spend a bit more text on explaining the algorithm, so that it is easy for practitioners to try it out.
review_title: Official Blind Review #1

final_score: rating: 6: Weak Accept
review_content: 
============================== Update after rebuttal =======================================================

I did not have any major concerns about the paper in my initial review, only some suggestions for improving the presentation. The authors have addressed most of these issues in their revision. I would like to keep my score as it is. The work seems simple and sound, but somewhat incremental. To have a more meaningful impact, I strongly encourage the authors to make an optimized implementation publicly available as an easy-to-use, plug-and-play type optimizer.

========================================================================================================

This paper proposes a new memory-efficient pre-conditioning scheme for stochastic optimizers. The basic idea is to store a coarse-grained pre-conditioner, expressed as a rank-one tensor product, instead of the full-dimensional pre-conditioner typically used in algorithms like AdaGrad. I am not very familiar with prior work in this literature, but the proposed approach seems simple and sound. 

A regret bound is provided for the proposed method, however the analysis here seems to be a straightforward application of the results and techniques from a few prior works. So, I was a bit surprised to see so much space devoted to the proofs. These can be safely moved to the appendix in my opinion. Moreover, the bound does not seem to be very useful in practice. For example, in simulations in Figure 3, the proposed ET1 performs better than AdaGrad, but the bound is not able to capture this at all. 

The presentation of the experimental results in section 5 can be improved in my opinion. The authors keep referring to “Figure 5” in this section, but I think this is a typo and these should be “Figure 2” instead. In Figure 2, please indicate on the figure itself what dark blue and light blue colors correspond to (smaller and larger models, respectively). 

In Appendix B, wall clock results are presented for different algorithms. These show that the proposed approach is slower than standard algorithms like Adam or AdaGrad. Please explicitly mention this result in the main text (last paragraph of section 5.1) and discuss why this is the case (is this because of the extra reshaping operations required in the updates?).
review_title: Official Blind Review #3

Meta-review: Post author rebuttal the score of this paper increased.
Discussions with reviewers were substantive and the AC recommends acceptance.
Decision: Accept (Poster)