PARQ: Piecewise-Affine Regularized Quantization

We thank the reviewer for recognizing the strength of our paper (sound methodology and being well written) and giving us constructive suggestions on having more benchmark evaluations. We agree that additional empirical evaluations, especially on modern language models, will make the paper stronger. Although we do not have time to finish such experiments during the rebuttal period, we will try to include experiments on some basic language models in the final version to demonstrate its applicability.

We agree that ablation study on the choice of $\rho_t^{-1}$ is very useful for better understanding the behavior of the algorithm. The particular curve in Figure 10 is of the sigmoid family. Specifically,

$\rho_t^{-1} = \frac{1}{1 + \exp(s(t-t_1))}$

where $t_1$ is the transition center and $s$ is the transition steepness. This schedule changes $\rho_t^{-1}$ roughly from 1 to 0 (taking value $1/2$ at the transition center $t_1$). Essentially this changes the slope $\rho_t$ from 1 to $+\infty$. Here $s>0$ is the steepness parameter: the larger $s$ is, the steeper transition it has. For example, $s=0.1$ makes the transition almost linear, and the one shown in Figure 10 is with $s=10$.

Notice that the above sigmoid curve for $\rho_t^{-1}$ is essentially exponentially decaying as the reviewer suggested. Equivalently, we have the slope itself, $\rho_t$, increases exponentially.

In our ablation study, we train a 2-bit DeiT-T model with PARQ. We study the choice of transition steepness $s \in \\{0.1, 1, 10, 20, 40, 80\\}$ and transition center $t_1 \in \\{0.25, 0.5, 0.75\\}$, the fraction of training progress at which transition occurs. This sweep reveals that a shallow $s = 1$ performs best for the model configuration. A later transition of $t = 0.75$ improves QAT performance for shallower steepness values $s \in \\{10, 20\\}$, while higher $s \in \\{40, 80\\}$ are relatively unaffected. We will add the above description and details of the ablation study to the appendix in the final version of the paper.

We thank the reviewer for the overall positive feedback to our paper and especially recognizing our main novelty and contributions. Here we mainly address the reviewer’s question on the regularization approach proposed in the following reference, which we call Ref [1] hereafter.

[1] Towards Accurate Network Quantization with Equivalent Smooth Regularizer, ECCV2022

Ref [1] proposes smoother regularizers for inducing quantization, more specifically, of sinusoidal shape. It is clearly based on the intuition that such regularizers can help trap the weights to clusters close to a set of (evenly spaced) discrete values. This is similar to using W-shaped regularizers (which are nonsmooth) but with smooth functions in order to have a unique gradient in optimization. Unfortunately, such functions violate both properties we desired for a good regularizer for quantization: nonsmoothness and convexity (see the two paragraphs starting from Line 157 in our paper). More specifically,

• Smooth curves like sinusoid can trap weights into separate clusters, but does not induce quantization (i.e., concentrate on discrete values). This is due to the fact that locally near the local minima, the sinusoid behaves like the squared Euclidean norm, being flat thus does not induce quantization. In contrast, nonsmooth regularizers such as $L_1$ or W-shaped or PAR are locally sharp and thus can force the weights to concentrate at the local minima. As a result, Ref[1] still needs to use additional rounding or STE steps in order to obtain discrete quantized values.
• Convexity gives better global properties that help to avoid local minima, which explains the popularity of $L_1$ regularization over nonconvex regularizers. Without convexity (such as W-shaped regularizer), it is very hard to establish any interesting convergence theory as we do in our paper.

Indeed, in the intersection of nonsmoothness and convexity, piece-wise affine function looks to be the only sensible choice. Also notice that using proximal-update of nonsmooth regularizers (instead of gradient update) avoids any problem due to the non-uniqueness of their (sub-)gradients, which is the main motivation to address with smooth regularizers by Ref[1].

Finally, we agree with the reviewer that adding appropriate guidance (explaining the purpose of each section at their beginnings) will make the paper easier to read. We will be able to add them in the final paper with one extra page allowed for the main text.

We thank the reviewer for recognizing our main contribution on the PAR regularization, the AProx method and proving its last-iterate convergence. We will work on better readability in revising the paper as suggested by the reviewer. Here are answers to the reviewer’s questions.

