Task-Circuit Quantization: Leveraging Knowledge Localization and Interpretability for Compression

Thank you for your review and for appreciating our “broad” experimental section, “one hyper-parameter” setup, and building a “conceptually fresh bridge” between mechanistic-interpretability and quantization. We proceed to address the remaining comments in your review:

Scaling to Larger Models. Thanks for this suggestion -- we have added an experiment to show that TCQ continues to outperform GPTQ [2] and SliM-LLM [3] in the 4.6x larger Qwen2.5-32B-Instruct model. We observe that TCQ recovers 87.93% of the unquantized model’s performance at 2-bit quantization, reducing model size by 8x and outperforming Slim-LLM by 10.00% and GPTQ by 21.96%. Additionally, TCQ quantization remains tractable on a single 4x NVIDIA RTX A6000 48G node in 6.93 hours. Here we use 16 bit gradients for TCQ.

Performance on Qwen2.5-32B-Instruct

Q1. Compute Footprint. We briefly covered quantization compute footprint in Appendix A. Based on your suggestion we expand the comparison to GPTQ, SliM-LLM and TCQ for 7B and 32B in the table below and will include it in our final paper. We compute these numbers for all methods in the setting of 2-bit quantization conditioned on MMLU_MCQA. TCQ quantization is significantly faster than SliMLLM for both 7B and 32B models.

While gradient capture is expensive, there are several reasons why we believe it is tractable in the context of TCQ. We also conduct new experiments to show methods we can use to drastically reduce the cost of gradient capture by 8x without affecting performance.

• The number of datapoints used for gradient capture can be reduced by 8x without affecting performance. See Table A below.
• Gradients only need to be captured once to freely quantize at varied % outliers and base bitness.
• In Appendix D, line 708-733, we detail a memory saving trick for gradient capture that offloads the last stage of gradient computation to the CPU, and suggest other possible engineering improvements.
• Previous works such as SqueezeLLM [5] and GWQ [4], which we compare against both require gradient capture and assume reasonable cost.

Table A: varying the number of calibration examples for gradient capture when compressing Llama-3-8B-Instruct to 3 bits on MMLU.

Q2. Inference Efficiency. In terms of inference efficiency we note that our framework produces weights in the format of SPQR’s setup [1], and inference times should be the same by using SPQR’s kernel. We note this in Appendix D.2 line 758. We view our main contributions as the TCQ metric and task-specific quantization method, and thus do not evaluate using a kernel. Details on inference speeds can be found in Table 4 of the SPQR paper [1].

Q3. Outlier Ratio Sensitivity. In Figure 3, we show a pareto plot where we evaluate TCQ with higher outlier percentages, showing the accuracy memory tradeoff as you increase the number of outliers used. We note that 2% outliers correspond to approximately .75 additional bit-width. In general, a higher outlier ratio is not advisable in terms of memory use. Noted in line 542, average bit-width can be calculated as bit-width = N − Nr + 32r where r is the ratio of outliers we use. Appendix F, Figure 4, contains an additional pareto comparison where we sweep outlier ratio from .5% to 2%.

[1] Dettmers et. al., SpQR: A SparseQuantized Representation for Near-Lossless LLM Weight Compression, June 2023. URL http://arxiv.org/abs/2306.03078. arXiv:2306.03078.

[2] Frantar et. al., GPTQ: Accurate PostTraining Quantization for Generative Pre-trained Transformers, March 2023a. URL http://arxiv.org/abs/2210.17323. arXiv:2210.17323 [cs].

[3] Huang et. al., SliM-LLM: Salience-Driven Mixed-Precision Quantization for Large Language Models, May 2024b. URL http://arxiv.org/abs/2405.14917. arXiv:2405.14917 [cs].

[4] Shao et. al., GWQ: Gradient-Aware Weight Quantization for Large Language Models, December 2024b. URL http://arxiv.org/abs/2411.00850. arXiv:2411.00850 [cs].

[5] Kim et. al., SqueezeLLM: Dense-and-Sparse Quantization, June 2024. URL http://arxiv.org/abs/2306.07629. arXiv:2306.07629 [cs].

Thank you for your review and for pointing out that TCQ is a “Interesting and reasonable method to conduct quantization”. We proceed to address your comments below.

Q1a: GWQ comparison. In section 4.5 and Table 5 of the paper, we do in fact mention and compare against the GWQ baseline [4]. We labeled this as “Sample Absolute Gradient” in Table 5 due to the fact that GWQ’s metric is equivalent to the gradient term of the TCQ metric. Thus, we already compare our TCQ outlier selection metric in an apples-to-apples setting to GWQ’s outlier selection metric. The TCQ metric outperforms GWQ metric (Sample Absolute Gradient) by 6.49% at 2 bits and 1.81% at 3 bits, evaluated on MMLU accuracy. Thanks for the comment, we will clarify this in the final version.

Q1b: QAT comparison. We propose a Post-Training-Quantization method and following SqueezeLLM [1], GPTQ [2], and Slim-LLM [3], view a comparison to QAT as out of scope. Quantization-Aware-Training’s increased cost -- requiring significantly more data points and fully training the model for many steps -- means that past work has treated it as a separate category.

Q2: QAL and MSG concrete examples. Concretely, we expect that when QAL is used, the average distance to gridlines will increase. Intuitively, given the same gradient importance, it is more valuable to preserve outliers far away from gridlines as they will be more heavily impacted when quantized. We compute this average distance metric and show that this is indeed the case.

This new experiment was conducted on Llama-3-8B-Instruct, comparing the average distance to gridlines for outliers identified with just the gradient to those identified with QAL. In the final version of the paper, we will plan to include this result in an expanded discussion on the concrete impacts of QAL and MSG.

