Matryoshka Quantization

We thank the reviewer for their thorough feedback. Below are our responses to the comments/questions:

Comparison with SOTA int2 methods: We would like to clarify that the main goal of the work is to propose an adaptive quantization method that generates a single model that can do well at any precision and provide accurate mixed precision models (Mix’n’match). Int2 result demonstrates that the technique can potentially help train a significantly more accurate low bit-width model, but it is not the main goal, so the baselines are also set accordingly.

For Mix’n’match, we compare against the baseline of learning quantized models independently and then stitching layers together to produce a model with the same per layer quantization. The table given below shows that MatQuant based Mix’n’match can be 21% more accurate than a Mix'n'Match between separately trained baselines.

AQLM is a complementary technique to MatQuant, so AQLM+MatQuant might provide a more accurate 2bit model than AQLM. However, as mentioned above, that is orthogonal to the goal of the work and we leave it for future investigation. The Matryoshka-style nesting is possible across multiple axes like the number of code books for each group and the number of elements in each code book (i.e., 2^bits).

Table: Mix'n'Match on MatQuant + OmniQuant for Gemma-2 9B. Mix'n'Match with MatQuant sub-models does substantially better than a Mix'n'Match between baseline models trained explicitly for 2 and 4 bits.

Relation to NF Quantization Data Type: MatQuant does not explicitly try to push the weight distribution towards an information-theoretic optimal quantization scheme. The fact that we slice bits in MatQuant to obtain sub-models incentivizes gradient descent to carefully optimize the MSBs. Since the 2-bit sub-model corresponds to the top-2 MSBs, MatQuant optimizes the top-2 MSBs to maximize the int2 sub-model’s quality and thus ends up pushing the peak of the weight distribution towards the right, and in the process does a more uniform allocation to the higher buckets. This is evident in Figure 1 (c), where we can see that the peaks for a model trained with MatQuant are lower than that for the baseline and the weight distribution in the case of MatQuant is a bit more uniform.

Datasets other than C4: Popular quantization papers such as GPTQ, QuIP, and SpinQuant utilize C4. To be consistent with popular literature, we use C4 for our experiments.

We hope that the rebuttal clarifies questions raised by the reviewer. We would be very happy to discuss any further questions about the work, and would really appreciate an appropriate increase in score if reviewers’ concerns are adequately addressed.

We thank the reviewer for their insightful review. Below are our responses to the comments/questions:

Training memory/compute requirements: For both QAT and OmniQuant, MatQuant can be up to 30% cheaper than the training three separate baselines (one for 2-bits, one for 4-bits, and one for 8-bits). By recomputing some of the intermediate tensors generated during the forward pass when needed for gradient computation, MatQuant can run in the same memory requirements as a single baseline run. So overall from a training cost point of view, MatQuant uses up to 30% less GPU hours, and is more effective than having separate training runs for different bit-widths.

Training curves: Our experiments indicate that MatQuant requires less samples/iterations to converge to the same training perplexity. Specifically,

1. 2-bit QAT on Gemma2 9B: MatQuant and S.P. Matquant achieves same training perplexity with 20M tokens as the baseline with 100M tokens
2. 2-bit OmniQuant on Gemma2 9B: MatQuant’s layerwise reconstruction error at 8M tokens and S.P. MatQuant’s layerwise reconstruction error at 5.8M matches that of baseline at 20M tokens.

Inference cost: A MatQuant trained 2/4/8-bit model will have the same inference cost as the baseline (QAT or OmniQuant) 2/4/8-bit models, i.e., inference budgets for homogeneous precisions is exactly the same as that of the baseline since we do not change the base quantization algorithm. With Mix’n’Match, we can further trade off quality with latency catering to several serving environments since our models lie on the pareto-optimal curve for memory-vs-quality.

Plugging MatQuant on the latest methods: we thank the reviewer for referring the latest QAT style techniques. Most SOTA QAT methods build on top of the baseline QAT setup by either adding auxiliary losses or distillation or some forms of pretraining. So we adopted a baseline general QAT setup in our experiments to demonstrate the core advantage of MatQuant. MatQuant is complementary to the above mentioned methods of auxiliary losses etc. Furthermore, as we have shown, MatQuant works with a wide variety of quantization techniques (QAT, OmniQuant). So we believe MatQuant can be then combined with the latest QAT style methods.

Scaling MatQuant to larger models: We performed preliminary experiments on Gemma-2 27B with QAT. As shown in the table below, we find that our observations from 2B and 9B scales hold for the 27B model as well. That is, MatQuant’s quantized models are quality neutral (or better by up to 5% for 2-bit) for all the bit widths that we trained for as well as for the interpolated bit-widths 3 and 6. We will add these results to the next version of the manuscript.

Table: MatQuant + QAT on Gemma-2 27B performs on par with the baseline for int4 and int8 and significantly outperforms it for int2. Also, the int3, int6 models obtained for free via interpolation perform comparably to the explicitly trained baselines.

We hope that the rebuttal clarified your questions about the paper and hopefully leads to an even more positive evaluation of the work. We are happy to further discuss if you have any questions about the paper.

We thank the reviewer for their insightful review. Below are our responses to the comments/questions:

Training memory/compute requirements: For both QAT and OmniQuant, MatQuant can be up to 30% cheaper than the training three separate baselines (one for 2-bits, one for 4-bits, and one for 8-bits). By recomputing some of the intermediate tensors generated during the forward pass when needed for gradient computation, MatQuant can run in the same memory requirements as a single baseline run. So overall from a training cost point of view, MatQuant uses up to 30% less GPU hours, and is more effective than having separate training runs for different bit-widths.

Training curves: Our experiments indicate that MatQuant requires less samples/iterations to converge to the same training perplexity. Specifically,

1. 2-bit QAT on Gemma2 9B: MatQuant and S.P. Matquant achieves same training perplexity with 20M tokens as the baseline with 100M tokens
2. 2-bit OmniQuant on Gemma2 9B: MatQuant’s layerwise reconstruction error at 8M tokens and S.P. MatQuant’s layerwise reconstruction error at 5.8M matches that of baseline at 20M tokens. We will add the training curves to the next version of the manuscript.

Single Precision (S. P.) MatQuant vs Baseline:

1. S.P MatQuant first quantizes the model to 8-bits and then slices the 2 MSBs. However, the baseline directly quantizes the model to 2 bits.
2. The scaling factor and the zero point in S.P. MatQuant are derived from the 8-bit quantized model. Here, the model has 6 additional bits to better optimize the scaling factor and the zero point since the quantized weight corresponds to only the first two bits and the model can use these additional 6 bits to improve the overall quality. We hypothesize this additional entropy (overparameterization) gives S. P. MatQuant an edge over the baseline.

Table 4: [8, 4, 2, 8 -> 4;2]: [8, 4, 2, 8 -> 4;2] means the standard cross entropy (or layer-wise reconstruction) is applied to the 8-bit, and the 4 and 2-bit sub models. In addition to this, the 4 and the 2-bit sub models are distilled from the 8-bit model’s logits (or layer output). The notation is present in the caption of Table 4, but we will further clean it up to make it easier to parse.

We hope that the rebuttal clarified your questions about the paper and hopefully leads to an even more positive evaluation of the work. We are happy to further discuss if you have any questions about the paper.