1. Line 160 the definition of $\text{dist}(w, Q^d)$ is $\text{dist}(w, Q^d)=\min_{v\in Q^d}||w-v||_2^2$, which appears in line 153.
2. We will add relevant fundamentals on quantization, especially on previous work using piecewise affine functions.

• ProxQuant introduced W-shaped regularization, which is piecewise affine but nonconvex, so it is hard to establish interesting convergence properties.
• Several other methods (including BinaryRelax) can be reformulated with equivalent piecewise affine regularizations, but the authors did not recognize the connection. In particular, the BinaryRelax equivalently uses the proximal map in Figure 9(b), which corresponds to a nonconvex piecewise affine regularization.
• Dockhorn et al (2021) focused on the piecewise affine proximal maps and made the connection with piecewise affine regularizations (PARs). But they did not realize that there exists a class of convex PARs that can lead to much stronger convergence guarantees (which is one of the main contributions of our paper).

3. We appreciate the reviewer’s comment regarding the regularization overhead. In our method, the regularization overhead is negligible compared to the cost of gradient computation. Specifically, in AProx, the additional computation arises from the proximal update described in Equation (11). However, this update leverages an explicit proximal mapping, as shown in Equation (7), which can be evaluated efficiently. In the practical implementation (PARQ), we apply LSBQ to determine quantization values by solving a constrained least-squares problem. It is important to note that this step is common across many QAT algorithms and therefore does not introduce any additional overhead specific to our method.

4. We thank the reviewer for suggesting to incorporate other quantization schemes such as LSQ+. We agree that adaptive schemes such as LSQ (Learned Step Size Quantization) and the LSQ+ extension provide more flexibility in using trainable quantization scale and offset parameters. It’s possible to replace LSBQ used in PARQ with LSQ+ for potential performance improvement, which does not impact the convex PAR theory and AProx algorithms convergence properties.

We will incorporate some of the above discussions in preparing for the final submission.

We thank the reviewer for recognizing our paper’s contributions in advancing QAT by bridging gaps between heuristic approaches and theoretical foundations. We address the reviewers comments and questions as follows:

1. “The baselines in the experiment are too old. Why is PARQ not compared with newer methods such as AdaRound[1] or N2UQ[2]?”

• The baselines we use are indeed some of the early works on QAT, especially BinaryConnet/STE, which is still the de facto standard in practice. Most recent works are extensions of BinaryConnect in different ways, relying on the fundamental interpretation of “Straight-Through Estimator.” As we explained in Section 1.1, STE is a misconception we try to correct through the convex PAR framework, and we provide a more principled interpretation to it.

• As QAT attracts more attention in recent years, there are tens of new methods proposed and published including AdaRound and N2UQ. Comprehensive comparison with recent QAT methods is not our intent in this paper, as most recent methods integrate various small additional tweaks beyond the fundamental ideas in order to boost empirical performance. And it can be unfair if we do not equip every method with similar bells and whistles. For example:

  • N2UQ (Nonuniform-to-Uniform Quantization) introduces a particular nonuniform quantizer, it is an alternative scheme for LSBQ we use to generate the quantization targets $Q= \\{ q_1,...,q_m \\}$, which is independent of the PAR and AProx algorithm. We can replace LSBQ in PARQ (Algorithm 1) with N2UQ and reapply the rest (PAR and AProx) to test the performance, but the results are not indicative of what we care most about the effectiveness of PAR or AProx.
  • AdaRound is actually a PTQ (post-training quantization) method, in a quite different category of QAT, see our remarks in Section 1 (second paragraph starting on Lines 39).

• Among the QAT methods beyond BinaryConnect/STE, we choose to compare with BinaryRelax because it follows a similar proximal gradient approach, and the proximal map is also piecewise affine but corresponding to nonconvex regularization as shown in Figure 9(b). See also the third bullet point in our response to Reviewer DkcD's comments.

2. “Can PARQ be applied to large language models? For example, the llama Family?”

• Yes, PARQ is a generic method that can be applied to train any machine learning models, including LLMs. We will try to include experiments on training a (relatively small) Llama model in the final version to demonstrate its applicability.

3. “Reiterate the greatest contribution of PARQ.”

