Layer-wise Quantization for Quantized Optimistic Dual Averaging

Thank you very much for your time and your comments on our work. We will address your concern as follows:

Q1: To my understanding, the QODA algorithm appears to be agnostic to the choice between layer-wise and global quantization. The authors present results for both approaches under the QODA-based extended Adam optimizer. However, the impact of the QODA algorithm itself does not seem to be explicitly ablated in the experiments. It would be helpful if the authors could provide results comparing Q-GenX training and QODA, both using the global quantization setting.

A1: In Appendix G, we have compared Q-GenX and QODA in full precision setting with no quantization for a bilinear game with the feedback that is is corrupted by relative noise. Under this setting, Q-GenX and QODA are equivalent to GEG and ODA, respectively. From Figure 2, we can observe that the performance of ODA is significantly better than GEG. In fact, it is known in the literature that GEG diverges for certain standard settings like this simple bilinear problem [1,2]. ODA with stepsize seperation can help to prevent this issue [1].

References:

[1] Hsieh et al., No-regret learning in games with noisy feedback: Faster rates and adaptivity via learning rate separation, NeuRIPS 2022

[2] Daskalakis et al., Training GANs with Optimism, ICLR 2018

Thank you for the detailed comments.

1: Regarding (1), we believe our significant speedup with 16 GPUs are sufficent for a theory paper on distributed VIs (DVIs). We list all the related theory DVIs works (in top ML conferences) and no. of GPUs used. [2] is the only one with 16 GPUs like us, while the rest use fewer or only simulated nodes on CPU.

Regarding (2), our main novelty is not to show the known heterogenity in the literature. This heterogenity provides the motivations for us to design our quantization method. Our focus is on providing the theoretical framework and bounds for layerwise quantization and introducing QODA for distributed VIs with strong theoretical guarantees.

2: There is a misunderstanding. Throughout the paper, we usually denote a vector that will be quantized later $\mathbf{v}$, simply represented by the tuple ($||\mathbf{v}||_q, s, u$).

This does not mean that the vector $\mathbf{v}$ is quantized as such. For quantization, we refer you to lines 188-192, where the quantization of $\mathbf{v}$ is defined with quantization random variable $q_\ell$. We will move the Section 2.1 to Section 3 to improve clarity.

3: The assumption 2.4 is standard in stochastic VIs [6, 7]. Classical work [7] list out several general operators in Section 2.2 that falls under this assumption, such as the Matrix Minimax problem [8,Section 2.2.2].

4: $\sigma_R$ is a constant that shows how the error vanishes near a solution of VI [6]. We will fix the writing as: ... and there exists $\sigma_R >0$ such that $\mathbb{E}\left[\|U(\mathbf{x}, \omega)\|^2_\ast\right] \leq \sigma_R \|A(\mathbf{x})\|_\ast^2$.

5+8: The key point is that our layer-wise quantization framework is versatile and independent of the underlying optimizer. For instance, we demonstrate its applicability with Transformer XL (Table 3). Meanwhile, QODA is a novel method we introduce with theoretical guarantees for DVIs, and its contribution is distinct from the broader layer-wise quantization framework. For Qn8, QODA is our novel method along with its guarantees for DVIs since it recovers the order optimal O(1/T) rate which Adam and GD cannot [8].

6: We simplify the notations as follows:

We consider an arbitrary layer type and remove index $m$. Let $\upsilon$ be the index of a level w.r.t an entry $u\in[0,1]$ such that $\ell_{\upsilon} \leq u < \ell_{\upsilon+1}$. Let $\xi(u) = (u - \ell_{\upsilon})/(\ell_{\upsilon+1} - \ell_{\upsilon})$ be the relative distance of $u$ to the level $\upsilon + 1$. For a sequence $\bf{\ell}$, define the random variable $q_{\bf{\ell}}(u)= \ell_{\upsilon} \text{ with probability }1-\xi(u)$ and $=\ell_{\upsilon+1} \text{ with probability }\xi(u)$. The rest will follow a similar simplication.

7: For optimistic methods, this extra variable serves as a predictive term. It is used to "look ahead" to the gradient in the next iteration based on previous gradient information. That is, this anticipatory step with $Y_t$ allows the update rule to be more informed, improving convergence speed and stability.

9: Our novelties here are to design quantization with ODA and to propose adaptive learning rates that leads to optimal guarantees (section 5) under milder assumptions. Firstly, our optimistic approach reduces one “extra” gradient step that extra gradient methods like Q-GenX [1] take, hence reduceing the communication burden by half. Secondly, we design the adaptive learning rates (equation (4)) to obtain guarantees with fewer assumptions.

