Paper: Reshape and Adapt for Output Quantization (RAOQ): Quantization-aware Training for In-memory Computing Systems (ICLR-2024)

Reviewer comment (Theory):
The proposed weight reshaping aims to improve the variance of the weight, but opt to use a complicated 4-th order regularization in Equ. (4). It is unclear why such regularization term is selected, and how is it comparing to regularizing the variance directly. Also the choice of regularization strength may need more study.

Author response:
**W-reshape: 2nd order vs. 4th order.** This is a great question. First, the use of the 4th-order kurtosis regularization is inspired by the prior work (Shkolnik et al., 2020) which has shown that such regularization is able to flatten the weight distributions. We leverage such observation, taking it one step further by encouraging the distributions to go beyond flattening to increase the weight variance and to be adjusted by both weight and quantization parameters. Moreover, we actually tested regularization directly on the variance, but found it is not as effective as the kurtosis loss and is harder to control (i.e., high sensitivity to the strength). We show a table below for a comparison (based on the accuracy of MobileNetv2) between the regularization of variance and kurtosis. We can see that kurtosis regularization demonstrates better and more stable behavior. 
| Regularization strength  | 0.00001  | 0.00004  | 0.0008  | 0.002 |
|--------------------------|----------|----------|---------|-------|
| Variance                 | 69.57    | 69.62    | 65.58   | 56.25 |
| Kurtosis                 | 69.94    | 70.02    | 70.24   | 69.65 |

Reviewer comment (Experiments):
My main question is on the setting of ADC quantization problem. Is the hardware specification known at the training time of the model? If the specification is known, why not directly perform a post-activation shift/scaling before ADC to maximize the usage of ADC bins?

Author response:
**Hardware specification.** We do need the hardware specifications for the memory dimension, bit-precisions of activations, weights, and ADCs for a target IMC application. Ideally, we would like to perform dynamic scaling/shifting before ADC quantization, but the analog signal processing nature of IMC prevents this, as analog signals directly feed into the ADC, without power-costly variable-gain amplifiers. This is precisely what makes the algorithmic setting of ADC quantization so challenging and distinct from previous weight/activation quantization settings. Thanks for the question.

Reviewer comment (Evaluation):
The paper does not provide a detailed analysis of the computational cost and memory requirements of the proposed method, which could be important factors to consider in practical applications.

Author response:
**Cost.** Based on our experiments, A-shift and W-reshape add negligible overhead during training. BitAug adds extra computations (i.e., computation and gradient accumulation of other ADC bit precisions), requiring 14% more GPU memory in the last training phase. Considering the challenging ADC quantization problem, we believe such overhead is justified to substantially restore accuracy performance. We now mention the training overhead in our updated manuscript. We would also like to emphasize that our target is to realize efficient inference on IMC applications, which are not affected by this training overhead.

Reviewer comment (Theory):
The paper does not provide a framework for the proposed method.

Author response:
**Framework.** We now provide the summary of our training framework as well as the flow of running IMC inference in Appendix H in our updated manuscript. Thank you for pointing it out. As mentioned in Section 5, all of the proposed methods are implemented in PyTorch. We also attach the example code in supplementary material for your reference.

Reviewer comment (Experiments):
Comparison with multiple QAT approaches is missing. Only one [Bhalgat] is considered.

Author response:
**Other QAT methods.** We now include a comparison between different QAT methods in Appendix G. As we can see, different QAT methods behave similarly in terms of ADC quantization and our methods are applicable to all of them.

Reviewer comment (Evaluation):
Fig. 5 compares energy efficiency of IMC vs. digital accelerators. Was this comparison done at iso-computational or network accuracy, e.g., did both architectures display the same misclassification rate for image classification?

Author response:
**IMC vs. digital accelerators.** Fig. 5 is just showing an energy comparison between IMC and a typical digital accelerator architecture. The motivation of this work is that, while IMC achieves higher efficiency, its accuracy is typically degraded due to ADC quantization. Thus, in this work, we aim to restore the degraded performance to iso-accuracy.

Reviewer comment (Evaluation):
What is the impact of RAOQ on weight and activation sparsity? Does it increase or decrease it or is there no effect? If sparsity is reduced then it will affect complexity reduction techniques that rely on it.

Author response:
**Sparsity.** The table below shows the sparsity for example models with and without applying RAOQ. As we can see, the proposed RAOQ does not affect the sparsity level.  
|   Methods | ResNet50  | MobileNetv2 | YOLOv5s | BERT-base |  
|--------------|-----------|-------------|---------|-----------|  
| Without RAOQ | 1.15%     | 0.88%       | 0.89%   | 1.11%     |  
| With RAOQ    | 1.20%     | 0.79%       | 1.01%   | 1.08%     |

