State-of-the-art models are now trained with billions of parameters, reaching hardware limits in terms of memory consumption. This has created a recent demand for memory-efficient optimizers. To this end, we investigate the limits and performance tradeoffs of memory-efficient adaptively preconditioned gradient methods. We propose extreme tensoring for high-dimensional stochastic optimization, showing that an optimizer needs very little memory to benefit from adaptive preconditioning. Our technique applies to arbitrary models (not necessarily with tensor-shaped parameters), and is accompanied by regret and convergence guarantees, which shed light on the tradeoffs between preconditioner quality and expressivity. On a large-scale NLP model, we reduce the optimizer memory overhead by three orders of magnitude, without degrading performance.
============= 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.