FlatQuant: Flatness Matters for LLM Quantization

Authors: Yuxuan Sun, Ruikang Liu, Haoli Bai, Han Bao, Kang Zhao, Yuening Li, JiaxinHu, Xianzhi Yu, Lu Hou, Chun Yuan, Xin Jiang, Wulong Liu, Jun Yao
Primary area: foundation or frontier models, including LLMs
Keywords: flatness, post-training quantization, LLM, pre-quantization transformation

Abstract:
Recently, quantization has been widely used for the compression and acceleration of large language models~(LLMs). 
Due to the outliers in LLMs, it is crucial to flatten weights and activations to minimize quantization error with the equally spaced quantization points. Prior research explores various pre-quantization transformations to suppress outliers, such as per-channel scaling and Hadamard transformation. However, we observe that these transformed weights and activations can still remain steep and outspread. In this paper, we propose FlatQuant (Fast and Learnable Affine Transformation), a new post-training quantization approach to enhance flatness of weights and activations. Our approach identifies optimal affine transformations tailored to each linear layer, calibrated in hours via a lightweight objective. To reduce runtime overhead, we apply Kronecker decomposition to the transformation matrices, and fuse all operations in FlatQuant into a single kernel. Extensive experiments show that FlatQuant sets up a new state-of-the-art quantization benchmark. For instance, it achieves less than 1% accuracy drop for W4A4 quantization on the LLaMA-3-70B model, surpassing SpinQuant by 7.5%. For inference latency, FlatQuant reduces the slowdown induced by pre-quantization transformation from 0.26x of QuaRot to merely 0.07x, bringing up to 2.3x speedup for prefill and 1.7x speedup for decoding, respectively. Code will be released upon acceptance.

Decision: reject
Reviewer scores: 8, 5, 3, 5, 5

Reviews and discussion:
<thread>
creation date:
1729117793253

thread_type:
official_review

<thread_object>
creator: reviewer_twue

<content>
summary:
FLATQUANT, a new post-training quantization method, enhances weight and activation flatness via optimal affine transformations, calibrated quickly with minimal runtime overhead using Kronecker decomposition. It sets a new state-of-the-art with W4A4 quantization and significantly improves inference speed.

strengths:
1. This paper sets a new state-of-the-art performance.
2. Extensive experiments validate its effectiveness.
3. Fast CUDA kernel implementation enhances its applicability.

weaknesses:
1. The manuscript does not clearly explain how the invertibility of $\mathbf{P}$ is maintained during training. Detailing the methods used to ensure this could improve understanding of the method.
2. I notice the use of Kronecker decomposition during inference and SVD during calibration, but the specific methodologies for Kronecker decomposition are not detailed. Elaborating on these could clarify the approach.
3. The potential approximation errors from Kronecker decomposition and their impact on model performance are not mentioned. Experimental results in this area would be beneficial.
4. The overhead of dynamic activation clipping during inference and the individual contributions of $\alpha_a$ and $\alpha_w$, should be explored.
5. The ablation study should include separate evaluations for "PS", "LCT", and "PS+LCT" to comprehensively demonstrate the improvements from each component.

questions:
1. The distribution of $\mathbf{W}_d$, $\mathbf{W}_q$, $\mathbf{W}_k$, and $\mathbf{W}_v$ and their corresponding activations is not shown. Visualizations like Fig. 2 for these weights and activations would be beneficial.
2. I am curious about the speedup ratio of the method on Hopper architecture GPUs (e.g., H100, H800).
3. Could the authors explain why the per-channel scaling is employed? A proper $\mathbf{P}'$ can contain the ability to balance outliers between the weights and activations, for example,  $\mathbf{P}'=\mathbf{P}\text{dig}(\mathbf{c})$.
</content>
</thread_object>

</thread>
<thread>
creation date:
1729662670471

thread_type:
official_review

<thread_object>
creator: reviewer_x7ak

<content>
summary:
This paper builds upon AffineQuant by leveraging Kronecker decomposition to design efficient kernels that integrate affine transformations and quantization. This approach helps mitigate the doubled computation and memory access overhead caused by the pre-quantization transformation in AffineQuant. Experimental results demonstrate that FlatQuant achieves superior accuracy compared to existing methods under the W4A4 configuration.

