DiscQuant: A Quantization Method for
Neural Networks Inspired by Discrepancy Theory

Modern deep learning models continue to grow in size, incurring greater challenges to train and serve these models. Post training compression methods have emerged which aim to make model inference faster and cheaper. Compressing after pretraining is desirable among practitioners who either cannot afford to train models themselves, or do not want to change the expensive training process too much. In this paper, we study post training quantization (PTQ) of the model weights. Quantization reduces the memory requirements of the model, and speeds up inference for LLMs under memory-bound settings such as the generation phase (as opposed to prefilling phase which is compute-bound) (Kwon et al., 2023).

The quantization problem can be divided into two overall steps: (1) Construct a good low bit-complexity representation for the weights (we colloquially call this the quantization grid), and (2) Round the original weights to values in the quantization grid. Within step (1), we also consider those methods which apply a transformation on the weights to better match the encoding format. There has been much recent work on weights-only PTQ for LLMs. To date, the vast majority of such research has been focused on step (1): constructing good low bit representations (Shao et al., 2024; Tseng et al., 2024a; Egiazarian et al., 2024). However, work on rounding methods is under-explored. To the best of our knowledge, Round-to-Nearest (RTN) and GPTQ (Hassibi et al., 1993; Frantar et al., 2022, 2023) are the primary rounding methods for LLM weight quantization. RTN is a simple baseline, and GPTQ is a data dependent method which aims to match the activations of the quantized model with that of the original model layer-by-layer.

Let $f(w;s)$ be the loss function of a neural network where $w$ are original pretrained weights and $s$ is an input sample; for example $f$ can be the usual cross-entropy loss on input $s$. To find a good rounding solution, we are looking for perturbations of the original weights $w\in\mathbb{R}^{n}$ that correspond to values in the quantization grid, and do not increase the loss $f$ too much. We further impose the constraint that we only round each parameter up or down, this ensures that we are not changing the original model weights too much. Then the set of allowed quantization points can be pictured as vertices of a hypercube $H$ around $w$. Let $\hat{w}\in\mathbb{R}^{n}$ be these perturbed weights, and $\Delta f=f(\hat{w};s)-f(w;s)$ be the resulting change in loss function for a sample $s$. We approximate $\Delta f$ via a first order Taylor expansion: $\Delta f\approx\langle\nabla_{w}f(w;s),\hat{w}-w\rangle$. Some prior works such as Nagel et al. (2020); Hassibi et al. (1993) assume the gradients of a pretrained model to be nearly zero, and focus on the second order terms. We show that this assumption is not always true, the average gradients are close to zero but per-sample gradients can be big; in fact the first order term is a good approximation to $\Delta f$ (see Figure 4).

Therefore, to incur a small $\Delta f$, we want $\langle\nabla_{w}f(w;s),\hat{w}-w\rangle\approx 0$ for $s$ sampled from the data distribution $\mathcal{D}_{\textrm{data}}$. Suppose we are given $m$ independent samples $s_{1},s_{2},\dots,s_{m}\sim\mathcal{D}_{\textrm{data}}$, we can impose the constraints $\langle\nabla_{w}f(w;s),\hat{w}-w\rangle=0$ which correspond to an affine subspace $V$ of dimension $n-m$. The intersection of the subspace $V$ and the hypercube $H$ is a convex polytope $K$. It can be shown that any vertex of $K$ should have at least $n-m$ fully rounded parameters, see Figure 1 for an illustration. Since the number of parameters $n\gg m$, any vertex of $K$ gives an almost fully rounded solution. Obviously this solution satisfies the linear constraints for the samples $s_{1},s_{2},\dots,s_{m}$. But will it generalize to unseen samples from the data distribution $\mathcal{D}_{\textrm{data}}$? We prove that it can generalize if the distribution of gradients $g=\nabla_{w}f(w;s)$ for $s\sim\mathcal{D}_{\textrm{data}}$ is approximately low rank. Let $\Sigma=\mathbb{E}_{s\sim\mathcal{D}_{\textrm{data}}}[gg^{T}]$ where $g=\nabla_{w}f(W;s)$ be the covariance matrix of gradients. We prove the following theorem; the algorithm and the proof draws on techniques from discrepancy theory, in particular the famous Lovett-Meka algorithm (Lovett and Meka, 2012).