• Our main contributions include: Introducing a class of convex, piece-affine regularizations (PAR) that can effectively induce weight quantization; Developed an aggregate proximal gradient (AProx) method and proved its last-iterate convergence (first of its kind); Provides a principled interpretation of the widely successful heuristic of STE. In summary, we developed a principled approach for QAT, “bridging the gaps between heuristic approaches and theoretical foundations.”

Again, as we state in the paper, the main goal of our experiments is to demonstrate that PARQ has competitive (not necessarily superior) performance against the de facto standard of BinaryConnect/STE. Indeed they become essentially the same algorithm if run for a long time, thanks to our theoretical connection. A comprehensive performance benchmark against recent QAT methods is beyond the scope of this paper, which requires taking care of many additional nuances for a fair comparison.

We thank the reviewer for recognizing our main contributions on convex regularization for inducing quantization and the AProx method with last-iterate convergence. Our response will focus on the experiment results and analysis.

We agree that we can make the discussion on experiments more comprehensive, especially with the one extra page allowed for the final version. In particular, we should emphasize that the main goal of our experiments is to demonstrate that PARQ obtains competitive performance, not necessarily always better than existing approaches such as STE or BinaryRelax (for which we provide novel, principled interpretation). There are several aspects to discuss:

• The experiment results presented are in TEST accuracy, following the convention of the ML community. We developed PARQ as a rigorous optimization/training framework for QAT, but our work does not yet address the generalization capability of the regularization, which is an interesting topic that we aim to investigate in future work.
• The final training losses for different algorithms are also very close, see Figures 10 and 12. Even for training loss, it is hard to guarantee that PARQ is always better than others due to the nonconvex loss functions in deep learning – the results depend on the random initializations and random mini-batches in training. Our convergence guarantees are developed for convex losses, which also imply similar behaviors around a local minimum in the nonconvex case, but in general cannot guarantee a better local minimum. Many relative comparisons in the Tables are statistically insignificant, especially given the small number of runs with random seeds.
• On the other hand, similar performances of STE and PARQ is somewhat expected, as STE can be interpreted as the asymptotic form of PARQ. We are not contrasting two drastically different approaches, rather to argue that the one with a sound principle is as good as the widely successful heuristic (but the principled approach enables guided further investigation). Similarly, BinaryRelax uses a similar proxmap as for PARQ (see Figure 9), but does not correspond to convex PAR (hence less theory support).
• We mainly commented on the very-low-bit cases (1 bit or ternary) due to the observed, relatively large, half-point improvements. The low-bit cases have much less freedom compared with more bits so it may be relatively easier to reveal differences between different local minima. The gradual evolution of PARQ from piecewise-affine soft quantization to hard quantization may help the training process to be more stable and more likely to converge to better local minima (See our comments in Section 6). Again there is no guarantee that this happens for every model we try (especially with a small number of trials).

We agree that our experiment comparisons are limited. However, it is not our intent to give a comprehensive comparison against many recent QAT methods, as many recent methods integrate with small additional tweaks beyond the fundamental ideas in order to boost empirical performance, and it can be unfair if we do not equip every method with similar bells and whistles. We limit our comparison to BinaryConnect (STE) and BinaryRelax because they have direct connections with our method as explained above. In addition, STE is still the de facto standard in practice despite many new methods and variations proposed.

• On "why the best results under each setting in Tables 1-3 are not all highlighted in bold"? We only highlight the best results that look to be statistically significant, meaning the difference between the means are apparently larger than their standard deviations.

Answers to Questions for Authors:

• For Questions 1 and 2, please see our itemized explanations/discussions above. In addition, we conjecture that more bit-widths and more complex networks may have the advantage of being over-parametrized, leading to very small differences between the local minima found by different methods.

• For Question 3, ProxQuant (Bai et al. 2018) proposed to use the W-shaped regularizer (nonconvex) and indeed use the Prox-SGD method. As we explain in Section 3, Prox-SGD will NOT produce a meaningful quantization effect as training evolves over time ($t\to\infty$). In order to fix this empirically, as described in their implementation, the regularization parameter $\lambda$ is changed to be growing linearly $\lambda_t = \lambda \cdot t$, without principle/theory justification. It turns out that it is consistent with our AProx algorithms, where an increasing regularization strength should be applied according to our theory. Then their actual implemented algorithm is similar to PARQ, except that they use a nonconvex PAR of W-shaped. So in the end, it is more close to BinaryRelax, which we included in our experiments.