strengths:
* The paper employs Kronecker decomposition to design efficient kernels that integrate affine transformations and quantization, reducing the doubled computational and memory overhead caused by the pre-quantization transformation in AffineQuant.

* The paper further enhances accuracy by introducing activation clipping through learning, which yields notable improvements.

weaknesses:
* Even with Kronecker decomposition, FlatQuant still requires a pre-quantization step before each linear input, leading to O(n) memory accesses and O(nlogn) computational overhead, where n is the model's feature dimension. For a single block of LLaMA, this transformation needs to be applied more than six times, increasing the model’s complexity. In contrast, SpinQuant requires only one pre-quantization transformation, which may offer a more significant advantage.  I would be very grateful if more comparisons on the additional computational complexity and parameter quantity improvement brought by Spinquant and FlatQuant could be provided.

* AffineQuant ensures that affine transformations remain invertible by employing diagonal dominance, but it’s unclear how FlatQuant maintains invertibility.  This information seems to be missing from the current paper and would be important for understanding the method's robustness. 

* It is also puzzling that the paper does not report zero-shot performance for weight-only quantization. As shown in other works, LWC tends to cause overfitting, and the perplexity metric may not reflect true performance.  I would be very grateful if more weight-only results could be provided.

questions:
* OmniQuant only trains smooth parameters, while FlatQuant adds the inverse affine matrix, pre-quantization, SVD decomposition, and learnable activation clipping operations. Strangely, the paper reports that FlatQuant is 10 times faster than OmniQuant, which raises concerns. A more equitable and transparent comparison will be employed if a more detailed breakdown of the timing comparisons is provided, including information on the implementation details and the hardware used for both methods.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730273388204

thread_type:
official_review

<thread_object>
creator: reviewer_rqeg

<content>
summary:
This paper proposes FLATQUANT, a fast learnable affine transformation used to enhance the smoothing of weights and activation values, while the authors propose the use of Kronecker decomposition and operator fusion to reduce the runtime burden. Extensive experiments demonstrate the effectiveness of FLATQUANT in the W4A4 case.

strengths:
1. The paper was well written and clearly presented;
2. The authors show the effectiveness of the algorithm through extensive experiments;
3. The authors reduce the runtime burden by means of Kronecker decomposition and operator fusion, which is meaningful for many scenarios.

weaknesses:
Although the author's method achieved significant gains, there were a few points that could not be ignored that contributed to my initial score:
1. The flatness and quantization error expressed in the authors' paper do not strictly correspond. A simple example [3,4,5,6] has significantly higher quantization error at 2-bit than [0,0,0,sqrt(86)]; in fact, the authors only optimize the loss in Eq. (4) in their paper, and there is no direct correlation with flatness;
2. In the paper, the authors mainly compare the method with QuaRot and SpinQuant, and from my point of view, the way that authors compare is unfair, because the computational overhead of FLATQUANT is totally different from QuaRot and SpinQuant. Taking QuaRot as an example, only the rotation matrix in Query and Key and the rotation matrix after SiLU need to participate in the online computation in QuaRot, while all other matrices can be merged into the weights, while FLATQUANT has a large number of matrix multiplications that need to participate in the matrix computation of the floating-point, which introduces quite a lot of computational burdens compared to QuaRot, so I think the author's comparison is totally unfair. Meanwhile, the authors did not clarify this part of the paper;
3. While I strongly agree that the author's use of Kronecker decomposition and operator fusion reduces the runtime burden, which is also valid for QuaRot, the author's Speed Up comparison of Online Transformations in Figure.6 leads readers to mistakenly believe that FLATQUANT's accuracy-performance of the In fact, on the one hand, QuaRot's additional floating-point computation is less than FLATQUANT's, and on the other hand, the Kronecker decomposition and operator fusion reduces Hadamard's runtime burden under GEMM computation (instead of Walsh-Hadamard Transformation), so I think the comparison in Figure.6 is biased and misleading.

