Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: The paper shows a working implementation of 8bit states for momentum and second momentum. This is achieved by using block-wise dynamic quantization for efficient compression and fast implementation. They show drastic improvement over f32 diagonal optimizers (mainly Adam) and improvement over a subset of previously proposed sub-linear memory optimizer (AdaFactor f32 variant).
main_review: ```Updated Score```: See my comments why: https://openreview.net/forum?id=shpkpVXzo3h&noteId=PWzRo82xR07 and hopefully authors address the additional comments in the final version.

Paper proposes block-wise quantization (dynamic) to reduce the states (momentum and second moment) in diagonal first order optimizers and successfully implements a 8-bit Adam implementation that works as well as its f32 variant while being memory efficient. Paper further identifies embedding layers as a source of instability and proposes layer norm to improve stability,  and leaves optimizer states in f32 for this layer.

Strengths:
+ Efficient 8bit implementation of Adam on GPUs
+ Careful work to reduce quantization errors, and improvement to existing algorithms (for example quantile estimation) 
+ Blockwise quantization is very neat way to deal with outliers

Weakness/Improvements:
+ Lack of comparison to previously established SOTA for low memory optimization (See comments on AdaFactor)
+ Results are at similarish (lower) batch sizes. It would have been interesting to see the method is applicable for large batch training. One could conjecture that heavy tailed nature noise (https://arxiv.org/abs/1912.03194) and its associated effects on the diagonal preconditioner might make quantization harder.
+ Ignores AdaGrad line of work

==AdaFactor comparisons== 

AdaFactor has the option to reduce the memory used by momentum states completely by replacing it with the cheaper adaptive gradient clipping, and has been used for large model training (see https://arxiv.org/pdf/2006.16668.pdf for hyper-parameters. (See beta1=0))
[1] https://github.com/tensorflow/lingvo/blob/master/lingvo/core/optimizer.py#L1044 

The paper includes the f32 variant, but do not include the bfloat16 variant [1] comparison
[2] https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/utils/adafactor.py See parameter encoding.

== AdaGrad and its variants == 
Paper currently misses citation of the entire AdaGrad-line of work, and comparision/implementation, I hope the authors can address these easily.

[1] https://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf 
[2] Memory efficient adaptive optimization https://proceedings.neurips.cc/paper/2019/file/8f1fa0193ca2b5d2fa0695827d8270e9-Paper.pdf
[3] https://openreview.net/pdf?id=SklKcRNYDH (ICLR 2020)

Does the proposed block quantization work well for AdaGrad variants as well? Reason to ask is, after reading the paper, it's not entirely clear if the proposed approach works for AdaGrad, SM3 and Extreme Tensoring as it accumulates statistics over the entire horizon and has a very different range of values (distribution of values) than Adam's second-moment statistics.

It would be valuable to the community if authors could address this in the paper by demonstrating improvements or showing negative results and challenges to this end.  For example, in this work authors leave compressing states for embedding layers (with layernorm) as future work.

== Some related citations for block wise quantization ==
For weights of a network; similar idea of block-wise quantization (but for binarization) has been employed in "Improving Bi-Real Net with block-wise quantization and multiple-steps binarization on activation", Duy H. Le; Tuan V. Pham


summary_of_the_review: Authors propose 8-bit optimizers for training neural networks, and demonstrated its usefulness in one regime. There are several important ablations (mentioned in the reviews, and I will revise the score accordingly) that are missing which are important it ascertain its wide utility, as well understand its limitations.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This paper addresses the very important problem of reducing the memory footprint of neural networks training. For that matter, the authors propose replacing standard optimizers with their 8-bit quantize counterparts. 
The proposed scheme of optimizer state quantization has three components (I) block-wise quantization which osolat4ed outliers impact on the error (ii) dynamic quantization which quantize both small and large values with high precision and (iii) stable embedding layer which improves the stability of the optimizer during training. 
main_review: **Pros:**
- The authors conducted a wide range of experiments on diverse NLP and computer vision tasks performing consistent memory footprint saving without performance degradation. 
- The proposed method is stand-alone and potentially can be applied in parallel with other compression techniques such as weight, activation quantization, and pruning. 

**Cons:**
- The proposed method is beneficial more to neural networks with a high amount of the parameters proportionally to the activations thus for convolutional neural networks the memory saving ratio is much smaller rather for transformer-based NNs.
- There are several broken references that should be fixed
summary_of_the_review: The proposed paper proposed 8-bit quantized optimizers counterparts for saving memory footprints during training. 
Empirical results show the proposed method's efficiency over a wide range of tasks and architectures. 
I believe that the proposed method introduces an important contribution towards the efficient training of NNs.
Hence I vote for its acceptance in a pre-rebuttal phase.
***********
Post-rebuttal discussion: 
I have fully satisfied with the author's feedback and revision version of the paper according to my and other reviewers' concerns. 
My score remains "8" and I am convinced that this manuscript would significantly contribute to ML community, hence I vote for his acceptance.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper proposes a non-linear block-wise quantization method to reduce the memory overhead of stateful optimizers, without sacrificing performance going from 32bits to 8 bit. The authors combine block-wise quantization with 2 methods to stabilize training: dynamic tree quantization and a stable embedding layer. Results on WMT, GLUE and Moco show the effectiveness of the method.

main_review: This paper shows very impressive empirical results and has an open-sourced codebase which makes it reproducible. Memory saving and time-saving are especially impressive. These empirical results themselves are very valuable to the community given there are not many open-sourced codebases, to begin with.  

The methodology of dynamic (tree) quantization and stable embedding are reasonable but not surprising to the quantization research community. The authors modified the dynamic tree quantization, which also appears to be reasonably motivated. Appendix E seems to support the argument despite there being few places where strong implicit assumptions were made:  mean of gradient zero over time, and constant variance. These should be thoroughly discussed in the appendix as well. 

Quesion:
1.  Why image tasks do not have as much saving as NLP tasks? In table 1, Image related tasks's savings are marginal versus NLP related tasks. Is there a way to make this work for Image, speech task as well?

There are a few minor points that do not affect my score:
- Several broken reference links throughout the text with ??
- Many attempts have been made on quantizing the optimizer states from 32-bit to 8-bit using non-linear quantization, but just not as successful maybe. E.g. <Sun et al. 2019, HFP8> <Pappalardo 2021, Brevitas><Li et al. 2020 End-to-end Quantized Training via Log-Barrier Extensions>  can do such jobs as well.

summary_of_the_review: This paper has very valuable shared resources, and it's worthwhile for the research community to notice.
review_title: null

Meta-review: This paper proposes Adam and Momentum optimizers, where the optimizer state variables are quantized to 8bit using block dynamics quantization. These modifications significantly improve the memory requirements of training models with many parameters (mainly, NLP models). These are useful contributions which will enable training even larger models than possible today. All reviewers were positive.
Decision: Accept (Spotlight)