From these insights we develop a practical rounding algorithm called DiscQuant. The Lovett-Meka algorithm does a random walk starting from the original weights until it converges to a vertex of $K$. Instead, we can find a vertex of $K$ by minimizing a linear function over the convex polytope $K$. DiscQuant uses stochastic gradient descent to minimize two objectives, one corresponding to low $\Delta f$, and the other corresponding to minimizing a linear function. We take a knowledge distillation approach for the first term, minimizing the KL divergence between the original and quantized model. These two losses are balanced with a regularization parameter $\lambda>0$:

Here $p_{w}(\cdot|z_{<i})$ is the next token distribution given prefix $z_{<i}$. An astute reader may notice that the first order approximation of the KL divergence in (1) is exactly zero, and how our discussion above applies. In Section 4 where we describe in detail our exact optimization objective, we also show that the second order term of KL divergence can be written as

So minimizing the KL divergence is a succinct way to impose constraints of the form $\langle\nabla_{w}\log p_{w}(t|z_{<i}),\hat{w}-w\rangle\approx 0$ or equivalently $\log p_{w}(t|z_{<i})\approx\log p_{\hat{w}}(t|z_{<i})$ where $t\sim p_{w}(\cdot|z_{<i})$ and $z\sim\mathcal{D}_{\textrm{data}}$. Therefore our framework still applies.

After every step of gradient descent, we project the weights back to the hypercube $H$. This ensures that the trajectory of DiscQuant remains within the convex polytope $K$ and eventually converges to a vertex of $K$ with almost all the coordinates rounded. Instead of picking a random direction $c$ to find a random vertex of $K$, we use a special $c^{*}$ which let’s us find the vertex closest to the original weights $w$ (see Section 4). We use RTN to round the few unrounded parameters left at the end of the optimization.

We perform extensive experiments which show the strength of our method: on models Phi-3-mini-4k-instruct and Meta-Llama-3.1-8B-Instruct, across a variety of evaluation tasks, and across the block scaling and incoherence processing quantization formats. DiscQuant is agnostic towards the quantization grid, and can therefore be composed with other quantization methods. Block scaling sets a bits parameter which determines the number of grid points, and a unique scaling parameter per groupsize weights (Frantar et al., 2023). Incoherence processing applies a random orthogonal transformation, which reduces the weight ranges and can make quantization easier (Chee et al., 2023; Tseng et al., 2024a). A subset of results can be found in Figure 2. Across tasks, models, and quantization levels, our method DiscQuant achieves superior compression over baselines GPTQ and RTN.

We summarize our main contributions:

• •

  Theoretical developments: We prove that it is possible to achieve generalization error $\leq\varepsilon$ on the true data distribution by rounding all but $\mathrm{poly}(\log n/\varepsilon)$ weights, so long as the gradients of the original model are approximately low rank.

• •

  Practical algorithm: We develop a simple and practical algorithm DiscQuant guided by our theoretical analysis. We perform extensive experiments on Phi-3-mini-4k-instruct and Meta-Llama-3.1-8B-Instruct, over block scaling and incoherence processing quantization formats, and a variety of evaluation tasks. Our method DiscQuant achieves superior or comparable quantization to the baselines GPTQ and RTN as can be seen from Figure 2.

In this paper we focus on weights-only PTQ. Quantization can also be applied to the activations or KV-cache (Ashkboos et al., 2024; Liu et al., 2024a, b). Other compression method such as pruning
(Frantar and Alistarh, 2023; Sun et al., 2023) are also outside the scope of this work. As discussed in the introduction, post training quantization can be divided into two overall steps: (1) Construct a good low bit-complexity representations for the weights (the quantization grid), and (2) Round the original weights to the values in the quantization grid. To this date, the vast majority of PTQ research for LLMs has focused on step (1). Note that determining a good compressed representation can involve both encoding formats, as well as transformations to ensure the weights better match the encoding format.