10: We selected 5-bit compression because our experiments indicated that it is the lowest bit-width at which we reliably recover baseline accuracy. We also tested 4-bit compression, the training performance was nearly identical. Our goal is to show that even with aggressive compression, our method maintains baseline performance while offering significant computational benefits.

11: We refer the reviewer to Figure 1. In lines 373-375 (right half page), we states that in Fig 1, QODA approach not only recovers the baseline accuracy but also improves the performance relative to Q-GenX.

Refs:

[1] Distributed extra-gradient with optimal complexity and communication guarantees, ICLR 2023,

[2] Distributed Methods with Compressed Communication for Solving Variational Inequalities, with Theoretical Guarantees, NeuRIPS 2022

[3] Byzantine-Tolerant Methods for Distributed Variational Inequalities, NeuRIPS 2023

[4] Stochastic Gradient Descent-Ascent: Unified Theory and New Efficient Methods, AISTATS 2023

[5] Communication-Efficient Gradient Descent-Accent Methods for Distributed Variational Inequalities: Unified Analysis and Local Updates, ICLR 2024

[6] Polyak, Introduction to optimization, 1987

[7] Solving Variational Inequalities with Stochastic Mirror-Prox algorithm, Stochastic Systems 2011

[8] Training GANs with Optimism, ICLR 2018

We would like to thank the reviewer for reviewing and giving us feedbacks. We will address your comments in a QnA temmplate.

Q1: The motivation in terms of layer-wise quantization (which is directly related to Mishchenko et al., 2024 or [4]) in training could be further clarified, as layer-wise quantization is not a well-known concept in this field.

A1: Based on the reviewer's feedback, we propose adding the following brief summary at the beginning of Section 3: As mentioned in the introduction, previous literature has shown that different types of layers in deep neural networks learn distinct hierarchical features—from low-level patterns to high-level semantic representations [1, 2]. These layers also vary in the number of parameters they contain and in their impact on final accuracy [3]. Thus, the key motivation for our layer-wise quantization scheme is to account for this heterogeneity and provide a more effective approach to quantizing neural network training.

We also want to stress that our method is not "directly related" to the block (p-)quantization approaches described in [4, 5]. In lines 58-59 on the right-hand side of the page, we emphasize that block (p-)quantization [4, 5] is fundamentally different from the layer-wise quantization presented in our paper, as detailed in Appendix A.2. These are three fundamental distinctions between block quantization and our layer-wise quantization:

• Each of our layer or block in this context has different adaptive sequences of levels (Section 3). This is why our method is named layer-wise. The works [4, 5] on the other hand applies the same p-quantization scheme $\text{Quant}_p$ to blocks with different sizes, implying that the nature and analysis of two methods are very different.

• The way the quantization is calculated for each block or layer are different. [4] study and provide guarantees for the following type of p-quantization (for all blocks): $\widetilde{\Delta}=\|\Delta\|_p \operatorname{sign}(\Delta) \circ \xi,$ where the $\xi$ are stacks of random Bernoulli variables. In our work, the sequence of levels for each layer is adaptively chosen according to the statistical heterogeneity over the course of training (refer to equation (MQV) in the main paper).

• The guarantee in [8, Theorem 3.3] only cover p-quantization rather block p-quantization. In our Theorem 5.1, we provide the quantization variance bound for any arbitrary sequence of levels for each layer in contrast to that for only levels based on p-quantization [4].

Q2: For GANs, do the authors observe any instabilities during training? Or did the authors attempt to train GANs on more challenging datasets besides CIFAR to observe the effectiveness of the method?

A2: We did not observe any additional training instabilities when applying our compression method to GANs. As demonstrated in the FID plots (Figure 1 of the paper), the training curves for QODA closely follow — and at times are even smoother than — the no compression baseline, indicating that our method does not introduce extra instability. We believe the results and speedup observed on CIFAR-10 and CIFAR-100 provide strong evidence that QODA is as stable as the baseline.

References:

[1] Zeiler et al., Visualizing and understanding convolutional networks, ECCV 2014.

[2] He et al., Deep residual learning for image recognition, CVPR 2016.

[3] Dutta et al., On the discrepancy between the theoretical analysis and practical implementations of compressed communication for distributed deep learning, AAAI 2020.

