PV-Tuning: Beyond Straight-Through Estimation for Extreme LLM Compression

Summary

The paper provides a new QAT algorithm for extreme compression of LLMs under various discrete weight representations (including vector quantization, uniform quantization). Traditional optimization techniques rely on straight-through estimation (STE) for optimization in the presence of discrete parameters. This paper provides an alternative to QAT called PV-tuning which resembles the EM algorithm, and performs alternating optimization. When specialized to standard uniform quantization, the algorithm does the following: in the P step, the scale parameter is updated using backprop, and V step performs rounding to the nearest scaled integer. One important aspect of the algorithm is that the V step only updates a few coordinates at a time. This is mainly done to ensure the algorithm doesn't get stuck at a sub-optimal point. Experimental results show that this new optimization algorithm leads to better quality models than existing PTQ techniques such as QuIP, AQLM, GPTQ on diverse models. Experiments on Llama 2 7B also show that this technique improves upon STE. In terms of compute efficiency, the proposed algorithm is 1.5x slower than standard fine-tuning.

Strengths

Improving STE is a very important problem and is key to making progress on extreme compression of LLMs. The current paper aims to tackle this problem. So, I believe the broad direction is interesting to the community. The proposed algorithm has an EM flavor, which is known to work well for estimating latent variables models. I found it quite interesting that the proposed algorithm seems to out-perform STE for 2-bit compression of Llama 7B. That being said, there are certain weaknesses in the empirical evaluation that needs to be addressed.

Weaknesses

- **Clarity** The presentation (especially the algorithmic, technical details) in the paper needs to be significantly improved. Instead of the non-uniform quantization example used in the paper, I would suggest using a more relevant example such as uniform quantization or vector quantization. In its current form, I found it hard to wrap my head around how the proposed algorithm can be used for vector quantization and uniform quantization (I had to look into the appendix to figure this out). This is probably because there is no single framework that can unify all the weight representations. Different representations require different definitions of P, V. - without the clean presentation, it is very hard to really understand how the proposed technique compares with existing QAT algorithms such as STE, at a conceptual level. - **Experimental Evaluation** Better experimental evaluation is needed to understand the real utility of the proposed approach. To be precise, the experiments in Table 2 compare PV tuning with PTQ techniques. But it is well known that QAT outperforms PTQ techniques. So, a better baseline to compare in Table 2 is QAT using STE (or other state of the art QAT techniques). - **Related Work** Uniform quantization, as done by GPTQ and QuIP, has been well studied in the literature. And several QAT techniques that improve upon STE have been developed for uniform quantization (packages such as Pytorch, JAX, TF have really good implementations of QAT for uniform quantization). It's a bit surprising that none of these works are discussed in the paper.

Questions

- It would be great if the authors can provide some explanation for why PV tuning outperforms STE for uniform quantization. - For the case of uniform quantization, it looks like PV tuning does some form of alternating minimization for learning the scale parameters and the quantized weights. In contrast, STE jointly optimizes for both scale parameters and quantized weights. Intuitively, shouldn't joint optimization perform better than alternating optimization? - It looks like PV tuning only updates a few coordinates in each iteration. Doesn't this slow down the convergence of the algorithm and increase the training cost (for each backprop we would be updating very few coordinates)? - If tau is set dynamically (as described in line 304), wouldn't it hurt the performance of optimizers like Adam which maintain coordinate-wise gradient statistics? - It looks like in experiments, STE is applied on top of PV-tuning. But this algorithm is never described in the paper. It would be great if the authors can provide some details about this algorithm?

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

See weaknesses and questions.

Summary

This paper proposes a novel PV-tuning algorithm for extreme compression of LLMs. In each iteration, two steps (P step and V step) are performed, each aimed at reducing the loss function. Different than existing methods that use straight-through estimator (STE), PV-tuning employs a subspace search strategy at the V step to enable large updates and overcome the training stagnation caused by discreteness. Extensive experiments demonstrate that the proposed method achieve superior accuracy on models such as LLaMA and Mistral.

Strengths

1. The proposed PV framework comes with theoretical guarantees. 2. The subspace descent in the V step is a novel approach to overcome the stagnation issue caused by discreteness, and it appears to work better than the popular STE method. 2. PV-tuning is validated through extensive experiments and consistently outperforms all existing methods for 1-bit and 2-bit vector quantization.

Weaknesses

1. PV-tuning is computationally more expensive than existing PTQ methods and requires more GPU resources on larger models. Could you report the runtime of PV-tuning in your expiriments? 2. The authors claim that STE-based optimization is not well justified and leads to poor practical performance. However, I am aware of a paper [1] that provided the theoretical justification for STE approach. - [1] Yin el al. Understanding Straight-Through Estimator in Training Activation Quantized Neural Nets, ICLR 2019.

Questions

