AffineQuant: Affine Transformation Quantization for Large Language Models

W1: When the paper shows that 4-bit quantization has a higher information density, the results emphasizing 2-bit quantization are misleading.

Reply: We admit the conclusion that 4 bit quantization has a higher information density than 2 bit quantization. To prevent any potential misinterpretation, we have removed references to the 2 bit quantization results from the abstract and the body of the paper, replacing them with the 4/4 bit results that incur no additional overhead. Our experimental results continue to outperform after modifying the quantization configuration.

W2: In practice it is necessary to store affine matrices, which doubles the model size and thus negates the FLOPs and memory benefits of quantization.

Reply: The modified AffineQuant doesn't need to store matrices because it can be fused into the weight and bias of LayerNorm and Linear. Specifically, we optimize only the diagonal elements of the affine matrix, enabling integration with the weight and bias of LayerNorm. Additionally, we apply the affine matrix between v and out_proj, as well as between q,k output activations, allowing fusion with the weight and bias of the q, k, v, out_proj Linear layer. We take the linear layer v and the previous LayerNorm as an example and formally,

Where $\gamma$, ${\gamma}^{'}$, $\beta$, ${\beta}^{'}$ are the LayerNorm parameters before and after the affine transformation. $diag(\cdot)$ is the diagonal matrix fetching operation. $A_{qkv}$ is the optimal affine matrix at the Linear layer $q$,$k$,$v$. $W_v$, $bias_v$ are the weights and bias of the Linear layer $v$. $A_{vo}^{-1}$ is the inverse of the optimal affine matrix between $v$ and out projection.

Therefore, AffineQuant does not incur additional FLOPs as well as memory overhead. Meanwhile, we compare the 4/4 bit quantization performance of LLaMA1&2 models on WikiText2 and C4 datasets in the following table.

Q1-2: How does the process work at inference time? How fast is inference for this method compared to baselines?

Reply: For the experiments in the table above we optimize the matrix after LayerNorm for diagonal elements only. This allows us to fuse the affine matrix with the LayerNorm weights and bias. For the remaining positions, we apply matrix transformations. After optimization, the matrix can be further fused with the weights and bias of the preceding linear layer. Consequently, AffineQuant can be achieved without introducing any additional overhead to model inference. We utilize the MLC-LLM [2] library to compile the LLaMA model at various scales under only-weight quantization configuration on a single A100, and the resulting inference speeds are presented in the following table.

Q3: scaling curves that show Pareto fronts in terms of performance per bit of total memory footprint.

Reply: In the Appendix, we provide the PPL vs. weight-memory Pareto-optimal curves for LLaMA1&2 models of different sizes in the 4/4 bit quantization configuration. The results clearly demonstrate that AffineQuant consistently outperforms the current State-Of-The-Art method, OmniQuant, without any additional overhead.

W1: Comparisons and analyses with FlexRound serve to emphasize AffineQuant's novelty and validity.

Reply: FlexRound [1] is a commendable paper. However, it is orthogonal to our approach. Because,

1. FlexRound proposes an element-by-element division weight rounding scheme designed to break away from the upper or lower rounding of Adaround [2], BRECQ [3]. AffineQuant proposes an affine transformation matrix to make the distribution of weights and activations more suitable for quantization.
2. FlexRound's rounding scheme can explore more aggressive rounding scenarios for a large magnitude of weight values to achieve the global reconstruction error optimization. AffineQuant proposes the gradual mask to provide stable computation of the inverse of the affine matrix during the optimization process.
3. The optimized weight rounding values are orthogonal to the optimized weight activation distribution. After using AffineQuant to fuse the optimal affine matrix into the original model, we can use FlexRound to minimize the reconstruction error by performing secondary reconstruction with the optimized model weight rounding values.

AffineQuant focuses on post-training quantization without fine-tuning, making the LoRA fine-tuning experiments in FlexRound not applicable to our approach. Furthermore, FlexRound does not provide clear information on whether it quantizes the input activations of both matmul operations for weight-activation quantization in LLM. In contrast, AffineQuant quantizes all input activations except for the softmax input activation of the second matmul. To compare their performance, we present the results of FlexRound in the 4/16 bit quantization configuration in the table below. It is important to note that AffineQuant, as shown in the table, only updates the diagonal elements of the matrix after LayerNorm to maintain model inference without any additional overhead. We place the comparison experiments in Appendix of the paper.

W2: Loss of some layers during MRE may not be directly related to improved model performance.

Reply: You raise a very good question. We have done sampling experiments based on this question. Specifically, we maintain consistent matrix initialization while randomly sampling the stability factor $\alpha$, which influences loss convergence, for LLaMA-7B and OPT-6.7B. Using AffineQuant, we obtain the performance of 4/4 bit quantized models based on the sampled solution. In the Appendix of our paper, we present scatter plots depicting the output loss of the last transformer block and the corresponding model performance. These plots demonstrate a significant positive correlation between loss and model performance, with correlation coefficients of 0.95,0.96 on OPT-6.7B and LLaMA-7B, respectively. Based on this observation, we conclude that the quantization loss of the last transformer block's output exhibits a strong correlation with overall model performance. Thanks again for your constructive comments.

Thanks again for your great efforts and constructive advice in reviewing this paper! With the discussion period drawing to a close, we expect your feedback and thoughts on our reply. We put a significant effort into our response, with several new experiments and discussions. We sincerely hope you can consider our reply in your assessment. We look forward to hearing from you, and we can further address unclear explanations and remaining concerns if any.