[4] Mishchenko et al., Distributed learning with compressed gradient differences. Optimization Methods and Software 2024.

[5] Wang et al., Theoretically better and numerically faster distributed optimization with smoothness-aware quantization techniques, NeuRIPS 2022.

Thank you very much for your time and your comments on our work. We will address comments in the following QnA format:

Q1: The paper does not analyze accuracy trade-offs-does layer-wise quantization affect generalization or model stability?

A1: We hope to clarify that in Figure 1, we show that our method can recover the accuracy of the full-precision baseline model. This suggests that layerwise quantization does not have an adverse impact on accuracy, while significantly reduces the time per optimization step (Table 1 and 2). Moreover, for the Transformer-XL task, Table 3 shows that layer-wise compression provides a better accuracy-compression trade-off than global compression, obtaining better perplexity score while improving training speed.

Q2: A comparison with other quantization techniques would highlight the benefit of this approach. How does QODA compare to quantization-aware training (QAT) and post-training quantization (PTQ)?

A2: In brief, the methods are complementary. Quantization-Aware Training (QAT) methods seek to produce models with quantized weights and activations during training, and as such, compress these elements during the training process. PTQ methods aim to do so in a single compression step, e.g. by using layer-wise solvers to find a good quantized weight assignment [1,2]. By contrast, we investigate gradient compression, with the goal of reducing the amount of communication employed during the training process itself. As such, the objectives are orthogonal, and QAT methods could be used in conjunction with QODA. We will add the above discussion with references [1,2] to related work section for future manuscripts.

Q3: I believe the Introduction could be improved a bit to provide the context. Paper also seems to be a bit overusing the notations. Writing could be improved a bit.

A3: Understanding the complexity of layer-wise quantization, we have tried to simplify the notations whenever possible (line 176-178, line 278-279). Since the paper is theory-focused, we try to provide the most rigorous model and anlysis. We will try to improve as follows:

• We make and will add the following schematic to help the visualization of our framework in the introduction: https://imgur.com/a/P3AFsPo, which highlights how layer-wise quantization can optimally choose different compression type depending on the importance of the layers.

• For notational simplicity, we consider an arbitrary layer type and remove index $m$ in the initial discussion of Section 3.1 as follows: (Fix a type $m$.) Let $\upsilon$ denote the index of a level with respect to an entry $u\in[0,1]$ such that $\ell_{\upsilon} \leq u < \ell_{\upsilon+1}$. Let $\xi(u) = (u - \ell_{\upsilon})/(\ell_{\upsilon+1} - \ell_{\upsilon})$ be the relative distance of $u$ to the level $\upsilon + 1$. For a sequence $\bf{\ell}$, we define the following random variable $q_{\bf{\ell}}(u) = \ell_{\upsilon} \text{ with probability } 1-\xi(u)$ and $\ell_{\upsilon+1} \text{ with probability } \xi(u)$. The later part in Section 3.1 will follow a similar simplication.

Q4: How does layer-wise quantization affect final model accuracy?

A4: We refer the reviewer to Figure 1, where we compare applying layer-wise quantization with the full-precision baseline. The figure indicates that the adaptive QODA approach not only recovers the full-precision baseline accuracy but also improves convergence relative to Q-GenX [3] (line 373-375 on the right half page). Moreover, Table 1 and 2 show that QODA achieves this accuracy under much faster time per optimization than the full-precision baseline.

Q5: Can QODA be integrated with mixed-precision methods to further improve efficiency? A discussion on QODA’s compatibility with mixed-precision training would be valuable.

A5: Yes, indeed. Mixed-precision methods seek to reduce the computational cost of one or more of the three matrix multiplications employed during training for e.g. linear layers (one on the forward and two on the backward pass). As noted above, QODA is complementary to such techniques, as it serves to reduce the precision in which the gradient is stored and transmitted. Thus, mixed-precision techniques do not directly address communication costs, but could be applied in conjunction with QODA, to reduce both computational and communication costs. From the reviewr feedback, we will elaborate this extension in the future manuscript.

Q6: Adding an ablation study on how different layers benefit from different quantization levels would strengthen the paper.

A6: We are currently running the experiments and will provide the result in the discussion period.

References:

[1] Ashkboos et al., EfQAT: An Efficient Framework for Quantization-Aware Training, 2024.

[2] Frantar et al., GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers, ICLR 2023.

[3] Ramezani-Kebrya et al., Distributed extra-gradient with optimal complexity and communication guarantees, ICLR 2023.