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

Authors: vladimir  malinovskii
denis  mazur
ivan  ilin
a i initiative
  kaust
denis  kuznedelev
konstantin  burlachenko
kai  yi
dan  alistarh
peter  richtarik
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 12:32:54

Abstract: There has been significant interest in "extreme" compression of large language models (LLMs), i.e. to 1-2 bits per parameter, which allows such models to be executed efficiently on resource-constrained devices.  
Existing work focused on improved one-shot quantization techniques and weight representations; yet, purely post-training  approaches are reaching diminishing returns in terms of the accuracy-vs-bit-width trade-off. State-of-the-art quantization methods such as QuIP# and AQLM include fine-tuning (part of) the compressed parameters over a limited amount of calibration data; however, such fine-tuning techniques over compressed weights often make exclusive use of straight-through estimators (STE), whose performance is not well-understood in this setting. 
In this work, we question the use of STE for extreme LLM compression, showing that it can be sub-optimal, and perform a systematic study of quantization-aware fine-tuning strategies for LLMs.
We propose PV-Tuning - a representation-agnostic framework that generalizes and improves upon existing fine-tuning strategies, and provides convergence guarantees in restricted cases.
On the practical side, when used for 1-2 bit vector quantization, PV-Tuning outperforms prior techniques for highly-performant models such as Llama and Mistral. 
Using PV-Tuning, we achieve the first Pareto-optimal quantization for Llama-2 family models at  2 bits per parameter.

Main paper: 
 introduction: 
 Recent years have seen the development of ever more capable large language models, attracting immense interest from both researchers and industry. One of the driving factors behind progress in this area is the availability of powerful open LLMs such as Llama [69] , Mistral [34, 35] , or Phi [41] . The main advantage of open LLMs is that they can be run and fine-tuned locally by end users; however, as state-of-the-art LLMs grow larger, they also become harder to run on commodity hardware. For instance, in order to fit the best available Llama-3 model on a consumer GPU, the model would have to be compressed to below 2 bits per parameter 1 .To achieve such "extreme" degrees of compression accurately, researchers have proposed a variety of techniques, which can be roughly categorized into i) better quantized weight representations and ii) better algorithms to learn these representations. The weight representations used for extreme quantization include group quantization [22, 20] , sparse high-precision outliers [17, 32] , incoherence processing of the weights [9, 70] , or additive and residual quantization [21, 72] . In turn, the calibration algorithms also vary between data-free methods [20] , layer-wise calibration [22, 18] , block-wise or global fine-tuning [21, 71] or even quantization-aware training [78, 75] . However, the weight representation and the fine-tuning algorithm are largely orthogonal: most popular quantized representations could be obtained layer-wise in one-shot, fine-tuned layer-wise to a variety of optimization objectives, or even trained entirely from scratch.Surprisingly, there is a clear disparity between the degree of interest shown to accurate one-shot2 quantization versus accurate fine-tuning. Specifically, one-shot quantization is very well-studied, to the extent that, as shown in Figure 2 , improvements in this direction are clearly saturating. At the same time, the impact of fine-tuning strategy is largely unknown: while many recent works use some form of fine-tuning [63, 21, 71] , they typically consider a single fine-tuning regimen based on straight-through estimation (STE) [6, 15] . Thus, given the multitude of representations considered, it is not at all clear whether current fine-tuning strategies are optimal.In this work, we analyze the problem of fine-tuning over highly-compressed weights from the optimization perspective. We begin by analyzing popular fine-tuning strategies for extreme LLM quantization. The key challenge in this context is that the quantized representations may contain both continuous and discrete variables: while continuous parameters, such as learnable scales or codebooks, can be optimized by backpropagation, the discrete parameters (e.g., integer assignments for the weights) cannot. Existing fine-tuning techniques either do not optimize over discrete parameters at all [71, 21] or fine-tune them using heuristics such as STE or stochastic rounding [3] . Unfortunately, these methods are not well-justified for weight quantization from the point of view of optimization theory, and, as we show in Section 3, can provide poor practical performance.We propose an alternative solution: instead of following heuristic gradient estimates, our approach follows the actual gradient of the objective in a small subspace of optimized parameters where it can be meaningfully improved. Following this insight, we formulate the PV-tuning framework for fine-tuning arbitrary quantized representations. We update both discrete and continuous components to minimize a global objective function, such as the KL divergence relative to the original model predictions. Our results show that this strategy leads to significant improvements across weight representations, achieving new state-of-the-art in compression-accuracy trade-offs.The main contributions of our work can be summarized as follows:1. We analyze the problem for training discrete quantized representations for better understanding of the limitations of existing optimization algorithms. We then propose a novel algorithm inspired by compressed gradient methods that addresses these limitations. When compared to straight-through estimation and stochastic rounding, our approach 1) can be shown to converge to a stable solution; and 2) this solution is significantly more accurate in practice. 2. We generalize the proposed algorithm into the PV-Tuning framework 3 , which can minimize a global objective function over a general quantized representation, by optimizing both continuous and discrete parameters via a variant of coordinate descent. 3. We demonstrate that PV-tuning can improve quantized model accuracy for leading existing approaches, including GPTQ and AQLM, on popular LLMs including Llama-2 & 3 and Mistral.Our procedure achieves state-of-the-art accuracy (measured through perplexity) in 1-and 2-bit quantization regimes while using the same amount of calibration data as the original algorithms. Importantly, the PV-tuned models use the same underlying weight representations, and are compatible with existing inference kernels. In terms of accuracy per model size, PV-tuning of vector quantization outperforms all prior techniques in the 1-3 bits/parameter range, and is the first to achieve Pareto-optimal quantization for Llama-2 models at around 2 bits per parameter. 
background: 
 Post-Training LLM Quantization (PTQ). There has been significant interest in PTQ methods [49, 25] that would scale to LLMs. Early work [17, 78, 50] used direct round-to-nearest (RTN) quantization over weight groups of well-chosen size. GPTQ [22] improved upon these results significantly via an accurate one-shot solver for minimizing layer-wise compression errors. Next, AWQ [42] improved upon these results by employing per-channel scaling to reduce the error on important weights while SqueezeLLM [36] implemented non-uniform quantization. QuIP [9] proposed a more accurate weight representation by leveraging incoherence matrices. Another line of works [18, 39] proposes an improved quantized weight representation, which saves a small fraction of outliers in full precision. Other recent works propose augmenting quantized representations with lowrank "adapters" that compensate quantization error [28, 84] . Recently, BiLLM [32] developed residual binarization that stores salient weights in progressively higher bitwidth, quantizing models to nearly 1 bit per parameter at non-catastrophic accuracy loss.Currently, the state-of-the-art methods in terms of accuracy-vs-size are QuIP# [71] and AQLM [21] .Both methods work roughly by mapping weight groups to points on highly-dimensional lattices, which are either chosen to satisfy some optimality properties (for QuIP#) or are learned (for AQLM). Interestingly, AQLM showed that fine-tuning the continuous parameters (codebooks) can improve accuracy significantly relative to pure one-shot compression; a variant of this approach was also adopted by QuIP#. PV-Tuning is compatible with both methods: as we show, it can lead to state-ofthe-art compression results for such representations.Fine-tuning over Quantized Weights. As mentioned above, the two SOTA quantization techniques apply fine-tuning, but only update continuous parameters, such as quantization scales. When optimizing over discrete parameter sets, a standard choice in deep learning is the Straight-Through Estimator (STE) [6, 15, 73] . Prior work on LLM compression proposed to update both continuous and discrete parameters, via STE, both for post-training quantization [78, 63] and for training quantized networks from scratch [32] . However, it was observed early on that STE leads to instability when fine-tuning heavily quantized LLMs [78] . While early results suggest that STE can perform well when training quantized models from scratch [44] , this behavior is yet to be validated for highly-performant multi-billion-parameter models, which are the focus of our work.In summary, the two standard approaches for fine-tuning quantized LLMs are 1) fine-tuning only over the continuous parameters, such as quantization scales, which heavily limits the number of trainable parameters; and 2) optimizing all parameters via the STE, which however is known to be quite noisy especially for extreme quantization. In this context, our work proposes alternative approaches in the post-training compression setting, which lead to state-of-the-art results relative to both options.3 Fine-Tuning Quantized ModelsIn this section, we study the problem of fine-tuning quantized models to minimize a global objective, such as cross-entropy. Section 3.1 formulates this problem from an optimization perspective and introduces our notation. In Section 3.2, we analyze several popular strategies for solving this problem and highlight some of their limitations. To circumvent these limitations, we propose an alternative optimization algorithm in Section 3.3 and discuss implementation details in Section 3.4. 