Let $f(w;s)$ be the loss function of a pre-trained neural network with weights $w\in\mathbb{R}^{n}$ on an input sample $s$ and let $\mathcal{D}_{\textrm{data}}$ be the sample data distribution. Suppose we are also given a (scalar) quantization grid $\mathcal{Q}=Q_{1}\times Q_{2}\times\dots\times Q_{n}$ where $Q_{j}\subset\mathbb{R}$ is a finite set of quantization points available to quantize the $j^{th}$ parameter.¹¹1The quantization grid $\mathcal{Q}$ can depend on $w$, like in Block Scaling (Frantar et al., 2023). So ideally, we should write $\mathcal{Q}_{w}$, but we ignore the dependence to simplify notation. In this work, we focus on scalar quantization which allows us to write the quantization grid as a product set, i.e., each parameter can be independently rounded to a finite set of available values. Alternatively, in vector quantization a group of $d$ variables are rounded together to one of a finite set of quantization points in $\mathbb{R}^{d}$, which has been used in some prior works (Tseng et al., 2024a; Egiazarian et al., 2024; van Baalen et al., 2024). Generalizing our method to vector quantizers is an interesting future research direction.

Our goal is to find a rounding $\hat{w}\in\mathcal{Q}$ of the original weights $w$ such $f(\hat{w};s)\approx f(w;s)$ where $s\sim\mathcal{D}_{\textrm{data}}$. We further impose the constraint that for each parameter $w_{j}$, we only round up or round down to the available values in $Q_{j}$, i.e., we only have two choices for $\hat{w}_{j}$ denoted by $w^{\textrm{up}}_{j},w^{\textrm{down}}_{j}\in Q_{j}$ where $w^{\textrm{up}}_{j}\leq w_{j}\leq w^{\textrm{down}}_{j}$.²²2If $w_{j}<\min Q_{j}$ or $w_{j}>\max Q_{j}$, we just set $w^{\textrm{up}}_{j}=w^{\textrm{down}}_{j}=\min Q_{j}$ or $\max Q_{j}$ respectively. We make this assumption because we don’t want to change any parameter of the original model too much during quantization, consider it an important property of algorithms we design. Using Taylor expansion:

Assuming that the quantization grid $\mathcal{Q}$ is fine enough and since we only round each parameter up or down, $\left\lVert\hat{w}-w\right\rVert$ is small and so we can ignore the higher order terms. We claim that the first order term is the dominant term. Prior works such as Nagel et al. (2020); Hassibi et al. (1993); LeCun et al. (1989) have assumed that the first order term can be assumed to be zero because the model is trained to convergence and focused on reducing the second order term. But the model being trained to convergence just means that average gradient over many samples from the distribution is nearly zero. But the gradients still have some variance and gradients w.r.t. individual samples from the data distribution are not approximately zero (see Table 1). Figure 4 demonstrates this by showing that the error term $\Delta f$ is well-correlated with the first order approximation $\left\langle\nabla_{w}f(w;s),\hat{w}-w\right\rangle$.³³3In the special case when $f$ is the KL distillation loss between the original model and quantized model, the first order term vanishes exactly. See Section 4 for why this analysis still applies.

So the goal now is to find a rounding $\hat{w}$ such that $\left\langle\nabla_{w}f(w;s),\hat{w}-w\right\rangle\approx 0$ for samples $s\sim\mathcal{D}_{\textrm{data}}.$ Suppose we sample $m$ samples $s_{1},s_{2},\dots,s_{m}\sim\mathcal{D}_{\textrm{data}}$ independently from the data distribution, where $m\ll n$. We now break our task into two parts of bounding the empirical error and generalization error as follows:

In this section, we will present DiscQuant, a simple and practical algorithm for rounding inspired by the theoretical insights in Section 3. Instead of trying to approximate the loss function of the pre-trained model, i.e., $f(\hat{w};s)\approx f(w;s)$, we will instead take a distillation approach and try to minimize the KL divergence between the next token distribution of the original model and the quantized model. Let $p_{w}(\cdot|z_{<i})$ be the distribution of the next token predicted by the original model given prefix $z_{<i}$ where $z\sim\mathcal{D}_{\textrm{data}}$ is a sample from the data distribution. We want $\textrm{error}(\hat{w})=\mathbb{E}_{z\sim\mathcal{D}_{\textrm{data}}}\mathbb{E}_{i}D_{KL}\left(p_{w}(\cdot|z_{<i})\left\|\right.p_{\hat{w}}(\cdot|z_{<i})\right)\approx 0$.