In summary, while I agree with the Kronecker decomposition and operator fusion approach used by the authors to minimize the burden on the online runtime, given the problems with the authors' motivation and experimental comparisons, I give my scores.

questions:
See weekness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730448099241

thread_type:
official_review

<thread_object>
creator: reviewer_ydzg

<content>
summary:
This paper propose FLATQUANT to enhance flatness of weight and activation. It consist of:

1) Learnable affine transformation;
2) Kronecker Decomposition to decompose the large transformation into 2 small matrix;
3) Per-Channel Scaling
4) Learnable Clipping Thresholds.

This paper is easy to follow. Experiments show relatively good performance compared to current SOTA like SpinQuant. Ablation study part explains  the effect of each technique.

strengths:
1. This paper propose Kronecker decomposition to decompose the large matrix P into 2 small matrix. The varying size of decomposition  have limited impact on the perplexity of generated text, and the speedup peaks when P1 and P2 are of equal size.
2. This paper propose learnable clipping threshold, where current works usually adopts grid search.
3. The figures in this paper helps easy understanding of the proposed FLATQUANT.

weaknesses:
1. In perplexity, learnable transformation and learnable scaling of FLATQUANT are borrowed from existing works. 
2. Learnable clipping threshold does not compare with existing CNN works like NWQ[1], also does not show detailed analysis on learnable clipping threshold  why not causing important outliers clipped, as Llm.int8() [2].
3. In speedup, Kronecker Decomposition still involves online transformation and can not achieve pure int computation.

[1] Leveraging Inter-Layer Dependency for Post -Training Quantization . Changbao Wang

[2] LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale. Tim Dettmers

questions:
1. What's the perplexity improvement of pure learnable clipping threshold? In table3, in perplexity, whether the difference of SpinQuant and FLATQUANT is the learnable clipping threshold? Is the perplexity improvement mainly from pure learnable clipping threshold? 
2. In Figure5, why the speedup curve is not almost-symmetricly distributed? The speedup of 2048 is substantially slower  than  2's.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730607734203

thread_type:
official_review

<thread_object>
creator: reviewer_m9t2

<content>
summary:
This paper introduces FlatQuant, a post-training quantization method for LLMs that focuses on flattening the distributions of weights and activations to reduce the quantization error. The key innovation is using learnable affine transformations optimized per layer, with Kronecker decomposition to reduce overhead. The authors also develop an efficient kernel that fuses the transformations with quantization.