Q3: Perplexity results. Thanks for pointing this out – we in fact include these results in Appendix G, Table 9, where we show perplexity results indicating that TCQ outperforms other methods in terms of language modeling across wikitext2, c4_new, and ptb_new datasets. We also point to this table in section 4, Datasets for Quantization Conditioning and Evaluation, line 207. We’ll move this pointer to the start of section 4 in the final version.

We include an updated version of these results on Llama-3-8B quantized with C4 as the conditioning dataset here for quick reference:

Note that this table differs slightly from Table 9, in that it includes Slim-LLM (which we initially excluded due to its long runtime).

[1] Sehoon Kim, Coleman Hooper, Amir Gholami, Zhen Dong, Xiuyu Li, Sheng Shen, Michael W. Mahoney, and Kurt Keutzer. SqueezeLLM: Dense-and-Sparse Quantization, June 2024. URL http://arxiv.org/abs/2306.07629. arXiv:2306.07629 [cs].

[2] Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. GPTQ: Accurate PostTraining Quantization for Generative Pre-trained Transformers, March 2023a. URL http://arxiv.org/abs/2210.17323. arXiv:2210.17323 [cs].

[3] Wei Huang, Haotong Qin, Yangdong Liu, Yawei Li, Xianglong Liu, Luca Benini, Michele Magno, and Xiaojuan Qi. SliM-LLM: Salience-Driven Mixed-Precision Quantization for Large Language Models, May 2024b. URL http://arxiv.org/abs/2405.14917. arXiv:2405.14917 [cs].

[4] Yihua Shao, Siyu Liang, Zijian Ling, Minxi Yan, Haiyang Liu, Siyu Chen, Ziyang Yan, Chenyu Zhang, Haotong Qin, Michele Magno, Yang Yang, Zhen Lei, Yan Wang, Jingcai Guo, Ling Shao, and Hao Tang. GWQ: Gradient-Aware Weight Quantization for Large Language Models, December 2024b. URL http://arxiv.org/abs/2411.00850. arXiv:2411.00850 [cs].

We thank you for your review, and are glad that you thought our paper was “well-structured” and had a “motivated methodology” and “comprehensive evaluation and ablation” studies.

Thank you for your careful review of our paper, and for highlighting our “simple and intuitive” approach to an “interesting and timely problem”.

Mechanistic Interpretability Connection. To clarify our connection to Mechanistic Interpretability, we trace the line of MI ideas motivating our work. Activation Patching, originally introduced in ROME [1] for localization and model editing, was applied for localization of relevant model components. Attribution Patching [2] was introduced as a gradient based reformulation of Activation Patching with significant efficiency improvements, and applied to automated circuit discovery [3, 4]. We take the efficiency insight and formulation from Attribution Patching and apply it for localization of task relevant weights, and define what we find as “weight circuits” to separate them from activation centric definitions of circuits. Nevertheless, we recognize that circuits can refer to manually discovered circuits and more intricate edge based processes and will include an extended discussion of this in our final paper. To that end, we would appreciate any elaboration on examples of “second-order interaction” that might be helpful to consider.

Question 1: which systems “share quantization”. All methods are conditioned in Table 1. Specifically, all methods we compare against depend on calibration data for quantization, and have their conditioning data replaced for each evaluation dataset. Additionally, with the exception of the uniform precision GPTQ, all methods use calibration data to determine their mixed precision outlier selection or layout, which we also replace. For example, for Slim-LLM, each seed in each square of Table 1 consisted of a ~12-18 hour run where calibration was conditioned on the specific dataset and specific seed. In Table 3, no additional conditioning of TCQ was done beyond that for other methods, in order to show the generality of our method. We thank you for pointing this out and will emphasize our apples-to-apples conditioning setup in the final version of this paper.

Question 2a: MSG for moderate weights. We thank you for the question. By “moderate weights”, we assume you are referring to weights where |W_ij| is small (please let us know if meant something else by this), in which case upweighting by |W_ij| would not fix approximation errors from QAL. However, as 0 is always a gridline and our quantization implementation uses a simple linear search strategy to shrink the quantization scales (discussed in lines 748-751), the majority of weights with small |W_ij| are already fairly well captured by gridlines, especially compared with weights with large |W_ij|. We briefly discuss this in Appendix D, Expanded Description of TCQ, lines 691-702.

Question 2b: exactness of QAL approximation. While we affirm that using a small percent of outliers makes the QAL estimate “informative”, it is not perfect. We note that we do not preserve the same ratio of outliers for each weight matrix. For example, when quantizing Qwen2.5-7B-Instruct, the 27th mlp down projection preserves 5.226e+04 outliers, while the 18th mlp down projection layer preserves an order of magnitude more outliers at 3.604e+05. The distribution is also skewed in a channel wise manner. We observe that in channels where the number of outliers is large, the gridlines can be altered. In those cases, MSG acts as a stabilizing agent by prioritizing a gridline-agnostic importance.

[1] Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and Editing Factual Associations in GPT, January 2023. URL http://arxiv.org/abs/2202.05262. arXiv:2202.05262 [cs].

[2] Neel Nanda. Attribution Patching: Activation Patching At Industrial Scale, 2022. URL https://www.neelnanda.io/mechanistic-interpretability/attribution-patching.

[3] Janos Kramar, Tom Lieberum, Rohin Shah, and Neel Nanda. AtP*: An efficient and scalable method for localizing LLM behaviour to components, March 2024. URL http://arxiv.org/abs/2403.00745. arXiv:2403.00745 [cs].

[4] Aaquib Syed, Can Rager, and Arthur Conmy. Attribution Patching Outperforms Automated Circuit Discovery, November 2023. URL http://arxiv.org/abs/2310.10348. arXiv:2310.10348 [cs].