problem description: 
 Consider the problem of minimizing objective (loss) ϕ, minEQUATIONwhere ϕ : R d → R is a differentiable function bounded from below (e.g., by zero), and R d c ⊂ R d is a set of all possible quantized weights that can be represented with a given quantization method. Without loss of generality 4 , we first analyze the case of scalar nonlinear quantization. In this scenario, c ∈ [d] := {1, 2, . . . , d} (typically c ≪ d), and R d c ⊂ R d is the set of all vectors in R d whose d entries take exactly c distinct values. In other words, the cardinality of the set V (x) := {x 1 , . . . , x d } is equal to c, and we can therefore write R d c := {x ∈ R d : |V (x)| = c}. x∈R d {ϕ(x) : V (x) ⊆ V (y)}. (3) Likewise, M V (y) ∈ R d ≤c and ϕ(M V (y)) ≤ ϕ(M V (y)) ≤ ϕ(y). Evaluating M V amounts to solving difficult discrete optimization problems with a search space of size |V (x)| d ≤ c d (exponential in d).x k+1 = M V (y k ) := arg min x∈R d ϕ(x) : V (x) ⊆ V (y k ) (V step: discrete) 5: end forOur key algorithmic idea, in its simplest form, is to optimize ϕ by alternating the P and V steps, i.e., iteratively applying the M P and M V operators. (We will propose several more practically-useful approximations and variations later; see Sections 3.2-3.3 and also Appendix B.) This resulting method, which we call the PV method, is formalized as Algorithm 1. Our key guarantee for the PV method is formalized in the next result. Theorem 3.1 (Convergence of the PV method). Assume ϕ is bounded below, and letx 0 ∈ R d c . Then (i) y k ∈ R d ≤c and x k ∈ R d ≤c for all k ≥ 0; (ii) ϕ(x k+1 ) ≤ ϕ(y k ) ≤ ϕ(x k ) for all k ≥ 0; and (iii) the sequence {ϕ(x k )} k≥0 converges.The proof can be found in Appendix A.1. Note that we do not claim that the method converges to a minimizer of ϕ; the optimization problem is too difficult for us to be able to guarantee this. However, as we shall see in the numerical results, we nevertheless obtain great empirical performance, especially when coupling the PV approach with some additional algorithmic tricks. This general approach is popular in "shallow" machine learning problems; for instance, if ϕ(x) = ∥x -z∥ 2 is the squared error with respect to some user-specified vector z, then the above algorithm recovers 1-dimensional K-means on the data vector z. Likewise, if ϕ(•) is the log-likelihood, then, depending on the choice of the set R d c , the approach is related to the EM algorithm [16] . In turn, we apply the PV method to obtaining highly-accurate quantized LLMs. Applying the PV method "as is", would be infeasible in practice: computing the P and V mappings requires solving difficult optimization problems especially due to LLM parameter scales. However, both mappings can be approximated. The P step can be reparameterized as an unconstrained optimization problem on the unique values in the weight matrix. Practically it means that the "codebooks" can be optimized using an automated differentiation engine (i.e. PyTorch). However, for many quantized representations, M P (x) can be approximated by one or more steps of GD, directly optimizing ϕ over the set V (x) of its c unique values. The c-dimensional gradient can be computed efficiently by backprop, as described in prior works [63, 71] . On the other hand, the V step (M V (•)) is more difficult to approximate as it involves searching a discrete space of size c d . We dedicate the next two sections to this task.The V mapping (3) can be approximated by solving a discrete least squares problem using an approximation of ϕ(x) around y:EQUATIONwhere L > 0 is a sufficiently large constant. Subsequently, we perform the V step using the simpler convex quadratic function ϕ y instead of the typically more complicated function ϕ:EQUATION≈ M V, ϕy (y)= arg minx∈R d ϕ y (x) : V (x) ⊆ V (y) .Our first lemma shows that we can replace ϕ y by a more convenient function ϕ y measuring the squared distance between x and y + := y -1 L ∇ϕ(y), the latter being the point obtained after taking a single GD step from y with learning rate 1 L , disregarding the constraint: Lemma 3.2. For any y ∈ R d ≤c we have M V, ϕy (y) = M V, ϕy (y), whereEQUATIONThe proof can be found in Appendix A.2. To summarize, the V step of the PV method (Algorithm 1), i.e., x = M V,ϕ (y), can be approximated via the "linearized V step"EQUATIONOur next lemma says that the above approximation is in a certain sense natural reasonable provided that ϕ is L-smooth5 on R d ≤c , i.e., provided thatEQUATION)Lemma 3.3 (Monotonicity). Let y ∈ R d ≤c . If ϕ is L-smooth on R d ≤c , then ϕ (M V,ϕ (y)) ≤ ϕ(x) ≤ ϕ(y), where x is the point obtained from y by the linearized V step (6) . Indeed, the point x obtained via the linearized V step can not have a worse loss than the previous point y. Of course, one hopes that the loss will strictly decrease so that the method makes progress. From a practical perspective, the key advantage of linearized V step is that it can be performed much faster compared to the vanilla V step. The proof of Lemma 3.3 can be found in Appendix A.3.Note that since ϕ y (x) is separable (see (8) ), each entry/weight of x can be optimized independently of others. For scalar quantization, each individual problem can be solved in O(log 2 (c)) time using binary search in sorted version of V (y). For vector quantization, there are specialized optimization procedures for efficiently minimizing the L 2 error (see Appendix D) Key challenge. The main caveat with linearized V step is that it may be impossible to make small gradient-based updates to low-bitwidth discrete weights. More specifically, in (6) , one must update the discrete assignments to approximate y k -1 L ∇ϕ(y k ). However, for low-bit weights, the desired update 1 L ∇ϕ(y k ) can be smaller than the lowest possible increment to obtain a quantized vector. As a result, the optimal solution to (6) is often y k itself. In such a situation, the algorithm will get stuck on y k , which is undesirable. This problem is especially pronounced in deep LLMs, where L can be very large, or, from a practitioner's point of view, where one needs a small learning rate. In practice, as we explore in Section 4.2, the lowest learning rate where the algorithm makes any updates at all is already too large for optimization, leading to divergence.Require: initial parameters x 0 ∈ R d c , objective function ϕ : R d → R , subspace size τ ∈ [d] 1: for k = 0, . . . , K -1 do 2:▷ P step: update V (x) by backprop 3:EQUATION7: Many popular strategies for discrete fine-tuning can be seen as attempts to reconcile coarse lowprecision weights with the need to make small updates. These include straight-through estimation, stochastic rounding, or adding regularizers that push the solution to (6) away from y k . We review straight-through estimation in Appendix E.1 and stochastic rounding in Appendix E.2.x k+1 = arg minx ϕ y k ,S k (x) :V (x)⊆V (y k ) 8: end forHere we ask the following question: Can we modify the PV method so as to force the V step to make a larger update? In other words, we need an optimization algorithm that updates quantized weights either by a sufficiently large increment, or not at all.A natural example of such an algorithm is coordinate descent (CD) [43, 58] , or more generally, subspace descent [26, 38] . Instead of updating all parameters by a small margin, CD in each iteration chooses a single parameter, and makes a large update instead. This strategy can be generalized to updating more parameters at the same time, which leads to subspace descent methods. 6 The parameters to be updated can be chosen either greedily, (e.g., several i ∈ [d] with the largest magnitude of the partial derivative |∇ i ϕ(•)|), or at random, or through a variety of other means.Let S k ⊂ [d] be the set of parameters/weights/coordinates we wish to update at iteration k. We choose |S k | = τ ≪ d. Let Z k : R d → R d be the linear mapping defined as follows:Z k (x) i = x i if i ∈ S k and Z k (x) i = 0 if i / ∈ S k .We now formulate the linearized subspace V step:EQUATIONand L S k > 0 is a smoothness parameter of ϕ associated with the subspace spanned by the parameters belonging to S k . This detail is important because L S k ≪ L when τ ≪ d. When estimating Lipschitz constants for real LLMs, we found that it is lower by at least one order of magnitude, making it possible to train with sufficiently large step sizes (see details and L S k estimates in Appendix F).Note that, necessarily, x + i = y i for i / ∈ S k . The remaining τ entries of x + can be identified exactly by searching a discrete space of size |V (y)| τ , which is feasible if c = O(1) and τ = O(1), for example.In practice, it means that the algorithm can apply large updates to quantized LLM weights, with the caveat that should only update a fraction of them at a time. This allows us to perform the linearized V step with sufficiently large "learning rate" to make non-trivial (i.e., x k+1 ̸ = y k ) improvements to quantized weights even without straight-through estimation or stochastic rounding.We formulate the full procedure in Algorithm 2. The algorithm performs the P step by directly optimizing V (x) (i.e., codebooks) by backprop as described in Section 3.1. For the V step, the algorithm greedily chooses a subset of τ quantized weights for update, then updates them using Eq. (8) . The arg top τ operator finds τ indices with the largest absolute gradient values and builds a subspace of R d ≤c where only these values can be changed, and the rest must be equal to y k .To speed up convergence, we use adaptive learning rates for both P and V steps. In Eq. 8, we replace ∇ϕ(y) with a single Adam [37] update, as depicted in Algorithm 3. In preliminary experiments, we found that this results in a significant convergence speedup. When choosing the subspace S k , we select weights based not on |∇ i ϕ(y)|, but on the magnitude of Adam update for that weight. For simplicity, we greedily choose the τ weights with the largest update norm within each weight matrix.This could be further improved through better techniques for choosing S k explored in Appendix Q. We also found that, despite the fact that PV-tuning by itself outperforms straight-through estimation, we could achieve slightly better accuracy by combining PV-tuning with straight-through estimation.We explore this in more detail in Section 4.2).We describe our approach for preparing the calibration data in Appendix G. We found that the preprocessing used in several recent PTQ works introduce a small bias when sampling the calibration data, leading to somewhat worse fine-tuning accuracy. For fairness, we always compare representations (Section 4.1) and algorithms (Section 4.2) using the same pre-processing.Fine-tuning efficiency. The most compute-intensive part of PV tuning is computing the gradients ∇ϕ(•), which is done through repeated forward and backward passes on an LLM. To reduce the number of gradient accumulations, we reuse gradients for P and V steps within one iteration. We use mixed precision, gradient checkpointing and batch accumulation to train more efficiently; for larger LLMs such as LLAMA 3 70B we also use sharding and optimizer offloading (see Appendix H).Our code can train 7B LLMs on a single GPU, while larger ones (e.g. 70B) fit into a single machine with 8×A100. In terms of wall-clock time, PV-tuning takes up to 1.5× longer than the fine-tuning procedure of [71] and requires additional memory in order to hold ∇ϕ(x).4 Experiments 