The main contributions are:
1. A theoretical framework highlighting the importance of flatness in LLM quantization
2. The FlatQuant method that achieves state-of-the-art W4A4 quantization results
3. An efficient implementation that provides up to 2.3x speedup for prefill and 1.7x for decoding, with minimal overhead from transformations (0.07x vs QuaRot's 0.26x)

The approach is validated across multiple model sizes and quantization settings with comprehensive experiments and ablations.

strengths:
1. The observation that flatness is the main influence on quantization error is insightful and well-supported. The motivation to improve flatness through learnable transformations is theoretically sound.

2. The Kronecker decomposition of the transformation matrix effectively reduces both memory and computation costs while maintaining performance. The training methodology is well-designed and practical.

3. The CUDA kernel implementation effectively reduces inference cost through fusion of operations.

4. The evaluation demonstrates that FlatQuant achieves impressive results, particularly: Less than 1% accuracy drop for W4A4 quantization on LLaMA-3-70B; Consistent improvements across different model sizes and tasks; Better speedups compared to existing methods like QuaRot.

weaknesses:
1. Figure 3 is confusing and needs better explanation. The caption should better explain the relationships between subfigures and clarify the flow of operations.

2. While the efficient implementation and speedup is impressive, theoretical complexity analysis is crucial for algorithm development. The evaluation lacks important theoretical analysis:
- No detailed computation complexity analysis
- Missing information about total parameter counts (including decomposed transform matrices)
- No comprehensive memory consumption analysis for inference

3. Limited model architecture coverage (only LLaMA family in this paper)

4. Citation gap for learnable clipping threshold work, such as "Duanmu H, Yuan Z, Li X, et al. SKVQ: Sliding-window Key and Value Cache Quantization for Large Language Models[J]. arXiv preprint arXiv:2405.06219, 2024."

questions:
1. Regarding the Kronecker decomposition: What is the perplexity performance without decomposition (direct training)?

2. Model generalization:
- How does the method perform on other model architectures (Gemini, Qwen, etc.)?
- Are there any architectural limitations or requirements?
- Would different model architectures require different decomposition strategies?

3. Training stability:
- How sensitive is the method to different initialization strategies?
- What is the impact of calibration data selection?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732354312644

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
The General Response

comment:
We sincerely thank all reviewers for their insightful comments and suggestions. 

We would like to highlight that **FlatQuant is a lightweight yet highly effective method for LLM quantization**. It achieves near-lossless W4A4 quantization results for several models (e.g., < 1% accuracy drop for LLaMA-3-70B) with minimal latency overhead (e.g. < 0.10x speedup loss), which is valuable for deploying W4A4-quantized LLMs in real-world applications.

In response to the reviewers' feedback, we have carefully revised the manuscript (with revisions highlighted in blue). Below is a summary of the main updates:

* **More Theoretical Analysis**. We provide a more comprehensive theoretical analysis of the memory overhead of pre-quantization transformations in FlatQuant, now included in Appendix C.5 (Line 1113~1118). In short, the additional affine transformation are pretty lightweight: for LLaMa-2 7B, they only account for 2.6% FLOPs of the forward computation, and take 0.11MB in size.

* **More Experimental Results**. 

  * Speedup of FlatQuant w/o kernel fusion in Figure 6. FlatQuant is on par with QuaRot in terms of latency even without kernel fusion, thanks to kronecker decomposition.

  * Experiments on Qwen models in Table 11.

  * The flatness visualizations of more linear layers in Figure 16-17.

  * A more detailed discussion of the relationship between flatness and FlatQuant in Appendix C.1 (Line 906~936).

* **Others**. We add more related works about learnable clipping in Appendix A (Line 712~715).
</content>
</thread_object>

</thread>
<thread>
creation date:
1733313620319

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
The General Response (part 2)

comment:
As the rebuttal period comes to a close, we would like to sincerely thank all the reviewers for their valuable comments and the time they dedicated to evaluating our paper. We are encouraged to see that all the reviewers recognized the contribution of our work, particularly in:
 
- introducing optimized decomposed affine transformations to enhance the flatness of weights and activations for effective low-bit LLM quantization [Reviewers m9T2, YdzG, tWUe]; 
- the speedup achieved through Kronecker decomposition and the fusion kernel [Reviewers m9T2, YdzG, RQEg, X7ak, tWUe]; 
- and our achievement of state-of-the-art W4A4 quantization results, demonstrated through extensive experiments [Reviewers m9T2, YdzG, X7ak, tWUe].

As Reviewer RQEg raised the concern about the latency comparison, and Reviewer X7ak raised a concern regarding the speedup without kernel fusion, we have provided thorough discussions. The results show that even without kernel fusion, the additional overhead is comparable to QuaRot/SpinQuant with superior accuracy improvements. With the fusion kernel, we achieve the best practices in W4A4 quantization for LLMs, offering both accuracy improvements and practical speedup. We believe our experiments to be fair and comprehensive. 

As the deadline is approaching, we are sincerely looking forward to the following-up feedback and re-consider your evaluations.

Thanks
Authors of FlatQuant
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper explores a post-training quantization method for LLMs. The main idea is to learn a linear transformation that can be applied to the weights and activations of LLMs, aiming to reduce activation outliers and improve flatness. Additionally, the authors propose a Kronecker-product-based approach to lower computational costs; they also introduce learnable clipping thresholds.

Main strengths:
- Solid results on LLM quantization.

Main weaknesses:
- Reviewers mentioned several related papers in this area (refer to individual review comments). It appears that the proposed method might be overly incremental and lacks sufficient and fair comparisons with these existing approaches.
- Several reviewers raised concerns that the proposed approach could introduce excessive computational overhead.
- The flatness achieved through the method requires more thorough explanation. At present, it seems to rely on experimental observations without adequate theoretical justification.