Reviewer comment (Evaluation):
What is the training overhead due to RAOQ? For large AI models, even if training is done once or occasionally, it is already complex. RAOQ will make it even more complicated.

Author response:
**Training overhead.** Based on our experiments, A-shift and W-reshape introduce negligible training overhead. BitAug does require 14% more GPU memory and slows down the training process by 1.5$\times$ for the last training phase, due to the computation and accumulation of gradients associated with different ADC bit precisions. We propose that such training overhead is an incurred cost for improving the accuracy as shown in Table 1. We now include the training overhead in our updated manuscript.

Reviewer comment (Evaluation):
Table 1 (main result) is missing error bars. It is hard to say if the results are statistically significant.

Author response:
**Error bars.** The results shown are averaged from 3 runs, which is observed to provide a statistically stable behavior. We kept the table compact without error bars for clarity. But, below we show some examples with error bars for 4-bit weights and activations with 8-bit ADCs, for your reference:  
|   &nbsp;       | ResNet50          | MobileNetv2     | BERT-base         | YOLOv5s         |  
|----------|-------------------|-----------------|-------------------|-----------------|  
| Accuracy | 76.27$\pm$0.05   | 70.46$\pm$0.07 | 87.67$\pm$0.07  | 33.49$\pm$0.03 |

Reviewer comment (Experiments):
How does RAOQ perform when subject to out-of-distribution (OOD) samples during inference? This is a very practical problem with BP-based training in general.

Author response:
**Out of distribution samples.** RAOQ does not modify the way the input is quantized, thus does not have a direct impact on OOD. Detailed analysis of OOD samples is out of the scope of this work, where we instead align to standard evaluation methodologies employed for QAT studies. We will consider the impact of OOD samples in our future studies. We thank the reviewer for the suggestion.

Reviewer comment (Writing):
Table 3 caption has a typo.

Author response:
**Typo.** We have fixed the typo in our updated manuscript. Thanks for pointing out.

Reviewer comment (Reproducibility):
Reproducibility issue: the paper discusses ADC quantization error, a new type of error that reflects the non-ideal property of analog signal processing; but modeling such error on DNN training is not trivial. Furthermore, the training procedure includes complex tuning of regularization and augmentation, which would also hinder the reproduction of the claimed benefits. Thus, to reproduce the impact of ADC quantization on the models and the effectiveness of RAOQ, it is desirable to provide a reference code that reveals the error modeling of ADC quantization and the hyper-parameter settings.

Author response:
**Reproducibility.** The modeling of ADC quantization has been done routinely in signal processing systems and is quite well understood. In fact, it follows quantization methods that have been used in DSP and deep-learning weight/activation quantization, and can be accurately modeled (Eq. 2). While RAOQ introduces new hyperparameters (i.e., $\lambda_A$ and $\lambda_\kappa$) like other QAT methods (e.g., [1-3]), they are straightforward to adjust. A global setting of $\lambda_A$ can be applied across all models as mentioned in Appendix E, and $\lambda_\kappa$ can be set by analyzing the quantized weight distributions as studied in Appendix A. To demonstrate and validate the effectiveness of RAOQ, we have attached example code for ResNet50 and MobileNetv2 as supplementary material for your reference, and will release more models along with publication.

Reviewer comment (Theory):
They ignore the impact of the distribution shift of W and X on quantization errors.

Author response:
**Effect of distribution shift of W and X on quantization errors.** Both A-shift and W-reshape are aimed at adjusting the output distribution to mitigate the effects of output (ADC) quantization. A-shift does not introduce any additional quantization errors on the activation quantization itself. The impact of W-reshape was studied and included in Appendix A.

Reviewer comment (Theory):
The claim that Var[Y] can be maximized by maximizing Var[W] and E[X^2] is weakly supported.

Author response:
**Support that var[Y] can be maximized by maximizing Var[W] and E[$X^2$].** We agree with the reviewer regarding the weak support of the formal claim around maximization - we did not intend to make the formal claim and have corrected it. Rather, the argument relies on analytical intuition and empirical validation. Analytically, such a relationship is clearly seen before training starts as weights and activations are independent of each other. As training proceeds, a complex correlation between weights and activations arises, along with a correlation between the elements of the activation tensors. This makes it difficult to derive a formal analytical maximization criterion. Instead, we seek to show empirically through the provided data (Fig. 2(c-d)) that the analytical intuition in terms of their proportionality continues to hold through training, and provides simple yet effective ways (i.e., W-reshape and A-shift) to increase Var[Y] to improve SQNR in the presence of ADC quantization. We have clarified this in the updated manuscript.

