Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 6: Marginally above acceptance threshold
review_content: Main idea: The paper developed a practical second-order preconditioned method, which is based on the Shampoo algorithm, to improve the wall-clock time compared with state-of-the-art first-order methods for training deep networks. 


1.	This paper is heavily drawn from the shampoo algorithm, algorithm designing, and even the later theoretical analysis (such as theorem 3). It appears that the iterative preconditioning method is only incremental improvement over the Shampoo algorithm. However, the iterative method is widely used in the similar spirit of the limited memory of BFGS or LU decomposition.
2.	The main text of this paper has only two lemmas. What is the main theorem of this paper?  
3.	Even though the paper mentioned wall-time everywhere, all the improvements shown in the figures, compared with Adagrad and Adam, are in terms of steps, such as Figure 2(a), Figure 3 - 5 and Table 1. The authors should directly show wall-clock time.
4.	Some minor comments: fonts in Figure 1 are too small; in other figures, the fonts in x-axis and y-axis are also very small; in figure 2 (b), the author should explain more about how to break down the latency of a single step into the parts mentioned in the figure.

##################################

Updates: I have changed my score after reading the author's responses.

review_title: The contributions of this paper are not so clear and appear to be incremental if only an iterative perconditioning method is proposed.

final_score: rating: 7: Good paper, accept
review_content: ######################################################################

1.  Paper Summary 

This work addresses practical challenges in applying full matrix pre-conditioner methods (such as Shampoo) on problems involving large datasets and architectures trained using a distributed setup.  In particular, this work presents a practical extension for the Shampoo algorithm by (1) using only a left or right preconditioner for large layers (2) computing inverse pth roots via coupled Newton iteration algorithms (3) distributing preconditioner computation across CPU cores in a CPU-GPU/TPU cluster and (4) delaying preconditioner computation to occur only once per several steps.  The proposed modifications lead to an implementation of Shampoo that consistently decreases the number of training steps and in certain cases provides a direct wall time improvement over Adagrad/Adam.    

######################################################################

2. Strengths

2.1. The proposed modifications do seem to make Shampoo practical when training large models on large datasets using a distributed setup, and the authors present evidence that their implementation does converge faster (in number of steps and occasionally wall time) than Adam and Adagrad.  The authors also present an example of CTR prediction where their method produces state of the art results.  

2.2. A few of the modifications are mostly well founded: namely, distributing preconditioner computation across unused CPUs and using the coupled Newton iteration for computing the inverse pth root.  

######################################################################


3. Limitations/Questions

3.1. Some of the modifications seem to be heuristic and I found it unclear why these should work well in practice.  In particular, while I found Lemma 1 an interesting theoretical result, it was unclear to me why the proposed preconditioner was actually a good approximation for the full preconditioned matrix from Adagrad.  Does this somehow follow directly from the regret bound provided in the appendix? It appears that these regret bounds only show that Shampoo with the provided extensions converges, but there could be still be a gap between the performance of the modified Shampoo and original Shampoo.  It would be helpful if the authors could comment on this.  I felt similarly about the delayed preconditioner computation: The fact that computing the preconditioner only once every several epochs is interesting, but I didn't quite follow why this should have worked in the first place.  Is there some further intuition the authors could provide for this phenomenon?  

3.2. While the authors demonstrate that their modifications worked well on the settings they described on the bottom of page 2, it appears that there are a lot of elements to tune (described in the appendix) to ensure that their method works well in these settings.  It would be helpful if the authors could clarify a bit more about the basic elements of their method that require significant tuning.  For example, while I understand that there are several works using learning rate scheduling to achieve state of the art results, in most experiments, Adam can be used with a learning rate of 1e-4.  In the proposed method, how often (in number of steps) should one compute the preconditioner matrix? I understand that there is a plot comparing this in Figure 3c, but on smaller datasets, do we actually need to compute the preconditioner often in order to get convergence? 