1. In my view, the PV algorithm seems similar to alternating minimization. Is there a connection between these two algorithms? 2. Why is the P step an unconstrained minimization problem, as there is a constraint $P(y) \subseteq P(x)$? 3. Is the linearized V step described in section 3.2 equivalent to projected gradient descent? What is the computational cost to compute the projection onto {$x: V(x) \subseteq V(y)$} in the V step ?

Rating

7

Confidence

3

Soundness

4

Presentation

4

Contribution

3

Limitations

N/A

Summary

This paper proposes a new PTQ fine-tuning algorithm called PV tuning. Recent SOTA LLM PTQ works include fine tuning steps to recover the original model on top of the actual quantization step. Fine tuning has been shown to be an effective and relatively cheap to run (vs full QAT) way to improve quantization quality. However, fine tuning for LLM PTQ has not been well explored, in part due to it only being recently introduced and also LLMs being expensive to do things with. My understanding of PV tuning is that it performs alternating optimization on a quantized LLM. Specifically, it alternates between optimizing continuous parameters (codebook values, layernorms, etc) and what is essentially a form of coordinate descent on codebook assignments to weight matrix entries.

Strengths

- PV tuning appears to achieve strong empirical results on a wide variety of models. Notably, PV tuning achieves good performance on Llama 3, which GPTQ-based methods tend to do poorly on. - PV tuning does not seem too costly to run vs existing fine tuning methods and should be compatible with a wide range of "base" quantization methods.

Weaknesses

- This paper is written in a dense and arcane way and is hard to follow. While there are some theorems, they do not appear to be very useful (eg 3.1 just claims that PV tuning converges, but without any statements on how good the final solution is or any bounds on error). Perhaps I am misunderstanding the paper completely, so it would be useful for the authors to give a TLDR version of the actual method. - How large is the subspace for the CD step? Iterating through all possible code assignment changes is exponential in the subspace size, so how do you perform the code assignment optimization in a tractable way? - The empirical results presented are based off of the AQLM quantization method. One main benefit of quantizing LLMs is to achieve speedups in memory bound scenarios. However, the AQLM configuration used for PV tuning is too slow for fast inference (relative to what the hardware *could* achieve) due to the size of the codebook. How well does PV tuning perform with a smaller codebook (eg a 10 bit 8D codebook, which would use 16KB)? The paper also mentioned PV tuning could be applied to QuIP#, which does achieve near memory bandwidth. Do you have any experiments showing this? QuIP# is based off of GPTQ which performs poorly on Llama 3, does PV tuning fix this if applied to QuIP# after GPTQ? - I started reviewing this paper a few weeks ago and lost the first copy I was marking up, but I remember there were some inconsistencies in the result tables with numbers reported in the baselines, and some of the baseline results in 4.1 seemed a bit suspicious. Please check that your numbers are indeed correct. The pareto optimality claim is also questionable, since the "frontier" is just for PV tuning. A better comparison would be PV tuning @ 2 bits vs FP16 @ 3 bits, and it doesn't seem (although I didn't plot this) that PV tuning is actually better at 2 bits than FP16 at 3 bits.

Questions

see above

Rating

5

Confidence

2

Soundness

2

Presentation

1

Contribution

3

Limitations

yes

Summary

The paper focuses on fine-tuning techniques over compressed weights and proposes PV-Tuning, a general framework that improves existing fine-tuning strategies and provides convergence guarantees. Experiments show that PV-Tuning outperforms prior techniques and achieves the first Pareto-optimal quantization for Llama-2 models at extreme compression. The idea is novel and reveals the importance of fine-tuning strategies.

Strengths

The technical contribution of the work seems solid. In particular, the authors conduct extensive experiments to demonstrate the effectiveness of PV-Tuning coupled with the state-of-the-art quantized representations.

Weaknesses

1. The paper introduces the principle and specific implementation of PV-Tuning in detail, but does not compare the differences in principle and specific implementation with the existing STE method and Stochastic Rounding method. It would be beneficial to include such a comparison for clarity. 2. Some notations in the paper may be confusing. In part 3.1, the expression “P(y) ⊇ P(x)” is used but it does not mean that P(x) is a subset of P(y), which is confusing because the ⊇ symbol usually means to include. Also in this part, it would be clear if there are some simple examples.

Questions

1. In the experiment part, besides the PPL and accuracy, is there any evidence that supports the STE is noisier than PV-Tuning? It would be more convincing if the paper could provide some experimental evidence. 2. From Table 1, it seems the naïve Linearized PV without subspace doesn’t achieve pleasing performance while subspace Linearized PV performs well. Does this mean that subspace is the key to performance improvement? Can you combine subspace and STE to conduct some ablation experiments?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Please see the Weaknesses and Questions.

Summary

The paper proposes the use of fine-tuning over highly-compressed models to achieve better model compression. By introducing the PV-tuning algorithm, the authors have achieved superior discrete parameter optimization compared to traditional STE methods. The effectiveness of PV-tuning has been validated by the authors using various base models and datasets.

Strengths