Reviewer comment (Experiments):
Fig. 3.b shows a compound effect of W-reshape and A-shift, but does not show an individual impact separately.

Author response:
**Individual impact of A-shift and W-reshape.** Thank you for pointing this out. We have modified Fig. 3b to include both their compound effect and their individual effects. As we can see, both of them contribute to the SQNR improvement.

Reviewer comment (Theory):
The rationale for enhancing performance by adding the loss terms for various bit precisions in BitAug to the original loss and gradient calculations seems insufficiently explained. The authors properly showed the difficulty of fine-tuning with ADC quantization, but there is little justification for why the proposed BitAug helps optimization.

Author response:
**Justify BitAug and impact on loss landscape.** We argue that BitAug assists the learning process by avoiding getting stuck at local minima and potentially maintaining a smoother loss landscape. We now include more details in Appendix F to show our analysis for justification and provide the requested loss surface after applying BitAug for 6-bit ADC.

Reviewer comment (Experiments):
The bit precision candidate selection in BitAug is empirical. According to Appendix C, the bit precision candidate set that shows good performance in MobileNetV2 is applied across all other networks, including BERT and ResNet.

Author response:
**BitAug (selection of candidates).** We agree that the selection of these candidates is empirical. But, the approach is to employ a narrow range around the target ADC precision, to optimize to this target, while employing high enough precision to assist in training and low enough precision to ensure adequate parameter adaptation to quantization errors. We would like to emphasize that, unlike weight/activation quantization, ADC quantization does not provide for any learnable quantization parameters, thus making the learning process more difficult. Hence, providing guidance with very different learning objectives would not help much (shown in Table 6). Thus, we mainly focus on the candidates around the target ADC precision. We provide more evidence to justify our arguments in Appendix C.

Reviewer comment (Novelty):
As the authors admitted, the Kurtosis regularization for weight quantization was already proposed by (Shkolnik et al., 2020) with a similar purpose of flattening weight distribution. Therefore, the novelty of W-reshape is not clear.

Author response:
**Novelty on W-reshape.** The target of this work is to develop methods to improve model accuracy with ADC quantization, and our use of W-reshape to realize this is distinct. Our work leverages the observation in (Shkolnik et al., 2020) that kurtosis regularization can be used to flatten the weight distribution, which was aimed to improve the weight quantization robustness. But, both the use and objective of such regularization are very different in our work. First, we are interested in the distribution of quantized weights rather than the floating-point weights. By applying kurtosis regularization on quantized weights, we add one more degree of freedom, allowing both weight parameters and quantization parameters to adapt to the desired distribution. Second, flattening (i.e. uniformly distributed) weights are not necessarily the best choice for output quantization. Thus, we push further, allowing weight distributions to be even endpoints-concentrated (e.g., Fig. 3(a)), specifically to benefit ADC quantization.

Reviewer comment (Evaluation):
Not all operations of the model are mapped to IMC. Based on Appendix E, while the research argues that the operations not mapped to IMC (e.g., the first and last layer in CNN, depth wise convolution in the MobileNet family, the second matrix multiplication inside the Multi Head Attention of BERT) represent a small fraction of the total operations, these operations are typically located intermittently within models. This may induce overhead (e.g., data transfer latency between the host and IMC).

Author response:
**Overhead from layers not mapped to IMC.** The typically envisioned architecture of IMC accelerators involves also using proximate digital acceleration units to handle layers that are not mapped to IMC (e.g., [1]). This avoids communication overheads to a host processor.

Reviewer comment (Reproducibility):
What is the precision of the short-cut (for ResNet, MobileNet, BERT)? (Since the output of MatMul/Conv is taken to the short-cut in pre-activation residual, ADC quantization would affect short-cut precision as well.)

Author response:
**Short-cut precision.** Short-cut connections are quantized and mapped to IMC in the same way as other layers.

Reviewer comment (Theory):
Regarding A-Shift, if enlarging E[X^2] is crucial, why don't we shift X more aggressively toward negative? (instead of stopping at zero-center, in Fig. 3.a(bottom))

Author response:
**Amount of shift in A-shift.** Great question. The amount of shift is limited by the minimum/maximum value that the quantized activation range can support. We have already maximized the shift within this limit.

Reviewer comment (Experiments):
For the same ADC bit-precision, how can BitAug adjust the loss surface? (e.g., can you compare loss surface of 6-bit ADC in Fig. 4 without and with BitAug?)

Author response:
**How can BitAug adjust the loss surface.** Please refer to our main reply to **Justify BitAug and impact on loss landscape**. We have plotted the loss surfaces for 6-bit ADC per your request in Fig. 4 (without BitAug) and Fig. 16 (with BitAug).