3.3. Point 3.2 leads to my next concern about the work.  Namely, if the preconditioner does need to be computed more frequently, this would lead to a significant increase in wall time on smaller datasets.  In particular, the authors even mention that there is in fact no wall time improvement on full ImageNet due to a lack of benefit of amortization (I'm assuming this is referring to the frequency of preconditioner computation prior to convergence, but please correct me if I'm mistaken).  Thus, the practicality of the proposed method seems limited to extremely large datasets, and may be of practical use only to large organizations with access to resources capable of handling such large datasets.  However, for several of these organizations, it may be easier to just parallelize a simpler method across more GPUs/TPUs.  Thus, I would be interested in hearing more of a discussion about concrete hardware/software improvements that would boost the performance of the current algorithm over currently used algorithms.  For example, is the main suggested improvement to provide better support for double precision computations? I feel that most other hardware improvements would benefit first order methods anyway.  


######################################################################

4. Score and Rationale


Updates: I have changed my score in light of the author responses and edits and now am in favor of accepting this work.  

Currently, I vote for rejecting.  However, my decision is borderline.  My main concerns with the work are (1) the practical relevance of the work as compared to first order methods and (2) the rationale around the heuristics proposed in the current method.  In particular, I am concerned that while the proposed method consistently improves the number of steps, it does not necessarily improve wall time computations even on ImageNet (which would be considered a large dataset by many organizations).  I am also unsure as to how software/hardware benefits could boost the performance of the current method without significantly boosting the performance of first order methods as well.  I am definitely open to changing my review provided that the authors are able to provide clarification around these and other points in section 3. 


######################################################################


5. Minor Comments

5.1. Apart from the points in section 3, it would be helpful to have a discussion section at the end of the work. 

5.2. I believe the y-axis label is missing in Figure 3 (though I believe it should just be -log(perplexity)).  


######################################################################

Post-Rebuttal Updates:

6. Thank you for running the experiments on small datasets and for clarifying my additional concerns.  In light of your clarifications, I do find this work to be a practical extension of Shampoo and am in favor of accepting the work.  I am still surprised by the fact that delayed pre-conditioner computation still leads to improved convergence over 1st order methods, but the presented empirical evidence demonstrates a consistent benefit in wall time across a number of settings. 
review_title: An interesting approach for making second order methods practical, but needs more clarification

final_score: rating: 7: Good paper, accept
review_content: The paper tackles the important problem of making second-order methods practical for stochastic optimization. The main contributions of the paper include an improvement to a previously proposed Shampoo algorithm that is, on a high level, an extension of AdaGrad to include a second matrix moment (such as covariance of gradients).

The paper addresses challenges along three axes. On the algorithmic level, computing the preconditioners is expensive both in terms of memory and compute and the authors suggest improvements theoretically motivated techniques to make this faster. On a numerical level, the authors put forth a case for investing in 64-bit arithmetic in hardware implementations required for computing the p-th root. This is an important point to make considering that the field has been trending towards the opposite direction and risks a local maxima. On the distributed implementation level, the authors implement the preconditioning computation (done every few steps) asynchronously on the host CPU owing to lack of double precision support on the accelerators attached to the host (GPU/TPU).

Experiments are run on a varied set of tasks (translation, CTR, NLP, vision) and they show fairly significant reduction in the number of optimization steps. Overall, the paper tackles the important problem of second order optimization in a more practical way. The motivations, main contribution and supporting experiments are well laid out.
review_title: Official Blind Review #1

Meta-review: Dear authors,

I like to topic of your paper very much. Indeed, your work is trying to show that 2nd order methods can be efficiently implemented in a distributed environment and can achieve improvement in training times.

However, having worked on distributed computing for many years, I personally think that reporting running time in your work is not very informative (without mentioning what hardware is used during computation), and one cannot understand the connection or reproduce your results. Also, it was not clear how the baselines were implemented, and how the hyper-parameters were tuned. It is also not clear why you haven't picked better benchmarks to compare your work. 

I think that addressing both the concerns from reviewers and the one mentioned above would improve the paper significantly. I would really like to see if accepted in the near future after these issues are fully addressed.

Decision: Reject