1. The authors have tested the effectiveness of PV-tuning under different experimental settings. The improvements are notably significant compared to the baselines. 2. The authors have found that fine-tuning should be employed to achieve better model compression rather than one-shot quantization. 3. The authors have proposed the PV-tuning method, addressing the issue of small gradient updates in discrete optimization.

Weaknesses

1. I would appreciate it if the authors could clarify the distinction between fine-tuning and one-shot quantization as mentioned in lines 37-43. What exactly does "one-shot" imply? Does it refer to layer-wise quantization? Conversely, what does the authors' claim of "fine-tuning layer-wise" mean? As I understand, quantization typically uses a small calibration set, yet the authors have utilized the entire RedPajama for calibration in section 4.1. Is the "fine-tuning" mentioned by the authors simply calibration on a larger scale? 2. It seems the main problem the authors aim to address (line 45) is the further fine-tuning of already highly-compressed models on a large-scale dataset (RedPajama), rather than the traditional quantization/compression methods applied to full precision models. I agree that this approach might be effective. However, if this is the case, there might be an issue with the experimental setup in 4.1. Firstly, Figure 2 should compare the effects of each method with and without fine-tuning; the current Figure 2 middle and right do not directly facilitate this comparison. Secondly, the authors should not only select quantized representations from different methods for fine-tuning but also apply different fine-tuning strategies, including STE, PV-tuning, GPTQ, etc., to demonstrate the advantages of PV-tuning. 3. In Table 2, PV-tuning indeed shows a significant effect. However, I am uncertain whether this is due to PV-tuning using more data (RedPajama) or because it is inherently more effective. Did the authors use the same calibration data for different methods in this experiment?

Questions

An important question: if we need to perform large-scale fine-tuning to achieve model quantization, why not maintain full precision during this large-scale fine-tuning and only apply low-precision quantization techniques after full precision fine-tuning?

Rating

5

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

Yes

full precision fine-tuning?

> (R3n7c) why not maintain full precision during this large-scale fine-tuning and only apply low-precision quantization techniques after full precision fine-tuning? > While we detail our reasoning in the direct response, we generally agree that it is interesting to test this empirically. To that end, we conducted full-precision fine-tuning and quantized the fine-tuned model as requested. The resulting perplexity scores can be found in Table 2 in the attached PDF (3n7c). Overall, it appears that SoTA LLMs are already well tuned on general data (that we use for calibration); fine-tuning them further lead to negligible improvements on some tasks to the small detriment of others. In turn, PV-Tuning also does not improve the original model quality, but it allows the quantized layers to “re-adjust” to each other’s quantization errors, leading to better accuracy. We thank the reviewer for this interesting suggestion. I'm having trouble understanding why this is a sensible thing to do / how such a setup would work. My understanding was that the goal of PV-tuning was to recover the original model through fine-tuning by tuning the quantized model on a calibration set. Before quantization, the model will always match itself on any dataset. Does this mean the tuning in PV-tuning actually minimizes autoregressive error on a dataset (eg perplexity), or does it minimize the error to the original model on a dataset? If its the former, then isn't PV-tuning essentially doing QAT to a *new* task?

more comments

- Regarding subspace selection, how do you do the code updates in a tractable way after selecting the subspace? Perhaps I am missing something, but this is still a discrete optimization problem after subspace selection, just a much smaller one. Is the "naive solution" that searches over all combinations cheap enough to do here? - Your main rebuttal said there were more experiments here with smaller codebooks, but I'm not seeing any. Did you forget to post part of the response? I'm interested in seeing how PV-tuning would work with a structured codebook, such as one that has sign symmetry. The 8 bit codebook result in the paper performs around the same as QuIP#, so it does appear that the quantizer quality still has an impact after PV-tuning. - Regarding the baselines, IIRC there were some FP16 numbers that didn't match. I don't have time to go back and check, so I will take your word that the tables are correct. - FP16 @ 3 bits just means if you plotted the FP16 perplexity numbers assuming the models were 3/16th the size. This represents what a lossless 3 bit model would achieve. Otherwise, if the 3 bit version of your quantizer just happens to be bad, then it would be very easy to outscale that.

Response to "full precision fine-tuning?"

Your original understanding is correct: PV-tuning aims to minimize the difference between original and compressed model output in terms of KL divergence, not the autoregressive loss. Another reviewer asked us to run an additional ablation experiment in a different setup. The meaning of this experiment requires context from our discussion with (R3n7c): > An important question: if we need to perform large-scale fine-tuning to achieve model quantization, why not maintain full precision during this large-scale fine-tuning and only apply low-precision quantization techniques after full precision fine-tuning? We conduct this experiment as a sanity check, to answer the question: how would the model fare *if we were to maintain full precision* during fine-tuning. In this (different) setup, it makes no sense to minimize KL (as you said), and for this reason, **we only minimize the autoregressive cross-entropy in this side-track experiment**. To the best of our understanding, the reviewer may have suggested the experiment as an extra ablation. We hope this clarifies the reviewer’s question, but we are happy to discuss this further.