Expanding $\textrm{error}(\hat{w})$ using Taylor series, we can see that first order term vanishes exactly and so the second order term is the dominant term (see Appendix D). By Lemma D.1, Hessian of $\textrm{error}(\hat{w})$ can be written as a covariance of gradients as:

Therefore

So minimizing $\textrm{error}(\hat{w})$ is a succinct way to impose constraints of the form $\langle\nabla_{w}\log p_{w}(t|z_{<i}),\hat{w}-w\rangle\approx 0$ or equivalently $\log p_{w}(t|z_{<i})\approx\log p_{\hat{w}}(t|z_{<i})$ where $t\sim p_{w}(\cdot|z_{<i})$ and $z\sim\mathcal{D}_{\textrm{data}}$. Therefore, we can use the same techniques developed in Section 3 to solve this as well. Assuming that the gradients are low rank, the set of $x$ satisfying these constraints (where $\hat{w}=w^{x}$) form an affine subspace $V$ of dimension $\geq n-m$ where $m$ is the number of samples. We are again interested in finding a vertex of the polytope $K=[0,1]^{n}\cap V$ which will have $\geq n-m$ integral coordinates. At this point, we could use the Lovett-Meka algorithm (Algorithm B.2) which has provable generalization guarantees. But explicitly calculating all the gradients and storing them is infeasible. Instead a simple heuristic way to find a random vertex of polytope $K$ is to minimize a random linear function. Let $c\in\mathbb{R}^{n}$ be some arbitrary vector; we will try to minimize the linear function $\left\langle c,x\right\rangle$ along with the KL divergence by taking a linear combination of them. The final optimization objective is shown in (3) where $\lambda>0$ is a regularization coefficient.

We solve the optimization problem (3) using projected stochastic gradient descent where we project $x$ to the hypercube after every gradient update. Optimizing (3) will keep us close the polytope $K$ and will approximately converge to a vertex of $K$ which is almost integral. We round whatever fractional coordinates are left using RTN to get a fully integral solution.

We use one additional heuristic to improve the performance of the algorithm in practice. Instead of choosing a random vertex of the polytope $K$ by choosing the vector $c$ at random, we will choose it carefully so as to find the vertex of the polytope $K$ which is closest to $y$ which is the interpolation point corresponding to the original model weights (i.e., $y$ such that $w^{y}=w$). We have:

where $c^{*}=(1-2y)$. Here we have used the fact that $x_{i}^{2}=x_{i}$ whenever $x_{i}\in\{0,1\}$ and since $x$ is almost integral, we can use the approximation in the summation above. With this approximation, minimizing $\left\lVert x-y\right\rVert^{2}$ over almost integral $x$ is equivalent to minimizing $\left\langle c^{*},x\right\rangle$. So in the DiscQuant algorithm, we use $c=c^{*}$ specifically instead of a random $c.$

We evaluate our method on the Phi-3-mini-4k-instruct (Abdin et al., 2024) and Meta-Llama-3.1-8B-Instruct (Dubey et al., 2024) models, and compare against GPTQ and greedy rounding (i.e. round-to-nearest, or RTN). We use the lm-evaluation-harness Gao et al. (2023) to evaluate on the Wikitext, GSM8k_cot 8-shot, MMLU 5-shot, ARC_Challenge 0-shot, PIQA 0-shot, HellaSwag 0-shot, and Winogrande 0-shot tasks. We report standard errors from lm-evaluation-harness. Wikitext measures perplexity, GSM8k is a generative task, and the remaining are multiple choice tasks. Note that generative tasks are typically more difficult than multiple choice tasks, and better reflect how the models are used in practice. See Appendix A for details on the hardware used, and hyper-parameter settings. Our method has similar memory requires as knowledge distillation, which also requires two copies of the model. We do not perform inference timing experiments; DiscQuant can optimize over a given quantization grid, so that we can utilize any pre-existing inference optimizations. For example, there are inference kernels for block scaling (Frantar et al., 2024) and incoherence processing (Tseng et al., 2024a). Ablations on the loss formulation are in Appendix A.