evaluating quantized representations with finetuning: 
 Before evaluating PV-tuning, we need to choose the quantized representation to be fine-tuned. We therefore compare popular weight representations from recent works on LLM quantization (see Section 2). To better isolate the effect of the weight representation, we evaluate them in three configurations: i) when quantizing a single LLM layer, in terms of MSE, ii) full model quantization in terms of perplexity without finetuning and iii) with finetuning.We compare several recently proposed quantized representations (see details in Appendix J):1. GPTQ: scalar uniform quantization with channel-wise and block-wise scales [22] ,2. SpQR: an extension of block-wise GPTQ with learned sparse outliers [18] ,3. VQ: basic vector quantization with a single codebook [72] with multi-step training. We run all three experiments on LLAMA 2 7B model [69] , calibrating on the RedPajama [13] dataset that best approximates the original pre-training data. When evaluating single layer errors, we report the L2 error in attention query projection outputs of a fixed transformer block, with other blocks exhibiting similar behavior. For full model evaluation, we report quantized model perplexity on WikiText-2 [45] dataset. We use the same data splits and preprocessing as in most recent PTQ works [22, 42, 18, 70, 21, 71] , including the biased preprocessing step that we mentioned in Section 3.4. For fine-tuning, we train continuous parameters only, using the approach from [71] . To compare these diverse representations, we evaluate their quantization errors as a function of average number of bits per parameter. To get a diverse set of bits per parameter, we vary the hyperparameters such as wbits, block size, codebook and group size for vector quantization and the rate of outliers.Figure 2 summarizes our findings. Overall, vector quantization methods (VQ, QuIP# and AQLM) outperform their scalar counterparts. Outliers and low-rank compensation both reduce error, but this improvement comes at the cost of extra bits per parameter. Interestingly, the improvement from outliers is significantly smaller when both methods have access to fine-tuning. Likewise, the improvement from using low-rank adapters also diminishes when comparing fine-tuned models, to a point where it no longer justifies the increase in model size. We provide a more detailed breakdown of results and hyperparameter configurations in Appendix J.Our main takeaway is that for sub 2 bits per parameter, the vector quantization (VQ) representation can achieve near-optimal quantization accuracy, whether or not it uses outliers, LoRC or incoherence processing. Naturally, this does not reduce the value of prior works since they were designed for different scenarios, typically with a higher number of bits per parameter.Next, we compare different fine-tuning strategies and ablate our PV-tuning protocol. We design our protocol to be representation-agnostic, i.e. compatible with different quantized representations. To showcase this, we pick three methods from the previous section: GPTQ, VQ and AQLM.These methods differ not only in their weight representations, but also in how they search for the optimal codes. Namely, GPTQ can scale the target weight and round it to nearest 2-bit integer. In turn, VQ quantizes weights as a group and must find the nearest vector from its codebook, and AQLM uses a multi-step beam search procedure to choose the best combination of codes from both codebooks. Our PV-Tuning implementation uses these search algorithms during the subspace linearized V step (find_nearest in Alg. 3). We describe the full PV configuration for each method in Appendix K.We compare PV tuning against several popular fine-tuning regimens found in the literature. Our first baseline is fine-tuning only continuous parameters, e.g., codebooks or input/output embeddings [71, 74] . The second baseline is training with Straight Through Estimation (STE) [75, 77] . We also test stochastic rounding as described in Appendix E.2. Finally, we evaluate PV tuning combined with STE, but otherwise the same configuration. We set the subspace size τ equal to the number of weights such that the update satisfies ∥x k+1 -x k ∥/∥x k ∥ ≤ 0.01, also known as known as trust ratio [81] .The results in Table 1 show that PV-Tuning consistently finds better quantized models, with STE coming consistently second. We explore this further by combining subspace updates with STE, which leads to slightly better perplexity and accuracy in most (but not all) setups. PV-Tuning over QuIP# In addition to these three configurations, we also apply PV-tuning to QuIP# [71] -a modification of vector quantization that applies Randomized Hadamard Transform (RHT) before quantization and uses fixed lattices instead of learned codebooks. We experiment with Llama-2 7B model quantized with QuIP# to 2 bits per weight and found that it is possible to significantly improve the model through PV-Tuning. For instance, PV-tuning improves WikiText-2 perplexity from 6.19 (QuIP# with built-in continuous fine-tuning) to 5.71 (PV-Tuning + STE). Since original 16-bit model has a perplexity of 5.13, this corresponds to almost halving the quantization error in terms of perplexity. We report additional details for QuIP# with PV-Tuning and full evaluation results in Appendix L and include it to Table 2 as "QuIP#+PV".On the choice of hyperparameters for 1-bit vector quantization. There are several possible hyperparameter configurations for vector quantization (VQ) that fall into 1-1.1 bit range. One can either use larger codebooks for longer groups (vectors), or smaller codebooks for shorter groups accordingly. In our main evaluations, we quantized vectors of 16 consecutive weights with 14-16 bit codebooks to fit into the desired bitwidth. However, we later found that it is more advantageous to choose smaller groups as well as codebooks. We found that 8-bit code per 8 weights outperforms 14-bit code per 16 weights despite having near-identical bitwidth (due to smaller codebooks). We report this configuration as "PV (gs8)" in Table 2 and provide additional experiments in Appendix M.Finally, we evaluate the resulting PV algorithm with a vector quantization backbone and KL objective on a range of popular LLM models. For this section, our goal is to evaluate our approach holistically for different models and target bit-widths, comparing against the best known baselines in common settings. To that end, we evaluate on LLAMA 2 & 3 [69] , MISTRAL 7B [34] and PHI-3 Mini-4k-Instruct [1] at 1-2.5 bits per parameter (averaged over all transformer layers).We report perplexity on WikiText-2 [45] and C4 [54] validation sets, zero-shot accuracy on Wino-Grande [60] , PiQA [67] , HellaSwag [83] , ARC-easy and ARC-challenge [12] via the LM Eval Harness [24] . We follow the exact evaluation setup from GPTQ [22] . We compare against QuIP [70] , BiLLM [32] , PB-LLM [62] , DB-LLM [10] , AQLM [21] , OneBit [77] , QuIP# [71] , the latter three using fine-tuning. For LLAMA 3, we use baselines from [33] and re-evaluate perplexity in our setup.Table 2 : Quantized model perplexity on WikiText-2↓ [45] & C4↓ [54] and the Average↑ accuracy on 5 zero-shot tasks [24] for various models and bitwidths. Arrows ↑ / ↓ mean higher / lower is better. OneBit, all of which use fine-tuning. The improvements on LLAMA 3 are also remarkable as this model is notoriously hard to compress [33] . We report additional evaluations in Appendix N.A key practical question concerns obtaining optimal quality for the target model size, where a smaller model compressed to 3-4 bits often dominates a larger model compressed to 1-bit. The best known Pareto-optimal bit-width for Llama 2 is 2.5 [21] : compressing a larger model to less than 2.5 bits per weight is inferior to a smaller model quantized to the same total number of bytes. From this perspective, PV-tuning pushes the Pareto-optimal frontier for LLAMA 2 to 2.0 bits. This is easiest to see in Table 12 : a 2-bit 13B model outperforms any 7B quantization and is comparable with the 16-bit 7B model. The same holds for the 2-bit 70B model.Fine-tuning efficiency. One limitation of our algorithm is that it requires more compute and memory during the fine-tuning procedure. The 7B models can be fine-tuned on a single GPU, our 70B runs require a server with 8×A100 or rely on RAM offloading. PV-Tuning shares this drawback with prior methods based on STE [21, 71] , as both methods need gradients w.r.t. dequantized weights. Our longest training run took 2 days on 8 GPUs to outperform all baselines and 8 days to fully converge.Inference speed. PV-Tuning does not change the underlying compressed representation, allowing us to reuse existing high-performance inference kernels. Specifically, VQ+PV can reuse efficient kernels from [21, 71] , while GPTQ+PV can use ExLlamaV2 kernels [14] . We report these inference speed evaluations in Appendix O. From a practitioner's point of view, PV-Tuning can significantly improve the accuracy of extreme (1-2 bit) quantized models, making it possible to deploy large LLMs on resource-constrained devices. As a proof of concept, we developed specialized inference engines for running vector-quantized models with PV-Tuning on mobile devices 7 or in the browser 8 . 
conclusions: 
 Limitations. We focused our effort on evaluating PV-Tuning with multiple setups and models, but spent relatively little effort tuning our algorithm for each specific setup. For instance, we always use constant learning rate and τ with no schedule, and always train on the same data. While this shows robustness of PV-Tuning, it also means that our results may be improved with better hyperparameters. For instance, Appendix M shows how PV-Tuning with 1-bit vector quantization can be improved by choosing smaller vector sizes, while Appendix L suggests that PV-Tuning can dramatically improve models quantized with QuIP# and may similarly be applied to other quantized representations. Furthermore, the algorithm could achieve better accuracy by simply training longer and on more data.Future work. This work opens several new research directions. The first is about how to choose S k : while we found that a greedy strategy works in practice, there may be fundamentally better ways. Another direction is applying PV-Tuning to other quantization niches: our evaluation focuses on extreme weight-only quantization, but the proposed algorithm can be used in weight + activation setting or KV cache quantization. Overall, PV-Tuning shows how an insight from optimization theory can improve LLM quantization and we are excited to see how this develops in future research. Part (i): First, by assumption, we know that and because y = x k satisfies the constraint P (y) ⊒ P (x k ), we conclude that ϕ(y k ) ≤ ϕ(x k ). Further, since x k+1 = arg min 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type computing workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: We report the hardware setting, calibration setting, time, and memory requirements in Section 4, which is sufficient for practitioners to reproduce our results. We omit some details, e.g. which runs were restarted due to unrelated server infrastructure issues.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute worker CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required mocomputingute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
 
 Appendix: x 0 ∈ R d ≤c . Assume that x k ∈ R d ≤c for some k ≥ 0. Since P (y k ) ⊒ P (x k ), this implies that y k ∈ R d ≤c . Next, since V (x k+1 ) ⊆ V (y k ),x∈R d {ϕ(x) : V (x) ⊆ V (y k )} and because x = y k satisfies the constraint V (y) ⊆ V (y k ), we conclude that ϕ(x k+1 ) ≤ ϕ(y k ). In summary, ϕ(x k+1 ) ≤ ϕ(y k ) ≤ ϕ(x k ).Part (iii): In view of part (ii), the sequence {ϕ(x k )} ∞ k=0 is non-increasing. By assumption, it is bounded below. Hence, it converges to its infimum:EQUATION= ϕ M V, ϕy (y)EQUATION= minEQUATION= minEQUATION) Since y ∈ R d ≤c , any x ∈ R d satisfying V (x) ⊆ V (y) must also satisfy x ∈ R d ≤c .So, in view of L-smoothness of ϕ on R d ≤c , we can bound the last expression in (9) from below via ϕ(x)(9)+(7)≥ minx∈R d {ϕ(x) : V (x) ⊆ V (y)} (3) = ϕ (M V,ϕ (y)) .Finally, since x = y satisfies the constraint V (x) ⊆ V (y), we can upper bound the same expression viaEQUATION≤ ϕ(y) + ⟨∇ϕ(y), y -y⟩ + L 2 ∥y -y∥ 2 = ϕ(y).We now introduce the pseudocode of an approximate PV meta-algorithm; the idea is to replace the P and V steps with some approximate computations to be defined later. y k ≈ M P,ϕ (x k ) 4:x k+1 ≈ M V,ϕ (y k ) (for example, we can use the method from Section B.1 or the method from Section B.2) 5: end for Next, we describe two new approximations of the V step.We now describe an algorithm computing an approximation to M V,ϕ (y):1. Start with some y ∈ R d c and choose sufficiently large L > 0, number of iterations T 2. Set z 0 = y 3. For t = 0, . . . , T -1 iterate:(i) Define ϕ z t (•) := • -z t -1 L ∇ϕ(z t ) 2 (ii) Set z t+1 = M V, ϕ z t (z t ) 4. Output: z TThe method is constructed so that z T ≈ M V,ϕ (y). If use this subroutine with T = 1 to approximate the V step in the PV method, we recover what we earlier called the linearized PV method. Choosing sufficiently large T ≥ 2 may be advantageous.We now describe a different algorithm for computing an approximation to M V,ϕ (y):1. Start with some y ∈ R d c and choose sufficiently large L > 0, number of iterations T 2. Choose a suitable decreasing sequence of positive scalars {α t }, with α 0 = 1 and lim t→∞ α t = 0 3. Set z 0 = y 4. For t = 0, . . . , T -1 iterate:(i) Define ϕ z t (•) := • -z t -1 L ∇ϕ(z t ) 2 (ii) Set z t+1 = (1 -α t ) M V, ϕ z t (z t ) + α t z 0 5. Output: z TThe method is constructed so that z T ≈ M V,ϕ (y). This approach is based on Halpern acceleration of fixed point methods [51] , and as such, may be sometimes effective.In Section 3.1, we define R d ≤c as a set of vectors with at most c unique items. This translates to the idea of k-means quantization, a scalar nonlinear quantization where each weight is rounded to one of c centroids found by clustering the weights. Below, we show how this can be generalized to linear quantization, vector quantization, additive quantization, and others.Linear quantization is the most basic and widely used type of quantization where weights are stored as integers, possibly multiplied by a scale and added to a zero point. The simplest way to account for this quantization type is to declare that weight values are integers up to c: V (x) = (0, 1, 2, ..., c -1). After that, one can treat scales / zero points as an extra non-quantized parameter, similar to biases or layer-normalized scales. This extra parameter interacts with weights by multiplication or addition, and hence it can be updated by backprop, similarly to other non-quantized weights. Equivalently, once can declare that V (x) = (0, s, 2s, ..., s • (c -1)) for arbitrary s ∈ R. Both options lead to equivalent fine-tuning algorithms where the V step does not change and the P step has an additional condition.Next, let us discuss vector quantization. Consider a quantization that splits x into 2-dimensional groups (non-intersecting pairs of adjacent weights) and encodes these weights as one of c 2dimensional codes that form its codebook. This can be viewed as two sets of weights (odd and even) quantized with scalar quantization, except that values in the two sets have the same partitioning P (x). In other words, if two values in the odd half belong to the same partition, the corresponding values in the other half also belong to the same partition, though the values themselves can be different. Alternatively, one can simply write down a version of R d ≤c , where V (•) is a set of 2-dimensional vectors, not scalars. Likewise, higher-dimensional vector quantization translates to higher-dimensional items in V (•).Both the P and V steps for vector-quantized weights follow the same general principle: P-step can be approximated by backprop with slightly more trainable parameters. In turn, the V step can be done by trying all values in V(x) and selecting the one with the lowest ϕ(•). A more compute-efficient version of the V step for this case is described in Appendix D.RVQ and Additive Quantization can be treated as learning two separate sets of vector-quantized parameters. However, a more efficient way would be to run the V step to find the best combination of codes via beam search [5] .Quantization with sparse outliers [17, 18, 42] can be seen as learning two distinct matrices with different definitions of R d c : one is quantized and the other is sparse (for outliers). Similarly, quantized weights with low-rank adapters (e.g. [28] ) can treat the adapter as an additional non-quantized parameter for the P step. This makes PV-tuning potentially extensible to neural network pruning.As we describe in Section 3.2, the linearized V step minimizes the squared error between quantized weights and the updated "target" weights. Here, we explain how one can compute and minimize the squared error efficiently in practice. To simplify the notation for this section, we define the objective as ∥x -B∥ 2 where B is the target vector set by the linearized V step. Following the definition of squared L2 norm, this objective can be re-written as follows:EQUATION)Consider the first term:∥x∥ 2 = d i=1 x 2 i . Since x ∈ R d ≤c ,this term is a sum of at most c unique terms. Abusing your notation, this can be rewritten as∥x∥ 2 = c i V i (x) 2 • |P i (x)|, where V i (x) is i-th unique element in x and |P i (x)| is the number of such elements.The second term is also a sum of c unique values:-2 • ⟨x, B⟩ = -2 d i=1 x i B i = -2 c i=1   V i (x) • i∈Pi(x) B i   .The third term does not depend on x.If you know the objective for some given x, you can efficiently compute ϕ for all neighboring x ∈ N 1 (x) where you only change one index. For the sake of formality, let us define the set of such neighboring x as follows:EQUATIONConsider one x ∈ N 1 (x) where only k-th value changed (i.e.EQUATIONNote that, for any ∀x ∈ N 1 (x), there are c 2 possible values for 4. for t = 0, . . . : 5. for x ∈ N 1 (x t ):x2 k -x 2 k and another c 2 unique values for 2 • (x k -x k ) • B k ,6. find k : xk ̸ = x t k (there's only one such k) 7. compute ϕ(x) = ϕ(x t ) + x2 k -x 2 k -2 • (x k -x k ) • B k 8.x t+1 := arg minx∈N1(x t ) ϕ(x) (minimum from array of pre-computed values)In practice, this can be extended from greedy (local) search to semi-greedy beam search. These practical algorithms are described in AQ, LSQ, and LSQ++. Algorithms for ∥A(x -B)∥ 2 are explained in AQLM and probably other works.In this section, we overview possible solution to the general problem of training / fine-tuning neural networks with quantized weights. We focus on strategies that train by gradient descent with additional measures to deal with coarse-grained weights.In principle, there are also gradient-free methods for quantized training, such as Evolution Strategies [56] , Bayesian Optimization [48] and others. However, these gradient-free methods have so far not gained popularity for large language model quantization. Adapting these methods to the scale and dimensionality of LLMs would likely require extra effort. Thus, we leave these methods outside the scope of our work and focus on gradient-based optimization.Reminder: gradient-based training of quantized models. We describe the general framework for training quantized weights in Sections 3.1 and 3.2. To summarize, the training algorithm computes the gradient w.r.t. de-quantized weights as though they were continuous, then uses these gradients to update continuous (P step) and discrete (V step) parameters. The core problem with this approach is that, when discrete parameters are very coarse (e.g. low-bit quantization), gradient updates are no longer large enough to make any changes and are lost to the "rounding error". We review two strategies for circumventing this problem: straight-through estimation and stochastic rounding.Straight-through gradient estimation is a technique for training neural networks with discrete components that ignores these discrete components during backpropagation. Its usage goes back to the Perceptron introduced by Rosenblatt [59] . There, the artificial neuron uses a step function as activation, but the training procedure treats this function as though it was identity. Subsequent works introduce variations of this idea, extend it to multi-layer networks [31] , discuss its convergence properties [76, 23, 80] . Other works use straight-through estimation or similar techniques to training neural network with quantized weights [30, 64, 65] .STE for LLM quantization. As we discuss in Section 2, straight-through estimation introduces an auxiliary non-quantized weight tensor that is updated using the gradients ∇ϕ(y) w.r.t. quantized weights. The quantized weights are then updated to best approximate this auxiliary buffer, usually in terms of L2 error. As a result, if an update to y -1 L ∇ϕ(y) is not large enough to change the parameter, it is still accumulated in a straight-through "buffer". Eventually, the cumulative effect of several such updates will be large enough that y k will no longer be the solution to Equation ( 6). This strategy prevents the algorithm from stalling, but it does so at the cost of convergence guarantees [80] . When applied to extreme LLM quantization (Section 4.2, straight-through estimation initially improves y k , but then stops improving and oscillates. We also tried several a variant of straight-through estimation [65] that introduce stochasticity to forward pass. When applied to extreme LLM quantization, this variant did not diverge like naive STE, but trained much slower and did not reach the same optimum as "deterministic" STE. We attribute this to the fact that adding noise during training can slow down convergence, which also applies to stochastic rounding (Appendix E.2).Stochastic (or probabilistic) rounding [82, 87, 2, 29] is one of the techniques that can circumvent stalling when training low-precision weights. To recall, the linearized V step (6) can be seen as rounding y + := y -1 L ∇ϕ(y) to the nearest quantized weight in R d c , which often happens to be y itself. To circumvent the problem of rounding back to y, one can instead round stochastically, to one of the two adjacent values that y + falls between. Let's denote these two adjacent values x l and x r for left and right. The probability of rounding is inversely proportional to the rounding error (distance), or, in terms of the objective,p(round to x l ) = ϕ(x l ) -1/2 ϕ(x l ) -1/2 + ϕ(x r ) -1/2 .This way,E p(round to x) x = y - 1 L ∇ϕ(y).The main drawback of stochastic rounding is t introduces noise, it changes the underlying optimization problem. Intuitively, if the optimal x ⋆ is adjacent to a significantly worse solution, the method may oscillate between rounding to either side. This rounding noise increases further as we consider lower quantization width. In Section 4.2 we exploit his phenomenon for real-world LLMs and find that stochastic rounding converges find significantly worse solutions, presumably because at every step, some portion of LLM weights will be rounded poorly. On top of that, when used for vector quantization, stochastic rounding is either intractable or biased.Stochastic rounding for vector quantization. To recall, stochastic rounding for non-vector quantization needs to find two quantized values: the nearest neighbor above the solution, and the nearest neighbor below it. It will then round to either of the two values inversely proportionally to their rounding errors.However, this intuition no longer works if you consider more complex quantization schemes such as vector quantization, additive quantization, quantized low-rank adapters, and others. In vector quantization, a group of weights is encoded jointly as one vector from a fixed set (usually called codebook or lattice). For simplicity, let us consider the case where the weight group size equals 2, i.e. weights are quantized in pairs.For a pair of two eights, we can no longer rely on the fact that they have one neighbor from above and one from below. Instead, they may have any number of adjacent "clusters" they can be rounded to. Intuitively, a pair of weights is a point in 2-dimensional that can have neighbors from left, right, top, bottom, and any diagonals. Formally, to determine a full list of neighbors, we can run Delaunay triangulation on all vectors from the codebook (or lattice) plus the target vector that needs to be rounded, then find all points that share a triangle with the target vector.Unfortunately, this procedure can be very expensive, especially for higher-dimensional vectors. A popular practical approximation to stochastic rounding for vector quantizations is to find K (e.g. 2) nearest vectors from the codebook, then use the probability formula from scalar stochastic rounding:p(round to x i ) = ϕ(x i ) -1/2 /( K j ϕ(x j ) -1/2 )However, unlike the original stochastic rounding, this approximation does not guarantee that E p(round to xi)EQUATION)For a (counter)example, if there is a high density of codes on one side of the target, all K (e.g. 2) nearest neighbors will be on the same side. As a result, this approximate stochastic rounding is biased and further changes the optimization result.Stochastic rounding with temperature. When used for LLM quantization, the main problem with stochastic rounding is that it introduces noise to the training procedure. This is important because modern LLMs [69, 68, 1] typically train without dropout or similar noise layers. This is because the dataset is huge and the training suffers not from overfitting, but from not fitting the data enough.Training with stochastic rounding makes optimization inherently noisy as if using dropout, making it harder to train. What is worse, low-bitwidth models produce more noise than high-bitwidth due to larger intervals between quantized values. This additional noise makes it difficult for the model to fit the training data tightly. For extreme 1-bit training, we often observed that the training objective (cross-entropy) would increase instead of decreasing due to sheer amount of rounding noise.To combat this issue, we introduce stochastic rounding with temperature τ (hyperparameter):p(round to x l ) = ϕ(x l ) -1/(2τ ) ϕ(x l ) -1/(2τ ) + ϕ(x r ) -1/(2τ ) .Setting τ < 1 results the algorithm keeping the original weights more often, which reduces the training noise at the cost of slower updates. In Table 1 (Section 4.2), we try τ ∈ {1, 0.5, 0.1, 0.01} and choose the best result for each setup that involves stochastic rounding. It can be shown to converge as long as τ is annealed, but it changes the underlying optimization problem similarly to dropout. In principle, it is possible to gradually anneal τ during training to make the algorithm unbiased in the limit.Finally, we can compare the two above strategies and our proposed approach from Section 3.3.From the optimization perspective, the most popular variant of straight-through estimation is known to not converge to a stable solution. While, in practice, STE can still significantly improve model quality (see Table 1 ), it is still a heuristic. In turn, stochastic rounding can be seen as an additional source of noise to stochastic gradient descent which can converge to a stable solution when τ is gradually reduced to 0. In contrast, subspace PV does not introduce noise or instability and does not require annealing.The most notable difference of PV-Tuning from STE is that the former tries to update all weights on every step, whereas our algorithm only updates a chosen subset. We give an explanation on why updating all weights is problematic at the end of Section 3.2.From the efficiency perspective, training with straight-through estimation requires storing an additional set of buffers on device memory to accumulate weight updates. Unlike the quantized weights, these buffers need to be stored in higher precision (half or full) to accumulate smaller updates that would be lost on quantized buffers. As a result, straight-through estimation requires additional memory for fine-tuning. Stochastic rounding and subspace PV (w/o STE) do not need these buffers, but they still need to accumulate high precision gradients w.r.t. de-quantized weights and store optimizer statistics for those weights in higher precision. To summarize, all methods require significantly more memory than naive (P-only) fine-tuning, but straight-through estimation has higher memory overhead.As for the computational overhead, both STE, stochastic rounding and subspace PV introduce additional computations and therefore increase step complexity. Of the three alternatives, the subspace PV algorithm is slightly faster since it only runs the discrete optimization on a small portion (subspace) of model weights per step, while stochastic rounding has somewhat higher overhead due to the complicated rounding procedure, especially for vector quantization. However, this overhead is small in practice: most of the LLM training time is spent accumulating the gradients on a large training batch, which is not affected by any of the three algorithms. In principle, it should be possible to reduce the compute / memory overhead both with technical improvements and better optimization algorithms, but we leave this investigation to future work.The classical definition of L-smoothness for differentiable function f : R d → R represented by requirement∥∇f (x) -∇f (y)∥ ≤ L∥x -y∥, ∀x, y ∈ R d .If function f (x) is twice continuously differentiable, then it is easy to show that function f is Lsmooth if and only if ∥∇ 2 f (x)∥ ≤ L. If striving to find the minimum value of L then via following the definition, one has to select L as L = maxx∈R d ∥∇ 2 f (x)∥ .If the domain of function f (x) is restricted to any subset S ⊆ R d , then the global L smooth constant can only decrease for a new function. It can be observed from the fact that ∀S ⊆ R d we have∥∇ 2 f (x)∥ = max v∈R d \0 ∇ 2 f (x) • v /∥v∥ ≥ max v∈S\0 ∇ 2 f (x) • v /∥v∥ , ∀x ∈ R d .Sparse sub-spaces satisfy this requirement; therefore, this theoretical observation is valid in this circumstance.Another observation is that the definition of L smooth constant has a global notion. However, for practical purposes, for the first-order training algorithm what matters is L-smoothness constant for the function f (x) with the domain restricted to the trajectory of iterates only. Unfortunately, the training process iterates follow the prior unknown path in R d space.Below we demonstrate two approximate schemas for evaluating L-smoothness constant for a function with a domain (artificially) restricted to the trajectory induced by iterates generated by Gradient Descent (GD) for functions f (x) with different subspace sizes and in general with different optimization trajectories, but with the same start iterate (model). We have performed 10 iterations of GD for training auto-regressive Llama-like LLama-160M [46] and TinyLlama-1.1B [85] models using automatic tokenizer from these models. We trained all linear layers in these models. The used step size for GD is 10 -4 .Schema I: Estimating L along the trajectory of iterates without capturing local curvature.After running GD for 10 iterations the L smooth constant has been estimated along trajectory s = {x 1 , x 2 , . . . , x 10 } with approximated asL := max xi,xj ∈z,xi̸ =xi ∥∇f (x i ) -∇f (x j )∥ ∥x i -x j ∥ .Results are presented in Tables 3, 4 . From them, we can see that L estimate along the trajectory of iterates have the same property as global L, namely during restricting subspace of training variables the L-the smooth constant is non-increasing, and in practice substantially decreasing. This schema exploits available information on gradient oracles in iterates in s and iterates s itself. This schema represents an estimation of upper bound L on the true value of L. Schema II: Estimating L along the sequence of iterates with capturing local curvature.The previous schema used a fixed sequence of iterates s = {x 1 , x 2 , . . . , x 10 } essentially estimate the L-smoothness constant along the piece-wise linear interpolated path along s. In the next schema, we approximate L-smoothness constant asL = max xi∈s ∥∇ 2 f (x i )∥ .Therefore this schema exploits a sequence of points s and selects the maximum in absolute values eigenvalue for all matrices ∇ 2 f (x i ). Computing any spectral information for a matrix with big dimensions can be challenging.The approximate schema that we have used to compute ∥∇ 2 f (x)∥ leverages several observable facts. First, ∥∇ 2 f (x)∥ = max(|λ i (∇ 2 f (x))|), where λ i is i -th eigenvalue for ∇ 2 f (x). Second, to identify the maximum eigenvalue in the absolute value we can use the normalized Power Iteration algorithm [47] , which requires execution only the hessian-vector product. Third, we can use Taylor expansion and forward difference approximation for ∇f (x + r):∇f (x + r) -∇f (x) = ∇ 2 f (x) • r + O ∥r∥ 2The K approximate hessian-vector product can be accomplished with K + 1 gradient oracle calls.For the experiment, we run Power Iteration for a prior known number of iterations equal to 10. In fact Power Iteration does not converge in case of degeneracy such as a situation when the matrix has two maximum eigenvalues in absolute values but with opposite signs, and the convergence rate is determined by the absolute value of the ratio of the second-largest-magnitude eigenvalue to the first. We ignore these aspects in our heuristic approximate Algorithm 5. Results are presented in Tables 5, 6 . From them, we can see that also this notion of L estimate along the set of iterates has the same property as global L, namely during restricting subspace of training variables the L-the smooth constant is non-increasing similar to previous estimation method. For a fair comparison, we run our algorithm using the same calibration data as the baseline algorithms, typically a sample from RedPajama [13] . However, the way we handle this calibration data differs from most baselines [21, 18, 71] .When analyzing their codebase, we found that these algorithms resample calibration data by taking a random excerpt of a fixed length from a random document in the calibration data, both sampled uniformly. However, with this approach, the data from longer documents (e.g. books) are underrepresented compared to shorter ones (e.g. news articles), which biases the calibration data.Upon further investigation, we believe that new methods blindly copied this code from each other, going back to GPTQ [22] and possibly further. This choice was harmless for GPTQ since it requires relatively little calibration data; however, full model fine-tuning like in QuIP# [9] and AQLM [21] , works better on unbiased data.To remove the bias, we use standard 11 LM preprocessing that concatenates all documents, then splits them into evenly sized chunks that become training sequences. The benefits from debiasing range from insignificant to as large as 0.15 perplexity for some models. To compensate for that, we run experiments with the same preprocessing protocol unless explicitly stated otherwise.When done naively, the longest operation is the discrete update ( 8) that runs discrete optimization on all LLM parameters. For scalar quantization, this step does simple rounding and runs nearly instantly.In turn, applying it to vector quantization requires solving a discrete optimization algorithm (e.g. beam search) for every group of weights. However, since equation ( 8) can only update weights within S k , we can skip discrete optimization for any weight that was not among the chosen few. As a result, when training models with up to 70 billion parameters, we search less than one billion times per step.The next longest operation is computing the gradients ∇ϕ(•), needed both for P and V steps. This involves running LLM multiple forward and backward passes on batches of texts and accumulating the gradients. To reduce the overhead from gradient computation, we compute the gradients once using mixed precision, then reuse these gradients for one P and one V step, respectively. In other words, we switch from alternating P and V steps to performing these steps simultaneously. We also reuse these gradients to update any parameters not affected by quantization: input embeddings, normalization layers, biases, etc.To limit VRAM usage, we use gradient checkpointing [27] , batch accumulation. For larger models, we also use parameter sharding 12 [55] and optimizer offloading [57] . We need these techniques so that smaller 7B LLMs can be trained on a single GPU and larger ones with 70B parameters fit into a single machine with 8×A100. Still, PV-tuning takes up to 1.5× longer than tuning continuous parameters and uses more memory (during training) to hold the gradients w.r.t. dequantized weights. Here, we describe some of the implementation details we used to optimize different quantized representations. For every optimizations, we check that this optimization improves the algorithm in both MSE and perplexity and does not require additional resources that would result in unfair comparison.Vector Quantziation. The original algorithm quantizes all weights a single pass over input channels. We found that it works slightly better if we make multiple such passes and, between passes, update codes by Adam to minimize the same objective [72] . This is resembles QuIP# with no RHT & lattices or AQLM with no additive quantization. For simplicity, we also use a single shared codebook (e.g. instead of groupwise codebooks).VQ+outliers To select outlier coordinates, we use https://github.com/fmfi-compbio/ admm-pruning that outperforms the SpQR outlier criterion [18] in both L2 error and perplexity (when both criteria are used alongside vector quantization). We re-run the ADMM procedure multiple times during optimization, resulting in an EM-like algorithm.VQ+lowrank. We experimented with two different initializations for low-rank correction: a) quantizing weight matrix, then training LoRC on quantization error, as in [79] and b) initializing LoRC to approximate the reference weight matrix, the applying vector quantization to LoRC errors. Of these two approaches, we found that the latter one produces a (slightly) better solution in both MSE and perplexity.Additional representation evaluations. Below, we report some additional quantized representation configurations that extend our evaluations from Section 4.1. For convenience, we report them both as per-method perplexity values in Table 7 and the combined plots in Figure 4 . On the (right), there is a plot that combines the first two, allowing for a better comparison of each method with and without fine-tuning. Compression algorithms without fine-tuning are represented with dashed lines, while algorithms with fine-tuning are represented with continuous lines.Table 8 : Evaluation of quantized LLAMA 2 models for 2.x bits per weight. We use the same setup as in the main paper (Section 4.2 with an extra baseline). As requested, we finetune the model in 16-bit precision for the same number of steps, then quantize it with AQLM, reported as "FT+AQLM". Finally, the "Finetuned" row corresponds to an uncompressed 16-bit model finetuned without quantization. We hypothesize that fine-tuning the model has little effect since we train on a dataset resembling its original pretraining data. K Additional Details for Section 4.2 VQ: vector quantization as a simple near-optimal algorithm. We use a single 16-bit codebook with group size 16 (over input dimension) and per-channel trainable scales over output dimension. During P step, we update the codebook, scales and non-quantized model layers by backprop. During V step, we try every code in the codebook and choose the one that minimizes (6) .GPTQ: scalar quantization with block-wise scales and zero points. We use 2-bit base codes and block size 128. During P step, we update the scales and non-quantized parameters by backprop. In turn, V step performs simple rounding to nearest (scaled) integer. 13 . AQLM: we perform scalar quantization with two 8-bit codebooks and group size 8 and channel-wise steps. This was originally published in [21] as the "speed-optimized" configuration capable of fast lookup-based inference. During P step, we update both codebooks, as well as scales an non-quantized parameters by backprop. On V step, we run beam search with beam size 8 with gradient-updated dequantized weight as target.Training. We minimize Kullback-Leibler divergence as our loss function for all three representations. More specifically, we fine-tune the quantized "student" model to approximate the predictions (logits) of a "teacher" -the same model prior to quantization. We fine-tune on the same RedPajama sample as in calibration. More specifically, we use the official one-billion-token sample 14 provided by the dataset authors [13] . We use a batch size of 2 20 (≈ 1M) tokens, split into batches of model-specific sequence length (e.g. 4096 tokens for LLAMA 2, 8192 for LLAMA 3). In early experiments, we found PV to be resilient to the choice of batch size, consistently training with up to 4× smaller batches.Hyperparameter tuning: we tune the hyperparameters for each method individually. For all algorithms, we tune learning rate on a logarithmic scale out of (1e-4, 3e-4, 1e-3, 3e-3, 1e-2). For methods involving discrete optimization, we tune learning rate for P and V step separately. The optimal configuration for STE and stochastic rounding is to use learning rate 3e-4 for both codes and codebooks. The optimal configuration for subspace linearized PV and the same with STE is to use learning rate 3e-4 for P step and 3e-3 for codes. Curiously, the subspace methods are stable even with larger step sizes for codes, e.g. 1e-2, whereas unrestricted methods (e.g. pure STE) are not.For stochastic rounding, we found that the unbiased rounding [53] causes the model quality to quickly degrade, likely due to the fact that the algorithm makes too many weight changes due to rounding. The results we reported in Table 1 use stochastic rounding with temperature 0.2. In other words, we measure the distances to 2 nearest bins and round to each bin proportionally to distance -5 . We also tried gradually annealing the rounding temperature to 0, but achieved only insignificant improvements in accuracy and perplexity (<0.01). To simplify evaluation, we do not use annealing in Table1.For PV-tuning and PV-tuning with STE, we always set τ to maximum number of weights such that updating them satisfies ||x k+1 -x k ||/||x k || < 0.01. We implement this by trying to update large portions of weights (in our implementation, we update 0.01 of all weights at a time) until the total change exceeds the constraint. Once it does, we rollback weights in the last chunk to x k until the constraint is satisfied. As an implementation quirk, we always allow at least one weight to change even if changing just one weight already exceeds 0.01 • ||x k ||. However, we didn't see this in practice.We also found that Lamb 15 [81] is more stable when training with large batch sizes, but converges to approximately the same accuracy. We use β 1 =0.9 and β 2 =0.95, same as in most LLM training configurations [69, 86, 61] . We do not use learning rate decay for simplicity. It is likely possibly to improve our results by annealing the learning rates during training or using a warmup. We intentionally avoid this to reduce the number of "moving parts" and simplify evaluation. Overall, we found that PV-tuning is about as sensitive to hyperparameters as continuous-only LLM fine-tuning [71, 21, 63] .In the codebase, we release several additional PV-tuned models with a more careful choice of training hyperparameters and calibration datasets: for instance, when quantizing instruction tuned models, we found that calibrating the model on chat-like data results in better accuracy. 13 We also found that we can perform V step by running the entire OPTQ calibration while using y -1 L ∇ϕ(y) as target weights, showing the flexibility of the general PV framework. This variant converges to the same values, but is much slower due to having to re-accumulate the L2 error hessians for each V step. 14 https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T-Sample 15 Lamb is a variant of Adam that limits the norm of weight updates relative to the norm of weights.QuIP# is a popular modification of classical vector quantization for neural network compression. Unlike standard VQ, QuIP# does not quantize weights directly, but instead applies a Randomized Hadamard Transform (RHT) and quantizes the projected weight representations. This is done by applying the Hadamard transform to the rows and columns of the m × n matrix and multiplying them by random vectors S V ∼ U{±1} n , S U ∼ U{±1} m after each transformation. These vectors do not need to be stored, but may be generated on the fly from a random seed.The main objective of this transformation is to ensure that the Hadamard-transformed weight matrix adheres to a normal distribution. With this in mind, QuIP# no longer needs to learn the optimal quantization "codebook", but may instead use the optimal lattice for quantizing normal distribution. As a result, QuIP# does not need to store the codebooks in memory, thereby reducing its overall bitwidth. For instance, for Llama-2 7B vector quantization with group size 8 and 16-bit codes, VQ produces a 2.29 bit model while QuIP# has slightly less than 2.01 bits per parameter.We apply PV-Tuning to QuIP# in a manner similar to Vector Quantization. We start with a model already quantized using the original QuIP# algorithm, specifically employing the official 2-bit quantization of the Llama 2 7B model 16 . We then optimize this model using the same hyperparameters as those for Vector Quantization (see Appendix K).During P step, we optimize quantization scales (SU, SV), as well as any non-quantized model parameters (embeddings, LM head, normalization scales and biases), but not the codebook, since QuIP# relies on a predefined lattice structure that cannot be trained directly. During V step, we update the discrete codes responsible for choosing a vector out of the said lattice using the same procedure as in standard vector quantization.We compare PV-Tuning against several other fine-tuning strategies: no fine-tuning, the built-in finetuning procedure proposed in [71] , as well as the technically improved version of that procedure using the observations described in in G (denoted as "improved FT"). We report the resulting perplexities and zero-shot accuracies in Table 9 and include them in other relevant tables as "QuIP#+PV".Table 9 : Evaluation of LLAMA-2 7B quantized using QuIP# with various fine-tuning strategies. We report perplexity on WikiText-2 [45] & C4 [54] and zero-shot accuracy. The Average is the mean accuracy of 5 zero-shot tasks. Primary metrics are Wiki2 (PPL), C4 (PPL) and Average (Accuracy). In general, we found that QuIP# shows similar if not better improvements from PV-Tuning as VQ and AQLM. Furthermore, when compared to only fine-tuning continuous parameters, QuIP# shows better relative improvement from V steps than traditional vector quantization. We attribute this to the fact that VQ can fine-tune its codebook even without PV-Tuning, whereas QuIP# relies on fixed codebooks that cannot be learned. As a result, discrete parameters (codes) take up a larger fraction of the total model size in QuIP#, allowing V steps to make more significant improvements. The detailed instructions for running these experiments can be found in our official implementation 17 , in a separate subsection of the common README file.As discussed above, there are several ways to achieve the same number of bits per parameter with vector quantization. In this section, we explore the effect of varying the group size: either quantizing groups of 16 weights with larger codebooks, or groups of 8 weights with smaller ones. The results with different group and codebook sizes are reported in Table 10 .Table 10 : Evaluation of LLAMA-2 7B quantized using VQ+PV with different group size (GS, also known as vector dimension) and code bits (CB, s.t. codebook size is 2 CB ). We report perplexity on WikiText-2 [45] & C4 [54] and zero-shot accuracy. The Average is the mean accuracy of 5 zero-shot tasks. Primary metrics are Wiki2 (PPL), C4 (PPL) and Average (Accuracy). Overall, we figure out that smaller groups provide better performance for the same model size. We hypothesize that PV-Tuning is better able to deal with these configurations because with smaller group size, discrete codes constitute a larger fraction of model parameters, which allows the V step of our algorithm to achieve more significant improvements. In the initial version of manuscript we performed all 1-bit experiments with suboptimal group size 16. More recent results for group size 8 are provided under a separate name "PV (gs 8)" in Table 2 .In this section, we report additional results for LLAMA 2 & 3, MISTRAL and PHI-3 and discuss baselines. In this section, we always evaluate PV-tuning for vector quantization, using 14-16 bits per codebook for a group of 8 or 16 weights, with each combination fitting a particualr niche. For instance, 16 bits per 8 weights is slightly over 2 bits per weight, whereas 14 bits per 16 weights is either at or below 1 bit per weight, depending on the model size.We use LLAMA 2 models as our main benchmark as they are well studied in the PTQ community.Here, we gather the latest state-of-the-art algorithms at the time of publication and group them according to their target number of bits, roughly 1-1.7 bits per weight (Table 11 ) and 2-2.5 bits per weight (Table 12 ).Both our training runs and almost all baselines use the same sample of RedPajama data from previous sections 18 . The only exception to this is OneBit that uses a corpora of LLM outputs gathered specifically for that paper [77] .The source code for this method was unavailable until very recently. The official link https://github.com/xuyuzhuang11/OneBit used to point to an empty repository until the code was released in a commit https://github.com/xuyuzhuang11/OneBit/commit/ 380a6aedc3c060993056ff50b79065e893be99ae on May 10th. Thus, unfortunately, we did not have time to make OneBit compatible with models except LLAMA 2 7B and 13B that were featured in the original paper.For LLAMA 3, we evaluate PV-tuning of vector quantization against the baselines introduced in [33] . Curiously, their paper seems to compute perplexity differently than our paper. Since our protocol matches with most prior works [22, 18, 21, 71] , we chose to re-evaluate the results from [33] with our perplexity code and not the other way around. We calibrate using the official code 19 and reuse published models where available.In this section, we demonstrate that PV-Tuning can achieve speedups by using fast inference kernels from the underlying quantized representation. Since our main experiments use vector quantization, we adopt simplified version of AQLM inference kernels with one codebook of size 16 for groups of 8 consecutive weights. This kernel is not written by us: it was added to the official AQLM implementation by an open-source contributor.We adapt this inference code to our codebase and switch it to using group size 16 to support out 1.1-1.58 bit models. We evaluate inference speeds on a single Nvidia RTX 3090 GPU using transformers with cuda graphs 20 . P The Choice of the Initial Point x 0 Algorithm 1 converges from any initial point x 0 ∈ R d c (3.1), but it might converge to a different final point x(x 0 ). In this section, we discuss two possible variants of instantiating the initial point. U = {u i ∈ R d } : |U | = c. 3: Sample d random elements x 0 i ∈ U : x 0 = (x 0 1 , • • • , x 0 d ) T such that |V (x 0 )| = c.Different random x 0 ∈ R d c provide different loss functions. To get more stable and smoothed results we ran the algorithm (1) with different initial points, generated by algorithm (6) . We will denote r as the number of runs with different random initialization. where ϕ : R d → R, and R d ≤c ⊂ R d is the set of all vectors in R d whose d entries take at most c distinct values. In other words, the cardinality of the setEQUATIONis at most c. For small experiments, we aim to minimize the following objectiveEQUATIONwhere a i ∈ R d , x ⋆ is a unique optimal point: ∇ϕ(x ⋆ ) = 0 andEQUATIONWe seta i = i /d for i ∈ [d]. Hence, we have the Lipschitz continuity for (18) with L = 1. Note that the Algorithm 1 does not guarantee to converge to x ⋆ . Let x denote the vector to which the Algorithm 1 converges.Q.2 Tiny-scale experiments (d = 6, c ∈ [1, 6] )We applied PV algorithm to the problem (16) with ϕ(x) being ( 18) (Fig. 6 ). Number of runs with different random initial points r = 50. The starting point x 0 is chosen randomly using the ng algorithm (6) .When c = 1 we converge just in one step of PV (red line on Fig. 6a ) because we have only one degree of freedom to play with and we fully utilize it on the first step. One can show that the solution for the case c = 1 will be x =d i=1 a i -1 d i=1 a i x ⋆ i .Note that as we increase the maximum number of unique elements c, the final loss ϕ(x) decreases up until the moment when c = d, when the loss is zero (the plot 6a is in logarithmic scale and that is why we cannot see the last line).We run PV algorithm (1) with different random starting points x 0 ∈ R d c (6) . Different starting points can lead to different local optimum. That is why we ran the algorithm (1) several times with different random initial points x 0 . A number of runs r = 50, we used this value for all further experiments. Each of the two steps of the PV algorithm contributes to the convergence. To observe this we plotted the loss function over the iterates and explicitly marked the progress of both P and V steps (Fig. 6b ). We can see that we have progress during each of these steps and one single P and V step is not enough to obtain a solution even in this very small and simple case.These simple experiments demonstrate that 1. larger c (smaller compress ratio) provides better final accuracy 2. several P and V steps are needed to converge to the solutionQ.3 Interpretation of P (y) ⊒ P (x k )As we mentioned before, we select the initial point x 0 ∈ R d c randomly such that it has c unique elements in its linear shell V (x 0 ). To understand the notation P (y) ⊒ P (x k ), defined in the chapter (3.1), let us consider the first step of the algorithm, specifically the P step.On the first step of the algorithm we select a new point by solving the following optimization problem:EQUATIONHere we have a random point x 0 , with some partition P (x 0 ) = {P 1 (x 0 ), • • • , P c (x 0 )} -c sets of equal elements from x 0 . We find y 0 such that minimizes ϕ(y 0 ) and have the same or smaller number of partitions: P (y 0 ) = {P 1 (y 0 ), • • • , P s (y 0 )}, where s ≤ c.Strict equality of two sets, P (x 0 ) and P (y 0 ) would mean that the linear shell -the number of unique elements of the vector x 0 is equal to the number of unique elements from y 0 (c = s). In our PV algorithm we allow y 0 to have smaller number of unique elements than x 0 , so we possibly merge some partitions of x 0 .It is worth to mention that in real experiments we observe that the number of unique elements did not change or the change is not significant. , where c = 2. We have P 1 (y) = {1, 2, 3, 7, 8} and P 2 (y) = {4, 5, 6}. So, P (y) = {{1, 2, 3, 7, 8}, {4, 5, 6}}. Notice that each element of P (x) is a subset of some element of P (y). For example, {2, 8} is a subset of {1, 2, 3, 7, 8} and {6} is a subset of {4, 5, 6}. Because of this, by our definition, P (y) ⊒ P (x).Q.5 Small-scale experiments (d = 100, c ∈ [1, 100])The problem of the algorithm (1) is that the V step requires the full parameter search which gives us the complexity O(c d ). Even for small tasks, this becomes unpractical to solve. Example Q.1. Let us take d = 100 and c = 10, then the complexity of one V step will be O(10 100 ). Modern computers can make roughly 10 petaflops or 10 16 calculations per second, hence we will have to wait ∼ 10 76 years to make a single V step.Let us consider special sets of function ϕ(x) that we will call separable functions. This class of functions should satisfy the assumption (Q.2). Assumption Q.2 (Separable function). The function ϕ(x) : R d → R can be written in the form ϕ(x) = d i=1 ϕ i (x i ), where ϕ i (•) is a mapping ϕ i (•) : R → R. Example Q.3. The objective function ( 18) is a sum of squares that can be written in the following form:EQUATIONwhereϕ i (x i ) = a i (x i -x ⋆ i ) 2 .Hence, the objective ( 18) is a separable function. One can show that for separable functions (Q.2) the algorithm (1) can be written in the form (7) . Hence, for separable functions, we can compute the V step in O(c • d) operations (instead of O(c d )), which makes the algorithm (1) practical to use. for i = 0, 1, . . . , d do 5:x k+1 i = arg min xi∈R ϕ i (x i ) : x i ∈ V (y k ) (Optimized V step) 6:end for 7: end for Then we ran the optimized PV algorithm (7) on the quadratic objective (18) . The results are presented in (Figure 7 ). Number of runs with different random initial points r = 50.For these experiments, we see the same results as for tiny scale with d = 6 (Q.2): larger c (smaller compress ratio) provides better final accuracy, and several P and V steps have to be done to obtain better solution. In addition, we observe that the PV algorithm does not decrease the dimensionality (number of unique elements) of V (x).As it was mentioned before, the algorithm (1) is conceptual only and cannot be used in the raw form in practice. That is why we need to move to the experiments with linearized V step. Q.6 Linearized PV Linearized V step (B.1) allows us to greatly reduce the cost of one V step. We ran Linearized PV on the quadratic objective (18) (Fig. 8 ). From the first experiment (Fig. 8a ) we can see that 1. increasing T provides slightly better convergence rates and approximately the same final accuracy. We saw the similar results on large-scale experiments. Hence, we can use T = 1 to save computations.2. Linearized PV algorithm (B.1) converges to a worse accuracy than the exact PV (1).3. Linearized PV has to make more iterations than PV to convergeThe second plot (Fig. 8b ) demonstrates the effect of multiple Linearized V step. We can see that the largest effect comes from the first V step.In previous section (Q.6) we have seen that Linearized PV algorithm converges to a worse accuracy than the exact PV (1).Linearized PV (B.1) with combination of sparse updates (3.3) is intended to mitigate this issue.On (Fig. 9a ) we can see comparison of three methods: exact PV (1) -red line, Linearized PV (B.1)blue line and Linearized PV with sparse updates (3.3) -green line.From this experiment we observe 1. Linearized PV with sparse updates converges to a better accuracy than Linearized PV 2. inearized PV + sparse updates has to make more iterations than Linearized PV and exact PV to convergeHence, this approach helps us to converge to a better accuracy, but with a price of larger number of iterations to converge. We can use different rules for choosing the subspace S k which produce different convergence rates and the final accuracy levels (9b). x 0 is random x 0 is random x 0 is random x 0 is random x 0 is random x 0 is random x 0 is random x 0 is random x 0 is random x 0 is random (b) Degradation of dimensionality of V (x) for Linearized PV algorithm (B.1). In large scale experiments we used sparse updates with S k being chosen greedily based on the absolute values of the entries of gradient vector |∇ i ϕ(y k )| -green line. Here we demonstrated that other strategies based on random uniform sampling or random proportional to |∇ i ϕ(y k )| sampling can be even more advanced providing better final accuracy at cost of larger number of iterations to converge -purple and orange lines (Fig. 9b ).We can see that Linearized PV with sparse updates can converge to even a better accuracy than the exact PV algorithm. The problem of Linearized PV algorithm is that we come to the local minimum and cannot get out of that minima because of the small value of the gradient (we are near to a real solution) and small stepsize.Linearized PV with sparse updates allows us to mitigate this problem by reducing the subspace from R d to S k in which we are solving this optimization problem. This allows us to greatly reduce the local Lipschitz constant from L to L S k and hence significantly increase the stepsize γ = 1 /L S k . We demonstrated how L S k changes for Linearized PV with different sparse updates sampling methods (Fig. 10a ).We can have degradation of |V (x k )| during both the P and V steps.Here are examples for both of them with optimizing simple quadratic objective ϕ(x) = d i=1 (x i -x ⋆ i ) (case when all a i , i ∈ [d] in ( 18) are equal to one).1. Degradation during the P step: Let x ⋆ = {0, 2, 1}, x 0 = {x, x, y}, so |V (x 0 )| = 2, then after the P step we will have y 0 = {1, 1, 1}, hence |V (y 0 )| = 1. 2. Degradation during the V step: Let x ⋆ = {2, 10, 0, 11}, x 0 = {x, y, y, z}, so |V (x 0 )| = 3, then after the P step we will have y 0 = {2, 5, 5, 11}. Finally, after the V step we have x 1 = {2, 11, 2, 11}, so |V (x 1 )| = 2.In real experiments we observe decreasing of |V (x)|. You can observe this phenomena with Linearized PV algorithm (10b). As we can see we have a big decreasing of |V (x)| during the first several iterations. We can use this gap to find even better solution with improved final accuracy and faster convergence rate.To add additional unique element in V (x) we considered the modification of V step: V + step, where we allow |V (x)| to become larger up to some upper bound.Let ϕ(•) is a mapping ϕ(•) : R d → R and the vector x ⋆ ∈ R d be the optimal point: ∇ϕ(x ⋆ ) = 0. Let the vector x ∈ R d ≤c and define the set W (x, x ⋆ , c) such that it contains at most c -|V (x)| unique elements from V (x ⋆ ) without elements from V ( and because x = y k satisfies the constraint V (x) ⊆ V (y k ) W (y k , x ⋆ , c), we conclude that ϕ(x k+1 ) ≤ ϕ(y k ).The rest of the proof is identical to (A.1).The main impact of our work, both positive and negative, is in the ability to deploy higher-quality LLMs to run on memory-limited devices like desktops, laptops, and phones. On the positive side, this would allow practitioners to develop offline LLM applications (e.g. translate service), lower-latency chat assistants that are not dependant on network latency, or privacy-sensitive LLM applications where the user's private data never leaves their device. Furthermore, this can facilitate the creation of free open-source software based on LLMs by eliminating the need to maintain costly inference servers on the backend. Since phones are everywhere and LLMs are powerful general-purpose tools, PV-tuned models could significantly impact how the general population uses LLMs to complete tasks.However, LLMs are still a dual-use technology with the potential for significant benefits and serious harm. Risks range from deliberate misuse (e.g. spam generation) and accidental misuse to negative economic side-effects. An upper bound on these risks is that PV tuning does not create new (potentially risky) LLM capabilities, merely making existing ones more accessible.Justification: We report error bars for small-scale experiments Q. For full fine-tuning runs, we do not include error bars since running those would be prohibitively costly for us.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar and then state that they have a 96%CI if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.Justification: Our research is focused on the base capability and accessibility of LLMs. While working on LLMs always has potential externalities, our specific work adheres to the ethics guidelines.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained? Answer: [NA] Justification: We did not conduct any research on human subjects. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.By "one-shot" we refer to methods that quantize quantize the model in a single pass over calibration data.The official implementation is available at https://github.com/Vahe1994/AQLM/tree/pv-tuning.We explain how this generalizes to other quantized representations in Appendix CIt is possible to consider different class of functions instead; e.g., Lipschitz functions. In such a case, we would us a different approximation. For simplicity of exposition, we work with L-smooth functions.Very closely related methods include block coordinate descent and compressed gradient descent with sparsification operators such as RandK or TopK[4,7].https://x.com/black_samorez/status/1821933255744016866https://galqiwi.github.io/aqlm-rs/about.htmlhttps://github.com/ggerganov/llama.cpphttps://reddit.com/r/LocalLaMAfrom e.g. https://github.com/huggingface/transformers/blob/main/examples/pytorch/ language-modeling/run_clm.pyWe use PyTorch FullyShardedDataParallel[52,40] and wrap discrete weights as in bitsandbytes[19]https://huggingface.co/relaxml/Llama-2-7b-E8P-2Bithttps://github.com/Vahe1994/AQLM/tree/pv-tuninghttps://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T-Samplehttps://github.com/Macaronlin/LLaMA3-QuantizationThe specific version of each library can be found in 'requirements.txt' 


Decision: The paper considers compression of LLMs through quantization, focusing on fine-tuning techniques over initially compressed weights so as to better approximate the performance of the unquantized model. The method PV-Tuning is developed and compared with existing baselines. 

The reviewers have noted that the method is novel, reveals the importance of discovering good fine-tuning strategies, extensive experiments validate the method and show improvement on baseline methods, and achieves good performance on Llama 3 (where GPTQ-based methods do not). 

The optimization is a mixed continuous-discrete case that complicates the solution. The PV-Tuning method is compared with and shown to outperform straight through estimation (STE), and some new insights into the problem and methods is gained as well. 

Although significant, the PV-tuning method seems to have complexity relative to existing fine-tuning methods and should be compatible with various initial quantization methods.  

The paper provides some convergence analysis and guarantees, contrasting with straight through estimation (STE) and stochastic rounding that generally lack such guarantees.  The reviewers noted that the performance analysis was not always clearly presented.

The rebuttal and discussion with reviewers has led to further interesting experiments and potential algorithmic improvements.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: Yes
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: Please see the Weaknesses and Questions.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: questions: see above
limitations: limitations: yes
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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 ?
limitations: limitations: N/A
paper_summary: paper_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.
strength_weakness: strength_weakness: - **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: 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?
limitations: limitations: See weaknesses and questions.

Comments: title: Thanks
Makes sense
Response to official comment by reviewer 3n7c
Response to "more comments"
Response to "full precision fine-tuning?"
more comments
full precision fine-tuning?
comment: Thank you for answering my questions. After reading the rebuttal, I keep my initial rating.
Thanks for the additional information. I am happy with the author's responses and will keep my score.
Thanks for clarifying. I guess all that experiment shows is that the fine-tuning dataset isn't the same as the pretraining dataset, which is obvious.
In terms of computation efficiency, PV-Tuning is roughly on par with STE since both of them are bottlenecked by accumulating gradients w.r.t. dequantized LLM weights. PV-Tuning with subspaces is slightly more efficient because it only computes V step for a small portion of parameters, but otherwise, the total runtime is comparable. The memory efficiency of STE and PV-tuning are equivalent. We discuss efficiency further in Appendix I for calibration and Appendix M for inference. Alternatively, if you meant efficiency as model accuracy, we analyze this in Section 4.2.
We are glad that the reviewer shows interest in our work and do our best to answer further questions:

> 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?

In general (for VQ), we do an exhaustive search for each code in the chosen subspace, but **not all combinations**. This is because our Linearized V step (Section 3.2) minimizes mean squared error with the adjusted weight matrix (Equation 5). This problem can be solved independently for each discrete code because the optimization objective is a sum over independently chosen discrete weights.

Still, this can be a costly procedure for LLMs with billions of parameters. However, it *becomes cheap since we only update a small fraction of codes*, i.e. the subspace. Thus, we run this costly optimization (1) for less than 1% of all weights and (2) in parallel on GPUs, making it practically feasible.

In the supplementary code, the discrete optimization takes less than 10% of the total training time when running on A100s. However, further improvement of this algorithm is an interesting direction for future research.

> 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.

We thank the reviewer for finding this out, as this is an important suggestion. We accidentally did not include these numbers when compiling the response. You may find our response below:

We conducted additional experiments using this configuration for Llama 2 7B in the same setup as the rest of our experiments (see Section 4.2). Specifically, we use a single 10-bit codebook with a group size of 8 consecutive weights, which leads to **1.26 bits per weight**, including codebooks. The resulting quantized model has a perplexity of 6.82 on WikiText-2 and 8.82 on C4. We compare these results with the previous configuration below:

| Method | Avg Bits | Wiki2 $\text{PPL}\downarrow$ | C4 $\text{PPL}\downarrow$ | ArcC $\text{Acc.}\uparrow$ | ArcE $\text{Acc.}\uparrow$ | HellaSwag $\text{Acc.}\uparrow$ | PiQA $\text{Acc.}\uparrow$ | WinoGrande $\text{Acc.}\uparrow$ | Average $\text{Acc.}\uparrow$ |
|---|---|---|---|---|---|---|---|---|---|
| $-$ | 16 | 5.12 | 6.63	| 43.43 | 76.3 | 57.14 	| 78.07 | 69.06 | 64.80 |
| PV-Tuning | 1.58 | 7.32 | 9.35 | 25.85 | 57.58 | 40.88 | 68.99 | 57.77 | 50.21 |
| PV-Tuning (10-bit codebook) | 1.26 | 6.82 | 8.82 | 32.94 | 68.27 | 49.01 | 74.76 | 64.17 | 57.83 |

**These are very promising results.**, we originally opted for 8- and 16-bit codes as they are easier to handle in modern hardware, but in principle, it is possible to pack 10-bit codes in a way that supports efficient inference. Following your suggestion, we will evaluate more configurations with small codebooks in the final version. Likewise, we will further explore structured codebooks as the reviewer suggested.

> 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.

Thank you for clarifying this. By this upper bound criterion, the Pareto-optimality would indeed be achieved at a somewhat higher bitwidth. In our submission, we used a popular criterion from prior works, but your upper bound criterion could bring more nuance to the Pareto-optimality. We will add the comparison to FP16 @ 3 bits in the final version of the paper.
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.
Thanks for your responses. Another questions: have you compared the efficiency of PV-Turning with STE or other baselines?
- 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.
> (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?
Thank the authors. After reading the rebuttal, I keep my initial scoring.