Summary
This paper proposes a quantization-aware training scheme for sequence-to-sequence models. During quantization of neural network weights, it is important to choose the **scale** by which float-64 or float-32 parameter values must be scaled, before obtaining a quantized representation using a low-precision format such as int8 or int4. The scalae factor corresponds to the dynamic range of the quantizer, i.e., the maximum and minimum allowable values within which the quantizer input must lie. If the quantizer input lies outside this range, the quantizer is said to be **saturated** and the input is **clipped** to the maximum or minimum value. This is, in general, undesirable, because clipping (or saturation) can lead to unbounded quantization errors, leading to instability of the training / inference process. The quantizer input values that lie outside this dynamic range (and consequently saturates the quantizer) are referred to as **outliers**.
As expected, it is quite important to choose the scale parameter (or hyper-parameter, if treated as so) appropriately, so that the effect of outliers is mitigated. The quantization scheme of this paper essentially proposed alternative options for choosing the dynamic range of the quantizer, which makes it **robust** to such outliers during quantization aware training, and hence, the title of this paper.
The authors consider specific frameworks for quantization neural network, referred to as STE-QAT (Straight-Through Estimator -- Quantization Aware Training) and PQN-QAT (Pseudo-Quantization Noise: Quantization Aware Training). In the latter strategy, the quantization noise is modeled as a uniformly distributed random variable over a bounded range, which is taken into account while computing the gradients during backpropagation. The range of this bounded random variable is proportional to the **scale** chosen for the quantizer, since larger range for a fixed number of bits implies larger quantization error. Clearly, a larger scale would increase robustness to outliers, but also decrease the utility of the trained and quantized model -- hence, the authors show via extensive numerical experiments that their proposed choice for the dynamic range is better than existing alternatives.
Please rectify if I am mistaken in my understanding of the paper. I would be more that happy to rectify it in that case.
Strengths
The paper considers an important problem of quantization of neural network weight parameters, from the perspective of outlier robustness (i.e., preventing quantizer unsaturation). the proposed quantization algorithm has tunable knobs (such as the hyper-parameter $p$), which allows the user the freedom to control the extent of robustness desired.
Whats quite commendable is the extensive numerical evaluations on sequence-to-sequence models, on different tasks and different neural network architectures. They also mention the limitations of LSC and other QAT methods, and how their proposed method(s) overcomes these shortcomings.
It is quite interesting that existing benchmarks, such as *Learnable Scale and Clip (LSC)*, wherein the **scale** parameter of the quantizer is learnt jointly during the training process. Intuitively, one would expect such a joint training process to converge to the optimal scale parameter and surpass the performance of any chosen/proposed scale parameter. But LSC methods, in general, perform poorly, because of the non-differentiability of the quantization operation, which makes it difficult to evaluate gradients with respect to the scale factor, leading to training instabilities.
Weaknesses
I have a few concerns with this work, and I would highly appreciate it if the authors could elaborate upon those. I would be more than happy to rectify my review and/or increase my score post the rebuttal period.
1. Is there anything specific in this paper that is concerned with sequence-to-sequence models? It seems that the general quantization aware training strategy can be extended to even other models. Please correct me if I am mistaken. And if there is anything specific with the proposed robustness-aware quantization scheme that particularly benefit seq2seq models, it would be appreciated if the authors highlighted it.
2. The proposed strategy is not analyzed much in detail. The authors did put some effort in explaining the intuitions behind them and their choice of the scale parameter, but in my opinion, it does lack some theoretical rigor. For instance, looking at the algorithm, it is not very clear what the choice of $p$ should be. Is it a hyper-parameter that needs to be tuned? The authors mention that the value $k = 2$ and $p = 8$ works best -- how did they arrive at this value? What was the set of $(k, p)$ values over which this hyper-parameter tuning was performed? Since this is quantization-aware **training**, isn't it expensive to retrain for every $(k, p)$ pair? A theoretical analysis would shed some light into this issue, and it will be highly appreciated if the authors could point something in this direction.
3. In general, the writing of the paper can be significantly improved. Several things were unclear during a first reading of the paper, such as follows:
(i) In the abstract, the authors mention "... accuracy suffers when there is insufficient regularization signal flowing back to the outliers..." It would be nice if this sentence is elaborated or completely avoided from the abstract. The notion of "outliers" isn't clear here and what does a "regularization signal" imply? It is my opinion that the abstract should be understandable to the general audience to most extent possible.
(ii) Some sentences like (page 2) "This allows for multiple outliers (i.e., scale candidates) within the channel to be simultaneously regularized in an E2E fashion, at every training iteration." -- is not very clear.
(iii) The authors mention that "uniform distribution does a good job in modeling quantization noise" -- It is important to note that this is **note true** in general, for the int8 or int4 (nearest neighbor) quantization. Since this is a strong modeling assumption on which the algorithm is heavily inspired, it would be highly appreciated if the authors make this explicitly clear. The uniform distribution assumption strictly holds true for **non-subtractively dithered quantization**. The origins go way back, see for example -- https://ieeexplore.ieee.org/abstract/document/4542554 (Lemma 1 and Lemma 2).
(iv) Several other statements will benefit from more deliberation such as "Mode 1 performs well when the ratio P=Q is small, where P is the number of parameters and Q is the number of scales in a weight tensor. That does not hold for quantization with
per tensor scale (Q = 1)." -- Is it just an empirical observation? What do you mean by number of scales -- does it mean number of rows, since each row has a scale? Please define terms like channel/sub-channel appropriately.
Questions
I have some more questions, and would appreciate if authors addressed that.
1. Does RAND Mode 3 even help with quantization? Or just robustness?
2. "On the other hand, PQN-QAT exposes weights to a randomly sampled noise even for weight values that would
incur zero quantization rounding error." -- is this a goo thing or a bad thing?
3. Some works do that by doing post-training quantization of pre-trained models, and subsequently fine tuning them (eg., https://arxiv.org/abs/2305.14314). This circumvents the issue of having to train from scratch. It would be nice to mention such approaches in the related work.
4. "A key drawback of creating dedicated scale parameters is that they are tuned to the particular precision (i.e., 4-bit aware model do not have scales for 8-bit inference)." -- How does this work address this issue, could you please elaborate?
Once again, I'd be happy to reconsider my review post-rebuttal.
Rating
5: marginally below the acceptance threshold
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.