Regards,

Authors

W1: Is the increased computational cost the reason for not reporting the results of the larger model?

Reply: For OPT-66B and LLaMA-65B models the optimization time is 82h and 78h respectively, which is much reduced compared to LLM-QAT [1] (80 GPU hours of fine-tuning time on LLaMA-7B). The ppl of OPT-66B and LLaMA-65B on WikiText2 are 10.19 and 8.85, respectively. we will report the larger model runtime comparison and the model performance in a future version.

Q1: What do the authors think about just fine-tuning the model after some standard quantization method in terms of a computation vs quantization performance tradeoff?

Reply: AffineQuant methods, based on post-training quantization, offer significant efficiency and performance advantages over fine-tuning methods. This holds true for both training and inference overhead.

1. In terms of training costs, Quantization Aware-Training (QAT) for LLM involves the fine-tuning process to address quantization errors. LLM-QAT [1] requires 80 GPU hours to fine-tune the model for a single A100 on LLaMA-7B. The time taken to generate the data used for fine-tuning the model on LLaMA-7B is not included in this count. In comparison, AffineQuant quantizes LLaMA-7B in just 6.5 hours on a single A100. Similarly, for LLaMA-65B, AffineQuant takes only 78 hours to obtain the quantized model, which is comparable to LLM-QAT on LLaMA-7B. However, LLM-QAT does not report the training cost of the 65B model. Despite the significantly expensive fine-tuning process, LLM-QAT does not demonstrate any advantage over AffineQuant in terms of experimental results. Please refer to the table below for further details.

2. For the inference cost, for the experiments in the table above we optimize the matrix after LayerNorm for diagonal elements only. This allows us to fuse the affine matrix with the LayerNorm weights and bias. For the remaining positions, we apply matrix transformations. After optimization, the matrix can be further fused with the weights and bias of the preceding linear layer. Consequently, AffineQuant can be achieved without introducing any additional overhead to model inference. We utilize the MLC-LLM [2] library to compile the LLaMA model at various scales under only-weight quantization configuration on a single A100, and the resulting inference speeds are presented in the following table.

[1] Liu et al., LLM-QAT: Data-Free Quantization Aware Training for Large Language Models.

[2] MLC-LLM, https://github.com/mlc-ai/mlc-llm.

W1: The lack of clarity in explaining the implementation of the core concept.

Reply:

1. We apologize for the lack of clarity in the presentation of the concept. As you mentioned, we have modified our description in the newest version:
2. Gradual Mask (GM) is a learning rate regulator that achieves its purpose by element-wise dot-producting with the matrix A.
3. Specifically, the impact of the GM matrix on the optimization process can be divided into two aspects. Here, we present the optimization process for the matrix A after incorporating the GM.

​ Where $\circ$ is the Hadamard product. $A_{e}$ and $GM_{e}$ are the matrices $A$ and Gradual Mask (GM) matrix in epoch $e$, respectively. $\eta$ is the learning rate of matrix $A$. $L$ is the optimization loss. The GM matrix effectively reduces the magnitude of non-principal diagonal elements in matrix $A$ during forward propagation when the stability factor $\alpha$ is less than $1$. This ensures the existence of a stable inverse matrix of $A^{*}$ in the optimization process during epoch $e$, as per the Levy-Desplanques theorem. In backward propagation, GM affects the learning rate $\eta$, thereby suppressing the update rate of non-primary diagonal elements in matrix $A$. Consequently, the impact of GM on $\eta$ ensures that matrix $A$ in epoch $e+1$ maintains strictly diagonally dominant, satisfying the Levy-Desplanques theorem.

W2: The hyper-parameters $b$ (bit-width) and $\alpha$ (stability factor) may introduce computational overhead.

Reply: Both the hyperparameter $b$ (bit-width) and $\alpha$ are determined by the requirements in different application scenarios. We thus conclude that, neither hyperparameter is chosen in a way that introduces additional overhead to the optimization process.

Q1: How the authors address the non-differentiable components in its optimization approach.

Reply: The gradient approximation is determined from Straight-Through Estimator (STE) [1], which is formally represented by the following equation:

Please kindly refer to the round_ste function in the anonymous link quantize/quantizer.py for the precise code implementation.

[1] Bengio~et al., Estimating or propagating gradients through stochastic neurons for conditional computation.

Thank you for your suggestion. We will implement adaptive adjustments to the stability factor $\alpha$.

Thank you for your response. I believe that incorporating an adaptive $\alpha$ could further enhance GM's performance. Specifically, I suggest considering bounding $\alpha$ at each iteration using the initial values of matrix $N$, the learning rate $\eta$, and the gradient value at epoch $h$. Given that $\eta$ decreases over epochs, adjusting the parameter $\alpha$ accordingly might provide additional benefits.

3. Meanwhile, we compare the 4/4 bit quantization performance of LLaMA1&2 models on WikiText2 and C4 datasets in the following table. We evaluate the quantization performance of LLaMA-7B, 13B, 30B on six zero-shot datasets using 4/4 bit quantization in the following table. There was no additional inference overhead in either experiment.

[1] Liu et al., LLM-QAT: Data-Free Quantization Aware Training for Large Language Models.

[2] MLC-LLM, https://github.com/mlc-ai/mlc-llm.