Matryoshka Quantization
Authors: Pranav Ajit Nair
Puranjay Datta
Jeff Dean
Prateek Jain
Aditya Kusupati
Venue: ICML 2025 poster (2025)
Keywords: Matryoshka
Quantization
LLMs
Deployment
Primary area: deep_learning->algorithms
TL;DR: Matryoshka Quantization (MatQuant) trains a single model that can operate at multiple bitwidths (e.g. Int8, Int4, Int2) simultaneously by leveraging the nested structure of integer data types, outperforming independent quantization of the same models

Abstract: Quantizing model weights is critical for reducing
the communication and inference costs of large
models. However, quantizing models – especially
to low precisions like int4 or int2 – requires a
trade-off in model quality; int2, in particular, is
known to severely degrade model quality. Consequently, practitioners are often forced to maintain
multiple models with different quantization levels or serve a single model that best satisfies the
quality-latency trade-off. On the other hand, integer data types, such as int8, inherently possess
a nested (Matryoshka) structure where smaller
bit-width integers, like int4 or int2, are nested
within the most significant bits. Leveraging this
insight, in this paper, we propose Matryoshka
Quantization (MatQuant), a novel multi-scale
quantization technique that alleviates the aforementioned challenge. This technique allows us to
train and maintain a single quantized model but
serve it with the precision demanded by the deployment. Furthermore, leveraging MatQuant’s
co-training and co-distillation, int2 precision models extracted by MatQuant outperform standard
int2 quantization by up to 4% and 7% with OmniQuant and QAT as base algorithms respectively.
Finally, we demonstrate that by using an extra bit
to represent outliers, a model with an effective
precision of 2.05-bit improves further by 6% with
OmniQuant as the base algorithm.

Matryoshka Quantization
Pranav Nair Puranjay Datta
Jeff Dean
Prateek Jain'
Aditya Kusupati
Abstract
Achiam al., 2023) in a large number of scenarios. Due to auto-regressive nature of LLMs, decode latency tends to dominate inference cost. Decode latency itself is dominated by communication cost of transferring model weights from high-bandwidth memory (HBM) to the SRAM or due to transferring weights/activations in a distributed cluster.
Quantizing model weights is critical for reducing the communication and inference costs of large models. However, quantizing models especially to low precisions like int4 or int2 requires a trade-off in model quality; int2, in particular, is known to severely degrade model quality. Conse- quently, practitioners are often forced to maintain multiple models with different quantization lev- els or serve a single model that best satisfies the quality-latency trade-off. On the other hand, in- teger data types, such as int8, inherently possess nested (Matryoshka) structure where smaller bit-width integers. like int4 int2, are nested within the most significant bits. Leveraging this insight, in this paper, we propose Matryoshka Quantization (MatQuant), a novel multi-scake quantization technique that alleviates the afore- mentioned challenge. This technique allows us to train and maintain single quantized model but serve it with the precision demanded by the de- ployment. Furthermore, leveraging MatQuant's co-training and co-distillation, int2 precision mod- els extracted by Mat Quant outperform standard int2 quantization by up to 4% and 7% with Omni- Quant and QAT as base algorithms respectively. Finally, we demonstrate that by using an extra bit to represent outliers, model with effective precision of bi05-bit improves further by 6% with OmniQuant as the base algorithm.
Quantizing weights and/or activations can significantly re- duce the overall communication load and is, therefore, one of the most popular techniques for reducing inference costs (Dettmers et al., 2022). While floating-poin represen- tations are standard for training, integer data types such as int8, int4, and int2 are appealing alternatives for inference. However, current methods for quantizing to these varying integer precisions typically treat each target precision as an independent optimization problem, leading to a collection of distinet models rather than single, versatile one. Fur- thermore, quantizing to extremely low precisions like int2 is known to be highly inaccurate. In this work, we pose the question whether both of the above challenges can be addressed; that is, can we train a single model from which we can extract multiple accurate lower-precision models? We answer this question in the affirmative by introducing Matryoshka Quantization (MatQuant), novel multi-scale training method that leverages the inherent nested (Ma- tryoshka) structure (Kusupati et al., 2022) within integer data types (Figure la). Specifically, slicing the most signifi- cant bits (MSBs) of an int8 int8-quantized weight can directly yield an int4 or int2 model. Existing quantization techniques often neglect this structure, which limits the potential for multi-scale adaptable models operating at various bit-widths with optimal performance.
Instead, MatQuant simultaneously optimizes model weights across multiple precision levels (e.g., int8, int4, int2). At a high level, we represent each model parame- ter at different precision levels using shared MSBs, and then jointly optimize the loss for each precision level. This allows us to develop a single quantized model that can effectively operate at any of the chosen bit-widths, offer- ing spectrum of nccuracy-vs-com options. MatQuant is a ceneral-purpose technique, applicable to most learning- based quantization methods, such as Quantization Aware Training (QAT) (Jacob et al., 2018) and OmniQuant (Shao et al., 2023).
1. Introduction
Due to their impressive performance, there is a strong push to deploy deep learning models, particularly large language models (LLMs) (G Team et al., 2024; Dubey et al., 2024;
Eqqual contribution GGooge DeepMind 'Google Research. Correspondence to: Pranav Nair <rranavajitnair@ google.com>> Puranjay Datta <puranjaydatta hetpoyyddatt@gggle.oooo @oomole..om..>..............mm.. Jeff Dean goef@@gogle.coom
Prateek Jain <prajain @owogee.cooo..... Aditya Kusupati <kusupati@ google.com>.
Proceedings of the 4nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Matryoshka Quantization
11 01 1001
Gemma-2 98 Weight Distribution for OmniQuant
12
Baseline MatQuant
int8
100
Crsit Cuve La.bie
00.8
QuaaQuaas
0.
(a)
0.0g
256
Quantized buckets
Gemma
Bape
Int4
Soceline VatQuart
Int2
70
Mantiant
Task Average
Frequency
Frequet
80
MatQuant
Mattuunt-Inteer.
50
Baseline
1.0
MinMax
40
Sliced inta
0%
Quantized buckets
Quantized buckets
Effective bits per FFN parameter
(c)
(h)
Figure 1. (a) MatQuant is multi-scale quantization training technique using the inherent Matryoshka structure of int8 int4 int2. (b) Empirical gains of Mat MatQuant on downstream tasks, especially > 8% for int2, on Gemma-2 9B with OmniQuant. (c) The right-shifted quantized weight distribution as a consequence of MatQuant's training mechanism that maximises accuracies across all precisions.
We demonstrate the efficacy of Mat Muanu when applied to quantizing the Feed-Forward Network (FFN) parame- ters of standard LLMS (Gemma-2 2B, 9B, and Mistral 7B) (Vaswani et al., 2017) typically, FFN is the main latency block hence the focus on improving the most signifi- cant component's 's latency. Our results show that MatQuant produces int8 and int4 models with comparable accuracy to independently trained baselines, despite the benefit of shared model parameters. Critically, the int2 models gen- erated by MatQuant significantly outperform their indi- vidually trained counterparts, with 4% higher accuracy on downstream tasks (Figure 1b). We also extend MatQuani to quantize all weights of a Transformer layer. In Figure Ic, we find that quantizing with MatQuant shifts the quantized weight distribution toward higher values, contributing to improved int2 performance. Finally, in Section 7, we also demonstrate that using an extra bit to represent outliers sig- nificantly boosts the performance for our sliced int2 models.
2. Related Work
Model weight quantization is extremely powerful and prevalent technique for making resource- resource-intensive neural networks suitable for deployment constraints - especially modern-day LLMs. Quantization algorithms can be cate- gorized as either learning-free free or learning- based. Learning- free methods use limited data to calibrate model parameters without relying on gradient descent. Learning ods, however, utilize gradient descent to update either model parameters or auxiliary parameters to aid in quantization.
meth-
Learning-free Quantization Methods. Naive quantiza- tion methods, such as MinMax, absmax, and zero-point quantization, aim to directly map the range of model weights to the target bit-width - see (Dettmers et al., 2022) for de- tailed background. Dettmers al. (2022) further improved this by identifying the need to handle outliers with higher precision than the rest of the model weights. The core principle of more recent learning-free quantization methods remains similar while improving various aspects of it and using small amounts of data for calibration. For example, GPTQ (Frantar et al., 2022) improves upon min-max quanti- zation by iterating over all the coordinates, quantizing them one at a time, and updating the remaining full-precision coordinates minimize the layer layer-wise activation recon- struction error. AWQ (Lin et al., 2023), SmoothQuant (Xiao et al., 2023), and AffineQuant (Ma et al., 2024) scale the weights and activations to reduce outliers, thus making them easier to quantize. QulP (Chee et al., 2024), Frame- Quant (Adepu al., 2024), and QuaRoT (Ashkboos et al.,
Beyond improving chosen precision performance, MatQuant allows for seamless extraction of interpolative bit-widths, such as intó and int3. Mat Quant. also admits a dense accuracy-vs-cos! trade-off by enabling layer-wise Mix'n 'n' Match of different precisions. Therefore, even if the hardware only supports int4 and int2, it's possible to serve models various effective precisions, tailored to the deployment environment. Overall, MatQuant and its variants present significant step toward developing multi-scale models with high flexibility and performance, pushing the boundaries of low-bit quantization for efficient
LLM inference.
2

---

Matryoshka Quantization
2024) multiply the weights and activations by orthonormal matrices before quantizing reduce the number of outliers. SqueezeLLM (Kim et al., 2024) uses clustering to obtain the optimal buckets for quantization, and CDQuant (Nair & Suggala, 2024) improves upon GPTQ by greedily choosing the coordinates to descend along. While learning-free ffree meth- ods are inexpensive and work well at higher bit-widths, they are often suboptimal in the low-precision regime, which benefits greatly from learning- based techniques.
types, especially in the context of quantization, which aids in the deployment of resource-intensive LLMs. Following this idea, Yu et al. (2019) have successfully trained sin- gle model that can do well at any precision. However, the experiments were limited to ConvNets and small Neural Networks. In this paper, we extend the idea of nested pre- cision to LLMs and show that it indeed works at scale. We also show that, for the first time, our models are quality neutral for intermediate precisions such as int3 and intó that we never trained for, and densely span the accuracy-ve -v-bbis trade-off. In Section 5.3, we show that even to train mod- els for a fixed target precision, having loss over the sliced bits of an 8-bit model does better than training a model explicitly for that precision, indicating that Mat Quant is a fundamentally better way to do low-bit quantization.
Learning-basee Quantization Methods. Quantization Aware Training (QAT) (Jacob et al., 2018; Abdolrashidi et al., 2021) is a logical approach to ensure that models easy to quantize during inference while retaining high accuracy. However, because QAT involves updating all the model parameters, its adoption for LLMs has been lim- ited. Several recent works improve the performance and efficiency of QAT. LLM-QAT (Liu et al., 2024a) and Bit- Distiller (Du et al., 2024) enhance QAT with knowledge dis- tillation from the full-precision model, EfficientQAT (Chen al., 2024) minimizes the block-wise reconstruction error before performing end-to-end training. This significantly reduces the time it takes for QAT to converge. On the other hand, some techniques significantly reduce the overhead by learning only the auxiliary parameters, such as scaling factors and zero- pero-points, that aid in quantization instead of updating the actual weight matrices. For example, Omni- Quant (Shao al., 2023) does not update the model pa- rameters; instead, it learns additional scales and shifting parameters (that aid with quantization) through gradient de- scent over the block-wise reconstruction eтror and achieves better accuracy than most QAT techniques. Likewise, Spin- Quant (Liu et al., 2024b) uses gradient descent to learn its rotation matrices. This class of learning-based quantization techniques (OmniQuant, SpinQuant, etc.) is widely adopted due their appeal of achieving QAT-level accuracy fraction of the cost.
3. Matryoshka Quantization
We introduce MatQuant, a general-purposs multi-scale training technique that works seamlessly with popular learning-based quantization methods such as Quantiza- tion Aware Training (QAT) (Jacob et al., 2018) and Om- niQuant (Shao et al., 2023). As long as the model or auxiliary parameters are optimized with gradient descent, MatQuant's multi-scale training technique can be used across chosen bit-widths, leveraging the inherent nested structure of integer data types. In this section, we will elab- orate on the preliminaries behind QAT and OmniQuant, alongside our novel proposed approach, MatQu
3.1. Preliminaries
3.1.1. QUANTIZATION AWARE TRAINING
Quantization Aware Training (QAT) learns a c-bit quantized model by optimizing for the end-to-end cross entropy loss using gradient descent. It uses the quantized weights for the forward pass and straight through estimator (STE) (Bengio et al., 2013) to propagate gradients through the quantization operator during the backward pass.
Multi-scale Training. Training across multiple data scales (resolutions) was heavily popularized in computer vision for both recognition and generation (Adelson al., 1984; Lin et al., 2017; Denton et al., 2015). More recently, the paradigm of multi-scale training has shifted to mod- els (Rippel et al., 2014; Yu et al., 2018; Kusupati et al., 2022; Devvrit et al., 2023), where the data remains the same, and models of varying capacity, all nested within one large model, are trained jointly. This joint, nested (Matryoshka- style) learning with varying model sizes results in a smooth accuracy- securacy-vs-compute vs-compute trade-off and is beneficial in many downstream applications and real-world deployments. How- ever, the most obvious structure with a nested nature is the bit structure of the integer data type. Given the suc- cess of multi-scale training for inputs, outputs, and model weights, it imperative to explore it further for integer data
To mathematicalyy formulate QAT, we define MinMax quan- tization of a real-valued vector w in c bits as follows:
w + x (1=t
Qum(w, e) = clamp
0, 2 -
(1)
max(w)
min (w)
min (w)
2c
where QMM(W, c) is the c-bit quantized version of w, is the scaling factor and z is the zero point.
Let WF represent weights of a Transformer LLM and let D = {( vı),.... (IN.yN)} be a labelled dataset where x and y, represent the input and output respectively. With
3

---

Matryoshka Quantization
LCE as the cross entropy loss, the optimization of QAT is:
(2022), MalQuant optimizes the quantization loss for sev- eral target bit-widths jointly.
mim 1 Σ Lce (F QuM (WF. c)), We N
(2)
To have a single model for various integer precisions, we nest smaller bit-widths into large ones - leveraging the in- herent Matryoshka nature of the integer data type. So, if we want to extract a r-bit model from c-bit model (0 <r < c), we can just slice out the most significant bits (MSBs) using right shift, followed by left shift of the same order. Formally, the r()".r) operator slices the most significant bits from a c-bit quantized vector gf:
sE[N]
where represents the LLM's forward pass.
3.1.2. OMNIQUANT
OmniQuant, unlike QAT, does not update the model pa- rameters. Instead, it learns additional scaling and shift- ing parameters through gradient descent over layer-wise L2 error reconstruction. These auxiliary parameters aid with quantization. Similar QAT, OmniQuant also uses straight through estimator during optimization. However, unlike QAT, OmniQuant operates with limited data, making it much more attractive for resource-scare settings.
(1200
g 20-
Si r) clamp
-0,2- * 2
0,,
(6)
Note that clamp(-) is required to curtail overflows generated by rounding. More details can be found in Appdendix A. Once we have this structure, we can optimize for several precisions by slicing the MSBs from the largest bit- width we are optimizing for. Let R = (ri, T2,.... TK) be the bit- widths we want to optimize for, represent the quantiza- tion function of the base algorithm (i.e., any bearning-based quantization scheme), L--) represent the loss function per- taining to the base algorithm, F(-) represent the forward pass required to compute the loss, 0 represent the set of model/auxiliary parameters we are optimizing for and let WF represent the model parameters. MatQuant's 's overall objective can be formulated as follows:
OmniQuaatt adds two learnable scales, 7 and B3 to MinMax quantization as follows:
w + +2 +=|.0,22 0,
Qumai (w, c) = clamp
2 -1011 B min(un)) (w)
a= Y. -=mxxxx w) - B min (w)
2
2c
OmniQuant also adds another set of learnable shifting and scaling parameters to the FFN's affine projections as follows:
min 1 N
E EE x CF((S(((()))))))))))))) (E(N)TER FER
XW +b + ((X-
- 2 ())2)))Q2 QQmmai (Wos)+b+8-N (4)
where 3/1 = y for QAT and v/ = F. X/) for Omni- Quant, and x² = x, for QAT and at = Xi for OmniQuant. X, is the loss reweighing factor for bit-width r.
where X E Raxd is the input to the affine transformation,
W E Rdddd is the linear projection associated with the affine transformation, b e R- is the bias vector, e R2 and s E Rd are learnable shift and scale parameters respectively.
In this work, default to training MatQuant with three bit-wi, R = {8, 4, 4. and subsequently perform a linear search over A, This process aims to optimize performance such that the model performs well across all targeted preci- sion levels. Further, while the focus of this paper is primarily on integer data types, we discuss the possibility of extending Mat Mal.Quatt floating-point representations in Section 5.5.
With the goal of optimizing the layer-wise L.2 error (where layer consists of an Attention block followed by an FFN block), OmniQuant's ' overall objective can be portrayed as follows:
min [[R(W))) Xi) F1(Qummi (Wz), XIIIB 33,66
(5)
A key point to note is that MatQuant primarily alters the quantized weight distributions across precision levels com- pared to the base quantization algorithm (OmniQuant or QAT). Figure Ic illustrates the differences in the quantized weight histograms obtained with and without Mat Quant on Gemma-2 9B using OmniQuant. Upon close observation, we find that all the distributions of Mat Quant are shifted to the right; that is, weights quantized with MatQuant tend to use more higher-valued weights. While this might not signif- icantly impact int8 or even int4 models, int2 models benefit from utilizing more of the possible quantized weights com- pared to the baseline. Because int2 favors higher-valued weights, this effect propagates to higher-valued weights for
where fi(-) represents the forward pass for single layer W/ represents the layer parameters and Xı represents the layer's input. Note that the above objective is optimized independently for each of the L Transformer layers.
3.2. MatQuant
MatQuant is general purpose framework to develop single model that can do well at any precision. It is a multi- scale training technique that works with most learning- based quantization schemes like QAT and OmniQuant discussed earlier. At its core, taking inspiration from Kusupati et al.
4

---

Matryoshka Quantization
Table 1. Mat Quant with OmniQuant across Gemma-2 2B, 9B and Mistral 7B models. MatQuant performs on par with the baseline for int4 and int8 while significantly outperforning it for int2. Even the int3, intó models obtained for free through interpolation from MalQuant perform comparably to the explicitly trained baselines. Task Avg. is average accuracy on the evaluation tasks (1) while log pplx (perplexity) is computed on C4 validation set (+).
Data type
Method
Gemma-2
Gemma-2 9B
Mistral 7B
OmniQuant
Task Avg.
log pplx.
Task Avg.
log pplx.
Task Avg.
log pplx.
b
68.21
2.551
74.38
2.418
73.99
2.110
Baseline
08.25
2.552
74.59
2.418
73.77
2.110
int8
MatQuant
68.02
2.570
74.05
2.438
73.65
2.125
Sliced int8
62.87
2.730
72.26
2,480
38.51
4.681
int4
Baseline
67.03
2.598
74.33
2.451
73.62
2.136
MuaQuant
66.58
2.618
73.83
2.491
73.06
2.153
Sliced int8
39.78
17.030
38,11
15.226
37.29
11.579
int2
Baseline
51.33
3.835
60.24
3.292
59.74
3.931
MatQuant
52.37
3.800
63.35
3.187
62.75
3.153
Sliced int8
67.72
2.497
74.64
2.353
73.00
2.071
int6
Baseline
68.06
2.554
74.23
2.420
74.10
2.112
MatQuant
67.52
2.574
73.92
2.440
73.63
2.127
Sliced int8
41.35
6.024
54.18
3.977
39.21
10.792
int3
Baseline 64.37
2.727
73.23
2.549
71.68
2.211
MatQuant
64.47
2.618
72.87
2.607
71.16
2.238
int4, and then to int8. This observation highlights the po- tential overparameterization and freedom in the int8 data type to accommodate the more stringent needs of int2 dur- ing joint training. We further explore the effects of this phenomenon in Section 5.3 to develop better standalone quantization technique for a single target precision.
tization methods: OmniQuant (Section 4.1) and QAT (Sec- tion 4.2). We demonstrate MaiQuant's efficiency on Transformer-based LLMs. Unless otherwise mentioned, our primary focus is weight only quantization within the parameter-intensivee FFN blocks of the Transformer layer.
For our experiments, we chose the default target quantiza- tion precisions to be int8, int4, and int2. Furthermore, we showcase the interpolative nature of MatQuant through evaluations on int6 and int3, as well as elastic ability to densely span the accuracy-vs-cost trade-off using layer-wise Mix' n 'Macch (Section 4.3). Finally, ablate on improving the performance of Mat Quant (Sections 5.1 and 5.2) and extend MatQuant to the quantization of FFN and Attention parameters. (Section 5.3). Further training and fine-grained evaluation details are in the Appendix.
3.2.1. INTERPOLATIVE BEHAVIOR
Slicing. Although we explicitly train MatQuant for three precisions (int8, int4, int2), we find that the resulting model, when quantized to interpolated bit-widths like int6 & int3 by slicing (Eq. 6) the int8 model, performs on par with baseline trained explicitly for that precision. It is also significantly better than slicing an int8 quantized model. We attribute this strong interpolation in bit-width space to MatQuant, and present more results in Sections 4.1 & 4.2.
Mix'n MatQuant also enables the use of dif- ferent precisions at different layers through layer-wise Mix'n 'n Match (Devvrit et al., 2023), even though we never trained for these combinatorial possibilities. These large number of models, obtained at no cost, densely span trade-off. We explore several Mix'n 'nn Match strategies and find that having higher preci- sion (int8) in the middle layers and a lower precision (int2) at the start and end is the most optimal among hundreds of possible models. See Section 4.3 for detailed experiments.
Models and Data. We experiment with Gemma- 2 (Gemma-Team, 2024) 2B, 9B, and Mistral 7B (Jiang et al., 2023) models. For OmniQuant experiments, we sample 128 examples with sequence length of 2048 from the C4 dataset (Raffel et al., 2020) and train using batch size of 4. We train for a total of 10M tokens for all models ex- cept the int2 baseline, where we train the model for 20M tokens (Shao et al., 2023). For QAT experiments, we sample fixed set of 100M tokens from the C4 dataset and train all our models using batch size of 16 and a sequence length of 8192 for a single epoch.
the accuracy-vs-memony
4. Experiments
Baselines.
For OmniQuant and QAT, our primary base- lines (referred to as "Baseline" in the tables and figures) are models trained explicitly for a given precision. When
In this section, we present an empirical evaluation of Mat Quant working with two popular learning- based quan-
5

---

Matryoshka Quantization
Table MatQuant with QAT across Gemma-2 2B, 9B and Mistral 7B models. MatQuant performs on par with the haseline for int4 and int8 while significantly outperforming it for int2. Even the int3, int6 models ohtained for free through interpolation from
MatQuant perform comparably to the explicitly trained baselines. Task Avg. is average accuracy on the evaluation tasks (1) while log ppix (perplexity) is computed on C4 validation set (J).
Data type
Method
Gemma-2
Gemma-2 9B
Mistral 7B
QAT
Task Avg.
log pplx.
Task Avg.
log pplx.
Task Avg.
log pplx.
16
68.21
2.551
74.38
2.418
73.99
2.110
Baseline
67.82
2.458
74.17
2.29
73.48
2.084
int8
MatQuant
67.44
2.449
74.52
2.262
72.58
2.104
Sliced int8
67.13
2.483
73.36
2.276
71.76
2.18
int4
Baseline
67.03
2.512
73.26
2.324
72.13
2.105
MatQuant
66.59
2.499
73.24
2.429
71.99
2.148
Sliced int8
39.27
10.217
40.40
7.259
37.41
9.573
in12
Buseline
47.74
3.433
56.02
2.923
34.95
2.699
MatQuant
52.20
3.055
62.29
2.265
61.97
2.524
Sliced int8
67.53
2.401
74.15
2.232
73.35
2.097
int6
Buseline
67.75
2.460
74.31
2.293
72.71
2.077
MatQuant
67.33
2.453
74.30
2.265
72.59
2.106
Sliced int8
59.56
2.882
68.70
2.512
64.33
2.493
int3
Buseline
61.75
2.678
69.9
2.43
68.82
2.197
MatQuant 60.76
2.734
70.41
2.429
67.16
2.324
interpolating the models trained with MatQuant for int6 and int3, we do not perform any additional training. How- ever, the baselines are trained explicitly for 6 and 3 bits respectively. We also compare against a sliced int8 Omni- Quant/QAT baseline model to the corresponding precision (referred to as "Sliced int8" in the tables).
Sliced Interpolation. Beyond the target quantization gran- ularities (int8, int4, and int2), MatQuant allows for bit- width interpolation to bit-widths not optimized during train- ing. We find that the accuracy of the int6 and int3 models obtained by slicing the MatQuant models is comparable to their explicitly trained baselines.
4.2. MatQuant with QAT
Evaluation Datasets. Following recent work (Frantar et al., 2022; Ma et al., 2024), we evaluate all the meth- ods based on log perplexity and average zero-shot accuracy across a collection of downstream tasks. We use C4's 's test set to calculate perplexity, and for downstream evaluations, we test on ARC-c, ARC-e (Clark et al., 2018), BoolQ (Clark et al., 2019), HellaSwag (Zellers et al., 2019), PIQA (Bisk et al., 2020), and Winogrande (Sakaguchi et al., 2020).
To further demonstrate the generality of MatQuant, we experiment on the same models using the popular QAT technique. Following the trend of experimental results with OmniQuant, we show in Table 2 that the models trained using MatQuant with QAT are comparable to the explicitly trained basclines for all the targeted bit- widths of int8 and int4. However, int2 quantized models using Mat MatQuant are 4. 4.46%, 6.27%, and 7.02% more accurate for Gemma-2 2B, 9B, and Mistral 7B, respectively.
4.1. MatQuant with OmniQuant
Table shows the efficacy of Mat when used with FFN-only OmniQuant and compared explicitly trained OmniQuant baselines for the target precisions, i.e., int8, int4, and int2, across all the models. While the average downstream accuracy of MatQuant for int8 and int4 quan- tization is within 0.5% of the corresponding independently trained baselines, the int2 quantized models of MatQuant are 1.04%,3 3.11%, and 3.01% more accurate for Gemma-2 2B, 9B, and Mistral 7B, respectively. Similar trends and improvements follow when measuring performance through validation log perplexity. Further, the quantized int4 and int2 models sliced from the int8 OmniQuant baseline suffer a significant drop in accuracy around int4, demonstrating that the nested structure of int8 is not well utilized.
Sliced Interpolation. Models trained using MatQuant with QAT exhibit strong interpolative performance similar to that of MatQuant with OmniQuant. We find that the accuracy of the int6 and int3 models obtained by slicing the MatQuant models is comparable to explicitly trained baselines for both interpolated bit-widths.
While OmniQuant only trains the auxiliary parameters needed for quantization, QAT also updates the weight pa- rameters. This potentially results in severe overfitting to the C4 subset used in the experiments. We observe this overfit- ting in all the experiments presented in Table 2, where the log perplexities improve for QAT compared to OmniQuant, while the downstream accuracies suffer. This also highlights
6

---

Matryoshka Quantization
Gemma-2 9B
Table 3. Design choice ablation for loss re- re-weighting of the target
75
bit-widths (int8, int4, int2) that MatQuant explicitly optimizes. Note that Mat Quanu (0, 0, 1) = Single Precision Mat MatQua
Task Average
Data type
Weightings
Gemma 2 2B
Gemma- 9B Mistral 7B
70
Task Avg.
MatQuant
(0.1,0.1,1]
68.02
71.05
73.27
65
Mixcn'Match
int8
[0.2, 1]
67.91
T3.91
73.44
(0.3,033.
68.01
73.88
73.56
MatQuant-Interp.
(0.4,0.4,1]
67.95
73.84
73.65
Baseline
[0.1,0.1.1]
66.58
T3.83
72.76
60
(0.2. 0.2,
67.47
73.8
73. 13
int4
[0.3, 0.3,
66.97
73:25
73.47
Effective bits per FFN parameter
[0.4,0.4.1)
67.48
74.32
73.60 66.66
[0.1,0.1,1]
52.37
63.35
63.25
Figure 2. Mix'n nn Match on Gemma-2 9B model trained using MatQuant with OmniQuant allows clastic accuracy-vs-cost model extraction for free during deployment.
[0.2,0.2.1]
51.88
64.04
63.99
int2
(0.3, 0.3, 1)
51.05
64.1
63.6
[0.4,0.4. 1)
51.69
61,98
62.75
5.1. Weightings (A,)) for MatQuant
the need for high-quality data for QAT to realize its bene- fits; otherwise, users are better off using resource-friende methods like OmniQuant.
Depending on the constraints, may wish to maximize the accuracy of one of the target bit-widths in MatQuant. Equation 7 provides general formulation of MatQuant that supports searching over the weight X, for bit-width
4.3. Layerwise Mix n Match
r. The results in Section 4 are with the weights that have balanced performance across target precisions. Table shows the weight multiplier ablation results for Gemma- 2 2B, 9B, and Mistral 7B. We find that higher relative value for X2 is essential in attaining good int2 performance. Increasing \2 Às to improve int8 and int4 models often results in accuracy drop for the int2 models. In general, we can see that a higher relative weight for a specific precision results in increased accuracy for that bit-width. We can consider re-weighting as scaling the importance of the bits during training, and finding optimal re- weightiing recipe is interesting research question.
Alongside the strong slicing-based interpolative properties, quantization with MatQuant also enables another form of elastic and interpolative behavior through Mix 'n' Match. Mix 'n' Match provides a mechanism obtain combina- torial number of strong models by using different quanti- zation granularities, from the target bit-widths - i.e., int8, int4, and int2 across layers. Figure 2 shows the ability of Mix'n 'n' Match to densely span the accuracy-vs-bits seerrrcy-s--bit-pe------e per-FFN- parameter (memory/cost) trade-off for the Gemma-2 9B model trained using MatQuant with OmniQuant. While there are many more feasible models, we only showcase the best models obtained through the strategy described in Section 3.2.1 I and further expanded in Appendix B. Inter- estingly, the Mix'n 'n' Match model, with a sub-4-bit effective width, is more accurate than the 4-bit sliced model. This opens up possibilitics for effective serving depending on hardware support. See Section 5.4 for further discussion.
5.2. Co-distillation for MatQuant
Given the nested nature of the models trained using MatQuant, we explored co-distillation, where the outputs from a higher-precision model are used as the target for the nower-precision nested model, either in standalone fashion or alongside the ground truth target (weighted equally). Table shows the effects of co-distillation applied to MatQuant with both OmniQuam and QAT on Gemma-2 9B. While int8 and int4 show no significant improvement, the nested int2 model benefits substantially from the int8 supervision, reaching 0.97% higher accuracy than the non- co-distilled MatQuant with OmniQuant. Co-distillation in Mai Quant opens up avenues for interesting design choices that can further leverage the inherent nested structure of integer data types.
5. Ablations and Discussion
In this section, we present design ablations to improve MatQuant. Section 5.1 discusses the effect of non-uniform weighting across target precisions (int8, int4, int2), and Section 5.2 explores co-distillation of lower precision lev- els (int4, int2) from a higher precision model (int8). Dur- ing the process of extending MatQuant to all Transformer parameters, not just the FFN block, we uncovered an in- teresting hybrid quantization algorithm (between Baseline and MatQuant). Section 5.3 further details this method, called Single Precision MatQuant, which stabilizes the QAT baseline for all the Transformer weights. Finally, we discuss MatQuant beyond integer data types and the con- siderations for effective deployment on current hardware.
5.3. Single Precision MatQuant
In Tables and 2, MatQuant performs on par with the explicitly trained baselines for int4, int8, and the
7

---

Matryoshka Quantization
Table 4. Design choice ablations for co-distillation within MalQuant. - y represents distilling the y-bit model from the x-hit model. We note that the accuracy for int2 has signifi- cantly improved while minimally impacting the other wit-wi.
Table 6. Extending Mat MatQuant with QAT to FFN + At- tention parameters. Baseline QAT destabilizes for int2
and int3 but improves significantly through MatQuant & Single Precision MatQuant.
Gemma-299 9B
OmniQuant
QAT
Data type
Method
Gemma 2 29
Mistral 7B
Data type
Config.
Task Avg.
log pplx.
Task Avg.
log pplx.
QAT
Task Avg.
log pplx.
Task Avg.
log pplx.
2 4, 2 2|
74.05
2.438
T4.52
2.262
b6ooa116
71.38
2,118
73.99
2.110
72.76
2.473
74.75
2.242
in18
Baseline
74.61
2.353
73.73
2.091
88, 2, - 2
73.99
2,355
74.87
2.240
in08
MatQuant
74.85
2.333
73.88
2.182
(8,4,2, 8 4:2]
73.85
2.437
74.81
2:210
[8,4.2] [8,4,8
73.83
2.491
73.24
2.295
Sliced int8
73.15
2.362
71.46
2.290
72.65
2.519
73.76
2.279
int4
Baseline
72.98
2.40
71.87
2.132
int4
18,4,2,8 2]
73.63
2.486
73.77
2.276
MatQuant
74.01
2.396
71.44
2.441
88,422, 4:
73.55
2.478
73.93
2.277
Sliced int8
38.97
23.467
35.06
10.640
Baseline
[8.4.2]
63.35
3.187
62.29
2.660
in12
48 + 2|
62.64
3.289
62.31
2.670
S.P. MatQuuant
45.69
3.780
35.35
7.761
int2
8,,1,2, - 2
62.91
3,138
62.70
2.673
Mat MaiQuant
44.19
3.826
38.35
10.971
88,4,2. 4:2]
61132
3.227
62.60 60
2.670
Sliced int8
74.49
2.290
73.61
2.104
int6
Baseline
71.65
2.357
73.72
2.093
MatQuant
74.57
2.340
74.01
2,161
interpolated int3 and int6 precisions. However, the int2 models show significant accuracy improvement. To investigate this, we conducted a simple ablation in MatQuant by removing the loss terms for int4 and int8 (i.e., R = {2} in Equation 7 or setting b Às 0) and present the results in Table 5. We call this ver- sion of MatQuant as Single Precision Mat Quant. With Single Precision MatQuant, we observe a further boost of up to 1.05%, in the accuracy of int2 models at a ~~2% accu- racy drop in the corresponding int4 and int8 models int2 is still nested within int8. This improvement likely stems from the six additional bits available during MatQuant- style training to optimize the int2 representation.
Sliced int8
64.19
2.895
39.01
6.018
Baseline
int3
S.P. MalQuant
67.68
2.520
67.59
2.335
MatQuant
63.63
2.937
40.55
4.776
icantly impact int2's performance. In fact, in the weight dis- tributions presented in Fig Ic, we observe that Mat results in model where larger number of weights are as- signed to the higher-valued buckets. Seemingly, Mat Quant and Single Precision MatQuant inherently are a better way of performing low-bit quantization.
FFN + Attention Weight Quantization. We present re- sults for FFN + Attention quantization for QAT in Table 6. For int8, int4 and the interpolated intó model, MatQuant performs on par with the Baseline. However, we found int2 and int3 to be very unstable while quantizing both, the FFN and the Attention parameters. Most recent works that do QAT for both the blocks (Chen et al., 2024; Liu et al., 2024a; Du et al., 2024) either do warm starting for the quantized parameters, or have additional distillation and auxiliary loss functions. In the naive setup of minimizing the loss with respect to the ground truth, we find QAT to be very unstable at lower precisions. On the other hand, both MatQuant and Single Precision Mat QaiQuant are very stable further high- lighting the benefits brought by MatQuant style training.
In the case of Single Precision MatQuant, gradient de- scent is free to tune these six additional bits to improve the overall quality of the int2 model. In MatQuant, since we have additional losses preserve the performance of the int4 and int8, the int2 performance is slightly worse than Single Precision MalQuant. However, since the int4 and int8 models are typically very close in accuracy to the bfloat 16 model, MatQuant can shift some of the weights to improve the int2 model. As int4 and int8 models have sub- stantially more quantized buckets than int2, we hypothesize that shifting some weights into adjacent buckets may not significantly affect their performance; however, it can signif-
5.4. Deployment Considerations
Table 5. Single Precision Mat Quant significantly improves upon the baseline for int2 and, times, outperforms MatQuant. Crucially, int8 and int4 performances of Single Precision MatQuant experience significant accuracy decrease (as shown in Tables 21 & 22) in Appendix G).
Current hardware accelerators have native support for serv- ing int8 and int4 quantized models. Additionally, custom- implemented CUDA kemels can can support various low- precision bit-widths, like int2 and int3 (Chee et al., 2024; Frantar et al., 2022). MatQuant can generate large num- ber of models at inference time. Depending on the serving environment, we can choose between Mix' 'n' n'Match models and homogeneous sliced models. For example, suppose the serving environment has a memory constraint equiva- 2.524
int2
Gemma-2 28
Gemma-
Mistral 78
Method
Task Avg.
log pplx.
Tisk Avg.
log pplx.
Task Avg.
log pplx.
QmniQquane
51.33
3.835
60.24
3.292
30.74
3.931
S.P. MaiQuant
53.
3,631
64.02
3.171
63 58
2,976
MusQuantt
52.97
3,800
63.35
3.187
62.75
3.153
QNT
47.74
3.403
56.02
2,923
54.85
2,699
S.P. MatQuant
52.08
3.054
62.66
2,656
61.48
2,509
MasQuant
3.055
62.29
2,600
61.97
lent to an int3 model but lacks optimized support for int3,
8

---

Matryoshka Quantization
Gemma-2 - 9B
while supporting int2. In this case, a Mix'n 'in'"' Match model with small performance drop when compared to the sliced int3 model could be deployed. More generally, as depicted in Figure 2, MatQuunt densely spans the memory-versus- accuracy curve and can be leveraged to obtain performant model for several serving constraints. MatQuant facilitates further research on hardware software co-design to effec- tively support elastic bit-widths on-the-fly during inference.
75
* +0V
Task Average
70
E.P. MatQuant
65
E.P. Mix'n'Match
E.P. MatQuant-Interp.
Baseline
60
5.5. Extension to Floating Point
6 Effective bits per FFN parameter
8
Extending Mal MatQuan Quant to floating-point representations, such as FP8 and FP4, presents significant challenges. Given that the exponent is encoded within the bit representation and contributes to the value as a power of (i.e., effectively log2), slicing it results in buckets whose sizes increase ponentially, unlike the integer case, where bucket sizes are constant. For example, slicing the first two bits from int8 yields buckets of 0, 64, 128, 192. Here, the bucket size (64) is constant; however, this would not be the case when slicing two exponent bits from FP8. This is a promising avenue for future research that could further unlock the benefits of Mat Quant, even during large-scale pretraining.
Figure 3. Mix'n 'n'' MaMach on Gemma-2 9B model trained using Extra Precision MatQuant with OmniQuant as the base algo- rithm allows elastic pareto-optimal accuracy-vs-cost model extrac- tion for free during deployment.
192). However, since the clipping operation is missing in Equation 8, 4 is never clipped down to 3, and r() is now 4 * 64 = 256 Thus, for certain int2 values in our final quantized model, we will have to store an extra bit. This is the case with int3, int4 and int6 as well where an extra bit is required to represent certain values. In Table 24 in Appendix H. we see that the fraction of parameters that fall into this extra bucket is very small. However, for our 2-bit models, this additional bucket gives significant improvements in performance, for example, in Table 24 int2 Gemma-2 9B's average downstream accuracy goes up by 5% when trained with an additional bucket (referred as Extra Precision MatQuant in Table 24 in Appendix H). This number is further boosted 6% with co-distillation, as evidenced by Table 25 in Appendix H. We hypothesize that this additional bucket helps with capturing the outliers and thus leads to a significant performance boost. As highlighted by recent work (Dettmers et al., 2023; Kim et al., 2024), it is crucial to store certain outliers full precision. Interestingly, we show that even a single bit enough to capture several of these outliers, especially for low bit quantization. Finally, note that this performance boost is not very evident in higher precisions where there are enough buckets to account for the outliers.
6. Conclusions
In this work, we presented MatQuant, novel multi-scale training technique that leverages the nested structure of in- teger data types to simultancously optimize model weight quantization across multiple precisions (int8, int4, and int2) within a single model. This general-purposs method, appli- cable bearning-based quantization techniques like Omni- Quant and QAT, produces models with comparable accu- racy to baselines for int8 and int4, while achieving signifi- cant improvements, up to 7% for int2 models. MaiQuant further enables bit-width interpolation and layer-wise mix- and-match for flexible accuracy-cost trade-offs, promising more efficient deployment of large models across various hardware settings. Finally, MatQuant also helped discover Single Precision MatQuant, which significantly improves standalone low-bit quantization.
7. Errata
Mix'n Match As shown in Figure 3 with strong int2 model (i.e., 2.050 bits on average), Extra Precision MatQuant Mix 'n' MMacch densely spans the Pareto-optimal accuracy-vs-bite betsbbitt-ee-FFN- per-FFF-aaaraeeette (memory/cost) trade-off for Gemma-2 9B model trained using MatQuant with Omni-Quant - sometimes even improving on the b model accuracy. Consequently, hardware supporting only int2 and int4 data types can still accommodate model with memory footprint similar to that of an int3 quantized model, and quality comparable or superior to int3; the additional bits required in the case of int2 can be packed into int2/int4. However, custom CUDA kernel would be required to enable sparse additions of these additional bits to the model weights.
In the first draft of the paper, we had a bug and used the following equation to train and quantize our models:
parameter
Sig (1 ge 20-r (2-0)
2
(8)
Equation 8 clearly allows an extra bucket to be included into the quantization range, i.e, r-bit model would have 2" +1 1 possible values instead of 2". For example, consider slicing the first two MSBs from an unsigned int8 value, 234. As per Equation 6, 234 first gets rounded to 4, following which it gets clipped to 3, and finally is scaled up to 3 64 = 192 (Note that MatQuant int2 allows for 0, 64, 128,

---

Matryoshka Quantization
Impact Statement
2404.00456. URL https: /tt.//diii.oo ottp::::////i...//......... a arxiv 2404.00455
This paper introduces novel technique designed to ad- vance the field of machine learning, specifically in the do- main of model compression and efficient deployment for large language models. By enabling the creation of ver- satile, multi-scale models that can operate across various bit-widths, our work has the potential to democratize ac- cess to these powerful technologies by making them more resource-efficient and deployable on a wider range of hard- ware. This could lead to positive impacts such as more sustainable AI systems and greater accessibility for users with limited computational resources. While there are poten- tial risks associated with the broad deployment of powerful Al systems, these are not unique to our work, and we believe the benefits of efficient and accessible AI through innova- tions like MatQuant have significant potential for socictal good. We encourage further investigation into how novel quantization techniques can play a role in future sustainable Al development.
Bengio, Y., Léonard, N., and Courville, A. Estimating or propagating gradients through stochastic neurons for con- ditional computation. arXiv preprint arXiv: 1eo..338.33 2013.
Bisk, Y., Zellers, R., Bras, R. L.., Gao, J., and Choi, Y. PIQA: reasoning about physical commonsense nat- ural language. In The Thirty- Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Confer- ence, IAAI 2020, The Tenth AAAI Symposium on Educa- tional Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020, pp. 7432- 7439. AAAI Press, 2020. doi: 10.1609/AAAL V34100 V34105. 6239. URL https: /tt:://ddo..e. http:///.oooogg....// 1600/ aaaii.
v34105. 6239.
Chee, J., Cai, Y., Kuleshov, V., and De Sa, C. M. Quip: 2-bit quantization of large language models with guarantees. Advances Neural Information Processing Systems, 36, 2024.
Acknowledgments
We are grateful to Varun Yerram, Shreya Pathak and Devvrit for assistance in setting up inference pipelines, Shivani Agrawal. Utku Evci, Praneeth Netrapalli, Rakesh Shiv- anna, Tom Duerig, Abhijit Ogale, Jon Shlens. Ali Farhadi and Rahul Sukthankar for helpful discussions, support and feedback.
Chen, M., Shao, W., Xu, P., Wang, J., Gao, P., Zhang, K., Qiao. Y., and Luo, P. Efficientqat: Efficient quantization-awan training for large language models. CoRR, abs/2407. 10062, 2024. doi: 10.48550/ARXIV. 2407. 1062. URL https: wtw.dioooooo.....oo 11062.
References
Clark, C., Lee, K., Chang, M., Kwiatkowski, T., Collins, M., and Toutanova, K. Boolq: Exploring the surpris- ing difficulty of natural yes/no questions. In Burstein, J., Doran, C., and Solorio, T. (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Lan- guage Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume (Long and Short Papers), pp. 2924-2936. Association for Computational URL
Abdolrashidi, A., Wang, L., Agrawal, S., Malmaud, J., Ry- bakov, O., Leichner, C., and Lew, L. Pareto-optimal quantized resnet is mostly 4-bit. In Proceedings of the JEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 3091-3099, 2021.
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I,, Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint drXiv:2303.08774,220 2023.
Linguistics, 2019. doi: 10.18655/V1//N19-1100 https://doi.. org / 10. 11w118656vvl/ /n1 90t111-13000
Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try are, the AI2 reasoning challenge. CoRR, ftts/180300555 2018. URL http://arxiv.. org abs//b///// aaa.a/a1...55.0 05457.
Adelson, E. H., Anderson, C. H., Bergen, J. R., Burt, J., and Ogden, J. M. Pyramid methods in image processing. RCA engineer, 29(6):33-44 1984.
Adepu. H., Zeng, Z., Zhang, L., and Singh, V. Framequant: Flexible low-bit quantization for transformers. arXiv preprint arXiv:2403.0600 2024.
Denton, E. I.., Chintala, S., Fergus, R., et al. Deep generative image models using a laplacian pyramid of adversarial networks. Advances in neural information processing systems, 28, 2015.
Ashkboos, S., Mohtashami, A.. Croci, M. L., Li, B., Jaggi, M., Alistarh, D., Hoefler, T., and Hensman, Quarot: Outlier-free 4-bit inference in rotated 1lms. CoRR, abs/2404.00456. 2024. doi: 10.48550/ARXIV.
Dettmers, T., Lewis, M., Belkada, Y., and Zettlemoyer, Gpt3. int8 O: 8-bit matrix multiplication for transformers
10

---

Matryoshka Quantization
at scale. Advances in Neural Information Processing Systems, 35:30318-30333 2022.
Stock, P., Scao, T. L.,, Lavril, T., Wang, T., Lacroix, T., and Sayed, W. E. Mistral 7b. CoRR, abs/23 10.0255,, 2023. doi: 10448550/ARXIV.231006662 URL https: / // //doi. .tt.:/1...488050 (/08550/axxiv.. 2310 06825.
Dettmers, T., Svirschevski, R., Egiazarian, V., Kuznedelev, D., Frantar, E.. Ashkboos, S., Borzumov, A., Hoefler, T., and Alistarh, D. Spqr: A sparse-quantized representation for near-lossless Ilm weight compression. arXiv preprint arXiv:2306.000 030788, 2023.
Kim, S., Hooper, C., Gholami, A., Dong, Z., Li, X., Shen, S., Mahoney, M. W., and Keutzer, K. Squeezellm: Dense-and-spars' quantization. In Forty-first Interna- tional Conference on Machine Learning. ICML, 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL https:/ httt:::///ppeneveeeww www.eeeiwwwwwwweeeeeeeee forum? id 0 jpbpFia8m.
Devvrit, F. Kudugunta, S., Kusupati, A., Dettmers, T., Chen, K., Dhillon, I., Tsvetkov, Y., Hajishirzi, H., Kakade, S., Farhadi, A., Jain, P., et al. Matformer: Nested transformer for elastic inference. arXiv preprint arXiv:2310.07707, 2023.
Kusupati, A., Bhatt, G., Rege, A., Wallingford, M., Sinha, A., Ramanujan, V., Howard-Snyder, W., Chen, K., Kakade, S., Jain, P., et al. Matryoshka representation learning. Advances in Neural Information Processing Systems, 55:30233-302433 2022.
Du, D., Zhang, Y., Cao, S., Guo, J., Cao, T., Chu, X., and Xu, N. Bitdistiller: Unleashing the potential of sub-4-bit Ilms via self distillliaa,,.. In Ku, L.. Martins, A., and Srikumar, V. (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguisties (Volume I: Long Papers), ACI. 2024, Bangkok, Thailand, August 11-16, 2024, pp. 102-116. Association for Computational Linguistics, 2024. doi: 10. 10.1865//V1200244CLLLL0. 7. htRtsttps: Utt.:ittpss/aaoooooo//0//..61133 acl sel-long...
Lin, J., Tang. J., Tang, H., Yang, S., Dang, X., and Han, S. Awq: Activation-aware weight quantization for Ilm compression and acceleration. arXiv preprint vrXiv:2306009777 2023.
Lin, T.-Y., Dollár, P., Girshick, R., He, K., Hariharan, B., and Belongie, S. Feature pyramid networks for object detection. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2117-2125, 2017.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang. A., Fan, A., et al. The Ilama 3 herd of models. arXiv preprint arXiv:2407.21782 2024.
Liu, Z., Oguz, B., Zhao, C., Chang, E., Stock, P., Mehdad, Y., Shi, Y., Krishnamoorthi, R., and Chandra, V. LLM- QAT: data-free quantization aware training for large lan- guage models. In Ku, L., Martins, A., and Srikumar, V. (eds.), Findings of the Association for Computational Linguistics, ACL 2024, Bangkok, Thailand and virtual meeting, August 11-16, 2024, pp. 467-484, Association for Computational Linguistics, 2024a. doi: 10. 18653/V1 2024.FINDINGS-ACL2226 URL httpt: //ooi.//oi../ org/ 10 ND.18653/v1///
Frantar, E., Ashkboos, S., Hoefler, T., and Alistarh, D. Gptq: Accurate post-training quantization for generative pre- trained transformers. arXiv preprint arXiv:2210.17322 2022.
G Team, G., Georgiev, P., Lei, V. I., Burnell, R., Bai. L., Gulati, A., Tanzer, G., Vincent, D., Pan, Z., Wang, S., et al. Gemini 1.5: Unlocking multimodal understand- ing across millions of tokens of context. arXiv preprint arXiv:2403.05530, 2024.
windings-acl.....
Liu, Z., Zhao, C., Fedorov, I., Soran, B., Choudhary, D., Krishnamoorthi, R., Chandra, V. Tian, Y., and Blankevoort, T. Spinquant: LLM quantization with learned rotations. CORR, abs/2405. 16406, 2024b. doi: URL https:/ //doi/i/ii.... 10.48550/ARXIV.205..1O 2405. 16406.
Gemma- Team. Gemma 2: Improving open language models at practical size. ArXiv, abs/2408.001 2024. URL http:::/apl..mmmmcccot.e semanticscholaa TD:
org 4 48554485553111
Jacob, B., Kligys, S., Chen, B., Zhu, M.. Tang, M.. Howard, A., Adam, H., and Kalenichenko, D. Quantization and training of neural networks for efficient integer- arithmetic-only inference. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 2704 2704-2713, 2018.
Ma, Y., Li, H., Zheng, X., Ling, F., Xiao, X., Wang. R., Wen, S., Chao, F,, and Ji, R. Affinequant: Affine transformation quantization for large language models. arXiv preprint arXiv:2403.I 125444 2024.
Nair, P. A. and Suggala, A. S. Cdquant: Accurate post- training weight quantization of large pre-trained models using greedy coordinate descent. CoRR, abs/2406.175. 2024. doi: 10.48550/ARXII2224000 17542. URL https: //doi //doi.org/10... 4805//.0...45/////
Jiang, A. Q. Sablayrolles, A., Mensch, A., Bamford, C., Chaplot, D. S., de Las Casas, D., Bressand, F., Lengyel, G., Lample, G., Saulnier, I.,, Lavaud, I.. R., Lachaux, M.,
2406. 17542.

---

Matryoshka Quantization
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21 (140):1-67, 2020.
In Korhonen, A., Traum, D. R., and Mårquez, L. (eds.), Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume 1: Long Papers, PP. 4791-4800. Association for Computational Linguistics, 2019. doi: 10. 1/P19 1472. URL https:// doi. ooi.org/10..... 1ttp:/10...663///
Rippel, O., Gelbart, M.. and Adams. Learning ordered representations with nested dropout. In International Con- ference on Machine Learning, pp. 1746-1754. PMLR, 2014.
Sakaguchi, K., Bras, R. L.., Bhagavatula, C., and Choi, Y. Winogrande: An adversarial winograd schema challenge at scale. In The Thirty-Fourthh AAAI Conference on Ar- rificial Intelligence, AAAI 2020, The Thirty-Seoond In- novative Applications of Artificial Intelligence Confer- ence, IAAI 2020, The Tenth AAAI Symposium on Educa- tional Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020, pp. 8732- 8740. AAAI Press, 2020. doi: 10. 00.16099/AAALV33155..
6399. URL https: http:///ooo../////////o 110910..1.1/a..... v34i05. 6399.
Shao, W., Chen, M., Zhang, Z., Xu, P., Zhao, L., Li, 7.,, Zhang, K., Gao, P., Qiao, Y., and Luo, P. Omniquant: Omnidirectionally calibrated quantization for large lan- guage models. arXiv preprint arXiv:2308. 131377 2023.
Sun, M., Liu, Z., Bair, A., and Kolter, J. Z. A sim- ple and effective pruming approach for large language models. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL https: //openreview net feet/foroe formm?ormmmmmmmm -id-n
Vaswani, A., Shazeer, N. M., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, Attention all you need. In Neural Informa- tion Processing Systems, 2017. URL https://app. htt.o/aar.rrgcoopuusi::::1111111 Io:poooo:p:p1100oa semanticschold wrapioooorrrrroooooooooooo
Xiao, G., Lin, J., Seznec, M., Wu, H., Demouth, J., and Han, S. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference on Machine Learning, pp. 38087-38099 PMLR. 2023.
Yu, H., Li, H., Shi, H., Huang, T. S., and Hua, G. Any- precision deep neural networks. ArXiv, abs/191 2019. URL https://app wwwppp.sssaaaatcccllssaa
Yu, J., Yang, L., Xu, N., Yang, J., and Huang, T. Slimmable neural networks. arXiv preprint arXiv: 0882889,8, 2018.
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi,
Y. Hellaswag: Can a machine really finish your sentence?
12

---

Matryoshka Quantization
A. Particulars of the Slicing Operation.
To extract a r-bit model from a c-bit model, we start by slicing out the most significant r - bits. We use 1 for the n bit if the (r + 1)th else, we use O. This is captured by the round function in Equation 6 and is done push values to higher buckets as we expect them be more informative (Sun al., 2024). For example, consider the the unsigned int8 value 53, The first two MSBs are Os. Naively slicing them would round down 53 to 0, however, bit corresponding to 32 is set, i.e., the (r + 1)th MSB, instead of rounding 53 down to O, we round it up to 1. The clamp(-) operation is also equally important. The rounding operation in Equation 6 will round 240 down to 4, however, unsigned int2 operates with only 0, 1, 2, 3. clamp(-) here would make sure that 4 is clamped down to 3.
want round it up 1. Since the
B. Addition Training Details
We run all our experiments on TPUv5e chips. For OmniQuant experiments, we use constant learning rate of le 3 and for QAT experiments, we linearly warmup the learning rate to 1e for 150 and use consine decay schedule thereafter. For OmniQuant experiments, we sample 128 examples with sequence length of 2048 from the C4 dataset (Raffel et al., 2020) and train using a batch size of 4. We train for total of 10M tokens for all models except the int2 baseline, where we train the model for 20M tokens (Shao et al., 2023). For Co-distillation experiments where OmniQuant is the base algorithm, we train for a total of 8.3M tokens. For QAT experiments, we sample a fixed set of 100M tokens from the C4 dataset and train all our models using a batch size of 16 and a sequence length of 8192 for single epoch. For Attn + FFN experiments with QAT, we sample fixed set of 300M tokens from C4 and train with batch size of 16 for a single epoch. We use (Ag, A2) 1.0) for all our Gemma experiments unless otherwise stated. In the case of Mistral 7B, for OmniQuant experiments, we use (As, X4, A2) (0.4, 0.4, 1.0), and for QAT experiments we use (As, A2) = (0.2, 0.2, 1.0). For all our Extra Precision MatQuant experiments, we use (As, As, A2) = (1.0, 1.0, 1.0).
Mix'n'Maren For fixed effective bits-per-FEN layer, where each layer was quantized to either int2, int4, int8, we explored four different quantization strategies: Pyramid, Reverse Pyramid, Increasing, and Decreasing, In the Pyramid strategy, the initial and final layers were quantized to int2, the central layers to int8, with int4 serving as an intermediate step. The Reverse Pyramid strategy followed the opposite approach, assigning int8 to the initial and final layers, int2 to the central layers, and int4 between. The Increasing and Decreasing strategies assigned bit precision ascending and descending order, respectively, across the layers. Our experimental results demonstrated that, for a given effective bits per FFN layer, the Pyramid strategy consistently outperformed the others. Allocating higher precision (int8) the middle layers helped preserve critical information, while the initial and final layers performed adequately with lower bit precision (int2 and int4), leading to more efficient and effective quantization scheme.
C. Detailed Downstream Evaluations for OmniQuant and QAT
Tables 8,8,9, 9, 10, 11, and 12 present downstream evaluation results on Gemma-2 2B, Gemma-2 9B and Mistral 7B with OmniQuant and QAT.
D. Detailed Downstream Evaluations for MatQuant Re-weighting
Tables 13, 15, and 14 present downstream evaluation results for OmniQuant reweighting experiments on Gemma-2 2B, Gemma-2 9B and Mistral 7B.
E. Detailed Downstream Evaluations for Co-Distillation
Tables 16 and 17 present the downstream evaluation and perplexity results for MatQuant with co-distillation Gemma-2 9B. We present results with both, OmniQuant and QAT as the base algorithms.
F. Detailed Evaluations for FFN + Attention Quantization
Tables 18 and 19 present the downstream evaluation and perplexity results for FFN + Attention quantization on Gemma-2 9B and Mistral 7B with OmniQuant and QAT.
13

---

Matryoshka Quantization
G. Detailed Evaluation for Single Precision MatQuant
Tables 20, 21, 22, and 23 present the downstream evaluation results comparing Single Precision MatQuant to Mat Quanu and the Baseline for int2 quantization of Gemma-2 2B, Gemma-2 9B and Mistral 7B with OmniQuant and QAT. Since Single Precision MatQuant slices 2 bits from an 8-bit model and computes loss only over the first bits, can evaluate the Single Precision MatQuant model trained for 2-bits int4 and int8. Downstream evaluation and perplexity results for this are presented in Tables 21 and 22.
We also plot the weight distribution for Single Precision Mal in Figure 4.
1el
Gemma-2 9B Weight Distribution for OmniQuant
Baseline
int8
1.0
S.P. MatQuant
0.8
0.2
0.0
256
Quantized buckets
Baseline
int4
Basdire
int2
S.P MaiQuant
S.lt Maqquant
00
Frequency
2.0
1.0
00
0.0 0
Quantized buckets
Quantized buckets
Figure 4. The Figure presents the weight distribution for Gemma-2 9B when trained with Single Precision MatQuant for int2 quantiza- tion. The right-shifted quantized weight distribution is a consequence of Single Precision MatQuant's training mechanism that heavily optimizes for the first MSBs the int8 representation.
H. Detailed Evaluation for Extra Precision Mat Muanu
Table 24 presents downstream and perplexity evaluation comparing MatQuant and Extra Precision MatQuant for Gemma-2 2B, 9B and Mistral 7B with OmniQuant as base algorithm. Table 25 presents codistillation ablation results for Extra Precision MatQuant on Gemma-2 9B with OmniQuant as base algorithm. Tables 26, 27, and 28 present downstream evaluation results for Extra Precision MatQuant when applied Gemma-2 2B, 9B, and Mistral 7B with OmniQuant as the base algorithm. Table 29 presents downstream evaluation and perplexity results for Extra Precision MatQuan co-distillation experiments on Gemma-2 9B with OmniQuant the base algorithm. Table 30 presents downstream task average and perplexity results comparing Single Precision MatQuant and Single Precision Extra Precision Mat Quant when applied to Gemma-2 2B, 9B and Mistral 7B with OmniQuant as the base algorithm.
14

---

Matryoshka Quantization
Table Table presents downstream evaluation results
MatQuant when applied OmniQuant on Gemma-2
type
Method
Gemma-2
OmniQuant
ARC-c
ARC-e
BoolQ
HellaSwag
PIQA
Winogrande
Average
b6
71.59
76.45
69.69
78.29
int8
Baseline
71.46
76.36
69.76
78.24
63.89
MatQuant
49.66
71.00
76.73
68.85
78.56
63.30
68.02
Sliced int8
11.35
66.12
72.02
62.31
75.79
59.43
62.87
int4
Baseline
18.46
70.96
74.22
67.66
77.26
63.61
67.03
QuaQuant
47.27
70.79
73.76
66.85
78.07
62.75
66.58
Sliced int8
23.55
27.66
59.63
24.09
51.58
52.17
39.78
int2
Baseline
31.31
53.58
62.2
40.78
66.05
54.06
51.33
QatQQuant
29.95
54.21
64.40
44.37
66.81
54.46
52.37
Sliced int8
18.72
71.13
76.06
69:12
75.45
62.83
int6
Baseline
49.32
71.76
76,48
69.52
78.56
62.75
68.06
Quaqu
48.89
70.50
75.69
68.89
78.40
62.75
67.52
Sliced int8
22.35
34.97
56.94
29.49
65.44
48.93
41.35
int3
Baseline
46.25
68.64 23.64
72.97
62.24
76.06
60.06
64.37
MatQuant
11.03
17.09
74.25
62.78
77.26
61.10
64.17
Table 8. Table presents the downstream evaluation results for MatQuant when applied to OmniQuant on Gemma 9B.
Data type
Method
Gemma-2 9B
OmniQuant
ARC-c
ARC-c
BoolQ
HellaSwag
PIQA
Winogrande
Average
16
58.96
77.57
83.33
77.31
81.12
67.96
74.38
Baseline
59.47
77.31
83.94
77.35
81.39
68.11
74.59
int8
MatQuantt
57.59
77.02
84.01
76.61
81.18
67.88
74.05
Sliced int8
35.80
82.32
73.56
int4
Baseline
38.79
78.37
83.55
76.71
81.45
67.09
74.33
MatQuuntt
58.02
78.11
83.24
76.08
80.96
66.54
73.83
Sliced int8
24.57
52.97
50.16
38.11
int2
Baseline
39.16
63.43
72.11
52.24
72.63
61.88
00.24
MatQuant
40.78
67.85
73.64
60.56
72.09
65.19
63.35
Sliced int8
59.04
77.61
84.62
77.10
81.18
68.27
74.64
int6
Baseline
59.22
77.27
83.21
77.1
81.12
67.48
74.23
MatQuant
37.25
76.94
81.04
76.63
81.31
67.32
73.92
Sliced int8
34.30
55.47
66.36
46.91
67.19
54.85
54.18
int3
Baseline
57.17
77.06
83.79
74.45
80.36
66.54
73.23
MatQuant
55.80
76.89
81.99
74.27
80.14
68.11
72.87
Table 9. Table presents the downstream evaluation results for MatQuant when applicd to OmniQuant Mistral 7B.
Data type
Method
Mistral 7B
OmniQuant
ARC-c
ARC-e
BoolQ
HellaSwag
PIQA
Winogrande
Average
b6loa116
84.4
Baseline
49.23
73.19
83.88
80.41
81.39
74.51
73.77
int8
MatQuant
49.06
72.52
81.74
79.21
81.45
74.90
73.65
Sliced int8
21.33
33.07
42.08
28.62
55.06
49.72
38.51
int4
Baseline
49.23
73.23
83.94
81.34
74.11
73.62
MatQuant
47.87
71.55
83.88
78.85
81.31
74.90
73.06
Sliced int8
24.32
23.44
49.72
24.71
51.74
49.80
37.29
int2
Baseline
36.
61.36
70.06
57.47
62.19
59.74
MatQuant
73.15
62.75
Sliced int8
48.21
71.09
83.21
79.93
81.28
74.27
73.00
in16
Baseline
50.26
73.65
81.04
80.55
81.06
74.43
74.1
MatQuant
49.40
72.47
84.68
79.52
81.34
74.35
73.63
Sliced int8
25.26
25.76
61.99
24.67
48.31
49.25
39.21
int3
Baseline
46.33
70.71
82.72
77.74
80.74
71.82
71.68
MatQuant
47.35
71.00
80.00
76,96
80.30
71.35
71.16
15

---

Matryoshka Quantization
Table 10. Table presents the downstream evaluation results for MatQuant when applied to QAT on Gemma-2 2B.
Data type
Method
Gemma-2 2B
QAT
ARC-c
ARC-e
BoolQ
HellaSwag
PIQA
Winogrande
Average
b9loat16
50.09
71.59
76.45
69.69
78.29
63.14
68.21
Baseline
47.78
70.66
75.08
69.92
78.35
65.11
67.82
int8
QuaQu
71.21
75.99
67.44
Sliced int8
46.16
69.53
75.35
68.49
78.18
65.04
67.13
in14
Baseline
46.16
71.59
73.73
68.72
78.62
63.38
67.03
MatQuart
69.58
75.84
77.80
66.59
Sliced int8
24.06
26.94
59.05
25.57
51.85
48.15
39.27
in12
Baseline
24.66
43.22
38.39
64.42
53.59
17.74
MatQuant
28.33
51.85
63.64
46.94
68.28
51.14
52.20
Sliced int8
47.87
70.83
74.25
69.80
77.86
G4.56
67.53
int6
Buseline
47.7
70.88
74.92
69.72
78.07
65.19
67.75
Mas Quant.
45.39
71.17
76.15
68.33
78.13
64.80
67.33
Sliced int8
37.97
62.67
64.71
58.01
74.27
59.75
59.56
int3
Buseline
39.68
65.28
67.00 07.03
62.68
77.04
58.8
61.75
Mas Quant.
36.95
66.20
64.25
61.03
75.19
60.93
60.76
Table 11. Table presents the downstream evaluation results for Mat Quant when applied to QAT on Gemma-2 9B.
Data type
Method
Gemma-2 9B
QAT
ARC-e
ARC-e
BoolQ
HellaSwag
PIQA
Winogrande
Average
16
58.96
77.57
83.33
77.31
81.12
67.96
74.38
Baseline
58.11
73.38
80.12
78.7
81.5
71.19
74.17
int8
MasQuant:
57.68
76.09
82.23
78.41
82.26
70.48
74.52
Sliced int8
56.91
75.17
78.78
77.02
81.18
71.11
73.36
int4
Buseline
56.91
75.42
75.38
78.06
81.39
72.38
73.26
Mas MuaQuant
56.66
75.72
77.55
77.30
81.23
70.96
73.24
Sliced int8
23.46
28.28
57.09
29.76
53.48
50.36
40.40
int2
Buseline
33.45
55.43
62,26
54.8
70.51
59.67
56.02
Mas Quant
41.21
66.84
65.41
63.61
75.41
61.25
62.29
Sliced int8
57.68
75.17
80.73
78.66
81.77
70.88
74.15
int6
Baseline
57.94
76.14
79.63
78.93
82.1
71.11
74..1
MarQuant 57:25
76.01
81.83
78.25
81.77
70.72
74.30
Sliced int8
50.60
67.85
75.54
71.07
79.11
68.03
68.70
int3
Baseline
53.07
75.04
66.61
74.94
80.03
09.69
69.9
Mat Quant.
51.19
71.80
78.69
73.18
79.49
68.11
70.41
Table 12, Table presents the downstream evaluation results for MatQuant when applied to QAT Mistral 7B
Data type
Method
Mistral 7B
QAT
ARC-
ARC-e
BoolQ
HellaSwag
PIQA
Winogrande
Average
b6
49.57
73.74
84.4
80.61
81.18
74.43
73.99
Baseline
48.59
71.63
82.45
81.69
81.18
75.06
73.48
int8
MarQuant
47.44
71.21
82.08
80.31
80.74
73.72
72.58
Sliced int8
47.61
70.41
80.21
79.74
79.98
72.61
71.76
int4
Baseline
47.27
70.02
81.28
78.95
81.12
73.56
72.13
Mac Quant
45.99
72.22
81.90
79.08
80.36
72.38
71.99
Sliced int8
24.40
25.97
47.52
24.66
50.27
51.62
$7.41
int2
Baseline
29.78
48. 23
64.55
55.11
70.84
61.25
541.95
Mac Quant
35.58
56.36
72.66
66.68
74.32
66.22
61.97
Sliced int8
48.55
71.76
82.57
81.67
$1.39
74.19
73.35
intó
Baseline
47.7
71.3
82.23
79.84
80.79
74.43
72.71
MusQuant
71.34
81.96
80.27
80.52
74.51
72.59
Sliced int8
38.99
61.11
72.54
65.65
77.48
70.24
64.33
in13
Baseline
44.5-1
67.97
73.98
76.31
79.65
70.48
68.82
MarQuant
40.10
62.42
79.05
73.82
77.31
70.24
67.16
16

---

Matryoshka Quantization
Table 13. Tables presents the downstream evaluation results on Gemma-2 2B for MatQuant loss reweighting when applied to Omni OmniQuant. Weightings: (xt y, 2) + (As, Aa, A2) (from Equation 7).
Gemma-2 2B
Data type
Weightings
ARC-c
ARC-e
BoolQ
HellaSwag
PIQA
Winogrande
Average
(0.1,0.1,1)
49.66
71
76.73
68.85
78.56
63.3
68.02
(0.2,0.2,
49.4
71.3
76.21
08.97
78.29
63.3
07.91
int8
0.33,033.1 1)
48.81
71.72
76.37
68.95
78.4
63.
68.
(0.4,0.4,1)
48.72
71.72
76.61
68.92
78.73
62.98
67.95
(0.5,0.5.1) 1)
49.06
71.34
76.15
68.86
78.45
62.98
67.81
(0.1,0.1.1)
47.27
70.79
73.76
66.85
78.07
62.75
66.58
(0.2,0.2,1)
48.63
76.06
68.11
77.97
63.06
67.47
int4
(0.3,0.3.1)
47.7
71.17
75.08
67.57
77.09
62.59
06.97
(0.4.04.1)
48.29
71.25
76.76
77.58
67.48
(0.5,0.5.1)
75.9
67.57
78.4
64.01
67.43
(0.1.0.1,1)
29.95
54.21
64.4
44.37
66.81
54.46
52.37
(0.2,0.2.1)
30.03
52.78
62.39
54.62
int2
(0.3.0.3,1) 1)
62.57
51.05
(0.4.0.4,
75
66.16
55.64
(0.5,0.5,
27.13
51.05
60.95
39.94
65.56
54.3
49.82
(0.1,0.1,1) 1)
48.89
70.5
75.69
08.89
78.4
62.75
67.52
(0.2,0.2.1)
49.32
70.96
75.87
68,93
78.29
62.67
67.67
int6
(0.3,0.3,1 1)
48.98
71.63
76.21
68.68
78.73
63.46
67.95
(0.4,0.4,
48.95
71.72
75.75
68.83
78.67
63.61
67.93
(0.5,0.5.1)
49.4
71.59
76.21
68.63
78.29
63.85
67.99
(0.1,0.1,1)
44.03
67.09
74.25
62.78
77.26
61.4
64.47
(0.2,0.2,1)
43.09
65.7
67.19
59.57
75.3
60.38
61.87
int3
(0.3. 0.3.0.3.
43.94
68.35
71.87
59.54
75.79
30.98
63.24
(0.4.0.4,1)
41.81
65.53
72.91
61.42
75.03
61.88
63.1
(0.5,0.5.
41.61
67.34
71.87
61.15
74.51
63.03
Table 14. Tables presents the downstream evaluation results on Gemma-2 2 9B for MatQuant loss reweighting when applied to Omni OmniQuant. Weightings: (xr, V. =) (As; X4, 22) (from Equation 7).
Gemma-2
Data type
Weightings
ARC-c
ARC-e
BoolQ
HellaSwag
PIQA
Winogrande
Average
(0.1,0.1, 1)
57.59
77.02
84.01
76.61
81.18
67.88
74.05
(0.2,0.2,
37.76
76.73
83.73
76.5
81.31
67.4
73.91
int8
(0.3, 0.3. 1)
37.94
76.64
83.36
76.56
81.01
67.8
73.88
(0.4.0.4,1) 1)
58.25
76.52
83.15
76.74
80.96
73.84
(0.5, 0.5, 1)
37.68
76.68
83.39
76.62
81.07
67.09
73.75
(0.1,0.1.1)
58.02
78.11
83.24
76.08
80.96
66.54
73.83
(0.2, 0.2.0.2.1
58.96
77.9
82.57
76.14
$1.07
66.14
73.8
int4
(0.3,0.3.1)
37-42
77.23
81.62
75.72
80.85
66.69
73.25
(0.4.0.4.1)
78.32
84.53
76.17
81.45
66.46
74.32
(0.5.0.5.1) 1)
37.08
84.65
81.56
(0.1.0.1.1)
40.78
67.85
73.64
72.09
65.19
63.35
(0.2,0.2,1)
40.53
67.97
75.57
60.83
72.25
in12
(0.3,0.3.1)
79.08
60.79
72.47
65.19
(0.4.0.4,
66.54
66.24
73.07
65.27
G1.98
(0.5.0.5,
40.02
66.16
69.08
60.54
73.23
64.58
62.32
(0.1,0.1,
37.25
76.91
81.04
76.63
81.34
67.32
73.92
(0.2,0.2,1)
37.25
76.6
83.79
76.46
81.12
67.64
73.81
int6
(0.3,0.3, 1)
58.7
76.98
83.09
76.63
80.69
67.32
73.9
(0.4,0.4.1
58. 18.28
76.43
83.15
76.76
81.18
67.09
73.81
(0.5.0.5.1)
58.28
76.3
83.33
76.68
81.18
66.93
73.78
(0.1,
55.8
76.89
81.99
74.27
80.14
68.11
72.87
(0.2,0.2.1)
34.69
76.56
79.79
73.92
79.92
66.77
71.94
int3
(0.3,0.3.1
56.48
77.58
83.09
73.71
80.69
67.32
73.14
(0.4.0.4.1) 1)
56.23
77.86
83.79
74.12
09.09
68.98
73.61
(0.5,0.5.l
34.35
76.3
83.67
74.21
80.09
72.77
17

---

Matryoshka Quantization
Table Tables presents the downstream evaluation results on Mistral 7B Mat Quant loss reweighting when applied to OmniQuant. Weightings: (ær, v. + (As, X4, A2) (from Equation 7).
Mistral 78
type
Weightings
ARC-c
ARC-e
BoolQ
HellaSwag
PIQA
Winogrande
Average
(0.1,0.1.1)
49.23
83.94
81.39
(0.2,0.2,1)
71.97
83.91
81.5
74.98
73.44
int8
(0.3,0.3,1)
49.32
72.39
81.43
79.24
81.23
74.74
73.56
00.4,0.411)
49.06
72.52
84.74
79.21
81.45
74.9
73.65
(0.5,0.5.1)
49.15
72.64
84.65
79.37
81.72
74.82
73.72
(0.1,0.1.1)
47.61
71.59
83.3
78.32
81.61
74.11
72.76
(0.2,0.2.1)
48.12
72:14
84.07
78.72
81.45
74.43
73.16
int4
(0.3,0.3, 1)
48.21
72.81
841.
79.02
81.18
75.22
73.47
(0.4,0.4, 1)
47.87
71.55
83.88
78.85
81.34
74.9
73.06
(0.5.0.5.1)
48.21
71.97
83.82
79.03
81.39
74.35
73.13
(0.1,0.1,
37.46
63.43
71.53
06,22
75.24
65.59
03.25
(0.2,0.2.1)
37.54
64.81
71.8
66.57
74.37
65.27
63.39
in12
(0.3,0.3.1)
37.46
62.92
75.35
67.2
74.43
64.25
63.6
(0.4,0.4,1) 1)
37.88
62.58
73.15
65.89
73.88
63.14
62.75
(0.5,0.5, 1)
37.29
62.75
09.36
61.99
72.36
64.25
61.83
(0.1,0.1,
49.57
71.72
83.76
78.87
$1.28
74.03
73.2
(0.2,0.2.1)
49.49
72.52
84.22
79.08
81.39
74.19
73.48
int6
(0.3.0.3.1)
83.85
(0.4,0.4,1)
49.4
81.68
79.52
81.34
74.35
73.63
(0.5.0.5.1)
49.4
72.39
84.31
79.5
81.28
74.27
73.52
(0.1,0.1.1)
70.54
(0.2,0.2,1)
43.94
67.85
81.56
76.55
72.61
70.38
int3
0..3.0.3.1)
45.39
67.89
80.92
77.13
80.47
72.06
70.64
00.4.0.4,1
47.35
80
76.96
80.3
71.35
71.16
(0.5.0.5, 1)
46.76
70.29
82.17
77.32
80.9
71.11
71.13
Table 16. Table presents the downstream evaluation and perplexity results for our MatQuant co-distillation experiments on Gemma-2 9B with OmniQuant.
OmniQuant
Gemma-2 9B
Data lype
Config.
ARC-c
ARC-e
BoolQ
HellaSwag
PIQA
Winogrande
Average
log pplx.
88,4,8 + 18,4,2.8 88,4228 4:2
57.51
76.26
83.30
73.35
80.74
65.43
72.76
2.473
int8
58.19
76.89
83.73
76.75
81.39
67.01
73.99
2,435
67.68
77.06
83.00
76.76
81.45
67.17
73.85
2.437
88,4,8 - [8,4,288
56.23
76.47
73.03
66.85
72.65
2.519
int4
57.51
76.73
76.23
67.09
8,,4228 + 4;2]
57.51
76.68
83.27
75.85
81.61
06.38
73.55
2.478
[8,4,8 -
38.14
66.50
76.73
59.70
71.11
63.69
62.64
3.289
int2
88,4228 - 2]
40.61
67.55
71.07
60.80
72.96
64.48
62.91
3.138
8,4,2,8 4:2]
42.75
69.65
74.40
60.53
72.42
66. 14
64..22
3.227
[8,4,8 2
57.59
76.30
83.55
73.41
80.85
65.51
72.87
2.469
int6
[8,4,,8 2]
58.28
76.85
83.43
76.91
81.18
67.01
73.94
2,438
8,4,2,8 4:2]
58.11
76.98
83.33
76.70
81.45
67.48
74.01
2.439
88448 +
52.30
75.25
78.26
71.08
79.49
65.35
70.29
2.651
int3
8,,,2,8
54.44
75.97
82,20
73.84
80.20
66.46
72.19
2 2.603
8, 4:22
54.44
76.26
81.90
73.89
79.92
65.75
72.03
2.604
18

---

Matryoshka Quantization
Table Table presents downstream evaluation and perplexity MatQuuat co-distillation experiments on Gemma-2 with QAT.
QAT
Gemma-2 9B
Data type
Config.
ARC-c
ARC-e
BoolQ
HellaSwag
PIQA
Winogrande
Averuge
log pplx.
88,4,8
57.68
76.09
82.60 20.
78.75
82.48
70.88
74.75
2.242
int8
[8.4,2,8 88,4,28 4:
57.76
76.35
81.50
79.13
82.43
72.06
74.87
58.19
76.05
78.92
71.90
[8.4,8 2
57.85
76.81
78.47
77.62
80.96
70.88
73.76
2.279
int4
18,4,2,8 -
57.08
75.88
78.47
77.65
81.34
72.22
73.77
2.276
88,4,28 4:2]
67.34
75.80
78.99
77.67
81.50
72.30
73.93
2.277
88,4,8 2
40.61
67.17
67.37
63.10
75.24
60.38
62.31
int2
[8,4,2.8 2]
40.53
66.71
63,29
62.35
62.70
[8,4,2, + 4:
40.10
66.37
67.86
63.14
75.08
63.06
62.60
2.670
[8,4,8 +
57.85
76.05
82.23
78.70
82.10
71.43
74.73
2.245
int6
[8,4,2,
58.11
75.93
82.14
79.10
82.26
71.19
74.79
2.243
88,,2,8 - 4:2]
58.19
75.67
81.31
78.80
82,15
71.27
74.56
2.243
[8,4,8 2
51.19
71.00
76.67
73.07
79.54
68.03
69.92
2,441
int3
8,4,2,8 2]
51.71
71.46
76.85
73.00
79.00
67.88
69.98
2,437
2,
4;
51.28
71.34
76.12
72.96
08.98
70.00
2.435
Table 18. Table presents the downstream evaluation results for MatQuan PFN + Attention quantization on Gemma-2 9B with QAT.
Data type
Method
Gemma- 9B
ARC-e
ARC-e
BoolQ
HellaSwag
PIQA
Winogrande
Average
16loat16
58.96
77.57
83.33
77.31
81.12
67.96
74.38
Baseline
58.62
77.02
83.43
79.01
81.34
68.27
74.61
int8
MatQuant
77.99
84.13
77.85
68.43
int8
57.42
76.01
80.86
76.34
80.03
68.27
73.15
int4
Baseline
56.06
74.96
79.27
77.83
80.25
69.53
72.9%
MatQuant
58.79
75.80
81.89
76.26
81.23
67.09
71.01
Sliced int8
26.37
25.34
58.10
25.60
19.08
49.33
38.97
int2
Bascline
S.P. MatQuunt
25.26
38.47
62.14
35.09
61.70
51.46
45.69
MatQuant
23.72
62.17
33.72
59.36
49.57
44.19
Sliced int8
58.53
77.10
83.00
78.81
81.07
68.43
71.19
int6
Baseline
58.87
77.06
83.12
78.81
81.23
68.82
74..65
MatQuant
58.00 .6
78.03
83.30
77.72
80.96
68.43
74.57
Sliced int8
44.71
65.28
71.56
65.25
75.84
62.51
64.19
int3
Baseline
S.P. MatQuant
18.55
71.25
68.38
72.12
79.00
66.77
67.68
MatQuant
43.34
61.91
75.96
65.20
75.46
59.91
63.63
19

---

Matryoshka Quantization
Table 19. Table presents the downstream evaluation results for MatQuant FFN Attention quantization Mistral 7B with QAT.
Data
Method
Mistral 7B
ARC-c
ARC-e
BoolQ
HellaSwag
PIQA
Winogrande
Averuge
bíloat16
19.57
73.74
81.4
80.61
81.18
71.43
73.99
Baseline
49.23
72.9
83.49
80.26
81.28
75.22
73.73
in18
MuaQuant
50.09
73.11
83.73
80.73
81.39
73.88
73.88
Sliced int8
45.99
71.55
81.19
76.90
80.58
72.53
71.46
in14
Baseline
48.04
71.72
78.87
78.93
80.36
73.32
71.87
MatQuant
70.29
77.34
80.25
72.53
Sliced int8
24.40
49.13
50.99
in12
Baseline
S.P. MatQuant
22.53
25.51
38,90
24.13
50.92
50.12
35.35
MatQuant
21.33
25.59
57.37
24.85
50.92
50.12
38.36
Sliced int8
73.53
80.28
80.96
74.98
int6
Baseline
49.32
73.4
82.48
80.24
81.28
75.61
73.72
MatQuant
50.00
73.78
83.55
80.74
81.66
74.51
74.04
Sliced in18
19.97
30.72
46.79
27.22
58.43
50.91
39.01
in13
Baseline
S.P. MatQQuatt
43.86
67.51
70.43
73.97
80.36
09.38
67.50
Mat Quant
20.82
33.42
53.30
27.77
56.76
49.25
40.55
Table 20. Table presents downstream evaluation and perplexity results for Single Precision MatQuant, comparing with Mat MatQuant and the Baseline for int2 quatization of Gemma-2 2B with OmniQuant and QAT.
int2
Gemma2-2B
Method
ARC-c
ARC-c
BoolQ
HellaSwag
PIQA
Winogrande
Task Avg.
log pplx.
MatQuant
29.78
57.70
63.39
44.32
68.66
56.67
53.42
3.631
OmniQuant
Baseline
31.31
53.58
62.2
40.78
66.05
54.06
51.33
3.835
MatQuant
29.95
54.21
64.40
44.37
54.46
52.37
S.P. MatQuaut
28.07
52.36
62.87
46.80
68.88
53.51
52.08
3.054
QAT
Baseline
21.66
43.22
62.17
38.39
64.42
53.59
47.74
3.433
MatQuaut
28.33
51.85
63.64
46.94
68.28
54.14
52.20
3.055
Table 21. Table presents downstream evaluation and perplexity results for Single Precision MatQuant, comparing with MatQuant
and the Baseline for int2, int4, int8 quatization of Gemma-2 9B with Baseline. Note that the model was trained with Single Precision MatQuant for int2; the
model were sliced post training.
Gemma-2 9B
Data type
Method
ARC-c
ARC-e
BoolQ
HellaSwag
PIQA
Winogrande
Averuge
log pplx.
S.P. MatQuaut
57.94
76.64
82.66
76.98
81.01
67.56
73.80
2.372
int8
Baseline
59.47
77.31
83.94
77.35
81.39
68.11
74.50
2.418
MatQuant
57.59
77.02
84.01
76.61
81.18
67.88
74.05
2.438
S.P. MatQuant
37.17
76.39
81.47
75.81
80.85
66.38
73.01
2,420
int4
Baseline
58.79
78.37
83.55
76.71
81.45
67.09
74.33
2.451
MatQuant
58.02
78.11
83.24
76.08
80.96
66.54
73.83
2.491
S.P. MatQuaut
40.44
66,75
77.92
60.42
72.52
66,06
64.02
3.171
int2
Baseline
39.16
63.43
72.11
52.24
72.63
61.88
60.24
3.292
MatQuant
40.78
67.85
73.64
60.56
72.09
65.19
63.35
3.187
20

---

Matryoshka Quantization
Table 22. Table presents downstream evaluation and perplexity results for Single Precision Mat comparing it with MaiQuant Note that the model trained with
and the Baseline for int2, int4, int8 quatization of Gemma-2 9B with Bascline.
Single Precision MatQuant for int2; the int4 and int8 model were sliced post training.
Gemma-2 9B
Data type
Method
ARC-e
ARC-e
BoolQ
HellaSwag
PIQA
Winogrande
Average
log pplx.
S.P. MatQuant
55.89
75.84
79.37
75.47
81.07
68.43
72.71
2.363
int8
Baseline
58.11
75.38
80.12
78.7
81.5
71.19
74.17
2.29
MarQuant
37.68
76.09
82.23
78.41
82.26
70.48
74.52
2.262
S.P. MatQuant 54.95
75.59
75.05
74.60
80.79
69.06
71.67
2.394
int4
Baseline
56.91
75.42
75.38
78.06
81.39
72.38
73.26
2.324
MarQuant
56.66
75.72
77.35
77.30
81.23
70.96
73.24
2.295
S.P. MatQuant
40.53
67.38
66.91
63.62
61.88
62.66
2.656
in12
Baseline
33.45
55.43
62.26
54.8
59.67
56.02
2.923
MatQuant
41.21
66.84
65.41
63.61
75.41
61.25
62.29
2.660
Table 23. Table presents downstream evaluation and perplexity results for Single Precision MatQuant, comparing it with MatQuant, and the Baseline for int2 quatization of Mistral 7B. Results are presented for both, OmniQuant and QAT as the base algorithms.
int2
Mistral
Method
ARC-c
ARC-e
BoolQ
HellaSwag
PIQA
Winogrande
Task Avg.
log pplx.
S.P. MatQuant
37.63
64.14
72.45
67.47
74.81
64.96
63.58
2.976
OmniQuant
Baseline
36.69
61.36
70.06
57.47
70.67
62.19
59.74
3,931
MatQuant
37.88
62.58
73.15
65.89
73.88
63.14
62.75
3,153
S.P. MatQuant
35.24
57.15
69.88
66.02
75.41
65.19
61.48
2.509
QAT
Baseline
29.78
48.23
64.5
55.11
61.25
54.95
2.694
MatQuant
35.58
56.36
72.66
66.68
74.32
66.22
61.97
2.524
Table 24. Results comparing MatQuant with Extra Precision MatQuant for Gemma-2 2B, 9B, and Mistral 7B, with OmniQuant the base algorithm. We find that for the 2-bit model, having an extra bucket significantly boosts the performance, however, this is not the case with the higher precisions.
Method
Gemma-2 2B
Gemma- 2 9B
Mistral 7B
OrniQuant
Avg. Bits
Task Avg.
log pplx.
Avg. Bits
Task Avg.
log pplx.
Avg. Bits
Task Avg.
log ppla.
bffoat16
68.21
2.551
74.38
2.118
73.99
2.110
MatQuaut
68.
2.570
74.05
2.438
73.65
2.125
Extra Precision MatQuant
8
67.85
2,580
74.33
2.446
73.46
2.132
MatQuant
4
66.58
2.618
73.83
2.491
73.06
2.153
Precision MatQuiat
4.023
66.54
2.617
74.26
2.470
73.13
2.155
MatQuaut
2
52.37
3.800
63.35
3.187
62.75
3.153
Extra Precision MatQuuat
2.052
35.70
3.355
2,050
68.25
2 2.823
2,051
65.99
2.569
MutQuant
6
67.52
2.574
6
73.92
2.440
73.63
2.127
Extra 1 Precision MatQuant
6.018
68.01
2.582
6.018
74.50
2.466
6.018
73.59
2.139
MatQuaut
61.17
2.618
72.87
71.16
Extra Precision MatQuant
3.031
63.24
2.757
3.029
73.25
2.535
3.030
71.55
2.228
21

---

Matryoshka Quantization
Table 25. Design choice ablations for co-distillation within Extra Precision MatQuant. represents distilling the y-bit model from the x-bit model. We note that the accuracy 2.050 avg. hits has significantly improved while minimally impacting the other bit-widths.
Gemma-2 9B
OmniQuant
MatQuant
E.P. Mat Quant
Avg. Bits
Config.
Task Avg.
log pplx.
Task Avg.
log pplx.
[8,4 88,4,8 2]
74.05
2.438
73.97
2.451
72.76
2.473
73.40
2.467
(8, (8,8)
[8, 2, -
73.99
2.435
73.46
2.466
[8, 8,4,2,8 + 4:
73.85
2.437
73.32
2.466
[8,4,2]
73.83
2.491
73.88
2.481
[8,4,8 +
72.65
2.519
73.84
2.488
(4,4.022)
88,42,8
73.63
2.486
73.01
2.495
[8, (8,4,2,8 4:2]
73.55
2.478
73.12
2,518
(8,4,2 [8, + 4,,
63.35
3.187
68.52
2.809
62.64
3.289
69.2
2.796
(2, 2.000
62.91
3.138
70.17
2.778
[8, 2,, 4:2]
64.32
3.227
69.72
2.804
Table 26. Table presents the downstream evaluation for Extra Precision MatQuant when applicd to OmniQuant on Gemma-2 2B.
Avg. Bits
Method
Gemma-2 2B
OmniQuant
ARC-c
ARC-c
BoolQ
HellaSwag
PIQA
Winogrande
Average
b6loat16
50.09
71.59
76.45
69.69
78.29
63.14
68.21
8
MatQuant.
49.66
71.00
76.73
68.85
78.56
63.30
68.02
8
Extra Precision MatQuant
48.04
71.8
75.78
67.64
78.07
63.22
67.42
MatQuant 47.27
70.79
73.76
66.85
78.07
62.75
66.58
4.023
Extra Procision MatQuaut
45.65
70.29
74.8
66.07
77.58
62.27
66.11
2
MatQuant.
29.95
54.21
64.40
44.37
66.81
54.46
52.37
2.052
Extra Precision MasQuant
34.39
59.64
62.69
52.11
69.86
55.56
55.71
6
MatQuant
48.89
70.50
75.69
68.89
78.40
62.75
67.52
6.018
Extra Precision MatQuant
47.1
71.46
76.02
67.47
77.91
63.61
67.26
3
MatQuant
44.03
67.09
74.25
62.78
77.26
61.40
64.47
3.031
Extra Precision MatQuant
44.45
68.56
69.11
62.28
75.95
62.59
63.82
Table 27, Table presents the downstream evaluation for Extra Precision MatQuant when applicd OmniQuant on Gemma-2 9B.
Avg.
Method
Gemma-2
OmniQuant
ARC-c
ARC-e
BoolQ
HellaSwag
PIQA
Winogrande
Average
16
58.96
77.57
83.33
77.31
81.12
67.96
74.38
8
MatQuant
57.59
77.02
84.01
76.61
81.18
67.88
74.05
8
Extra Precision MatQuant
58.11
78.03
83.27
76.17
81.18
67.09
73.97
4
MaiQuant
58.02
78.11
83.24
76.08
80.96
66.54
73.83
4.022
Extra Procision MatQuant
57.25
77.36
84.86
75.52
81.5
66.77
73.88
2
MatQuant
40.78
67.85
73.64
60.56
72.09
65.19
63.35
2.050
Extra Precision MatQuant
48.72
72.18
79.2
68.11
76.17
66.77
68.52
6
MatQuant
57, 37.25
76.94
84.04
76.63
81.34
67.32
73.92
6.018
Extra Precision MaiQuant
58.87
78.03
83.61
76.18
81.45
67.09
74.21
3
MatQuant
55.80
76.89
81.99
74.27
80.14
68.11
72.87
3.029
Extra Precision MatQuant
55.46
76.14
84.04
74.49
80.14
67.32
72.93
22

---

Matryoshka Quantization
Table 28. Table presents the downstream evaluation results for Extra Precision MalQuanl when applied to OmniQuant on Mistral 7B.
Data type
Method
Mistral 7B
OmniQuant
ARC-c
ARC-c
BoolQ
HellaSwag
PIQA
Winogrande
Average
bíloat16
49.57
73.74
84.4
80.61
81.18
74.43
73.99
8
MatQuant
49 06
72.52
84.74
79.21
81.45
74.90
73.65
8
Extra Prerision MatQuant
48.04
73.44
84.13
79.37
81.12
74.66
73.46
MatQuant 47.87
71.55
83.88
78.85
81.34
74.90
73.06
4.022
Extra Precision MatQuant
48.21
72.69
83.49
78.82
81.12
74.43
73.13
2
MatQuant
37 88
62.58
73.15
65.89
73.88
63.14
62.75
2.051
Extra Prerision MuaQuant
41.38
67.42
71.62
71.98
77.86
65.67
65.99
6
MuaQuant:
49.40
72.47
84.68
79.52
81.34
74.35
73.63
6.018
Extra Precision MatQuant
48.46
72.98
84.07
79.64
81.18
75.22
73.59
3
MatQuant
47.35
71.00
80.00
76.96
80.30
71.35
71.16
3,030
Extra Precision MatQuant
45.65
71.21
80.43
78.31
81.07
72.61
71.55
Table 29. Table presents the downstream evaluation and perplexity results for our Extra Precision Mat Quant co-distillation experiments Gemma-2 9B with OmniQuant.
OmniQuant
Gemma-2 9B
Avg. Bits
Config.
ARC-c
ARC-e
BoolQ
IeellaSwag
PIQA
Winogrande
Average
log pplx.
8,4,8 + 2] 8. 8,4,2.8 2
57.59
77.27
81.83
75.48
81.01
67.25
73.4
2.467
8
57.17
77.36
82.2
75.82
80.96
67.25
73.46
2.466
8,4,2,8 4:2]
56.4
77.82
82.32
75.02
80.63
67.72
73.32
2.466
8,4,8 2]
57.68
78.45
82.97
75.5
80.85
67.56
73.84
2.488
4.022
88,4,28 2
57.51
77.61
80.46
74.74
81.12
66.61
73.01
2.495
|8,4,2,8 - 4;22
56.57
77.99
82.54
74.77
80.58
66.3
73.12
2.518
8,4,8 2] 88,42.8 88,4,2, 4;
48.81
74.03
81.65
68.1
77.48
65.11
69.2
2.796
2.050
49.15
75.34
83.12
68.79
77.64
67.01
70.17
2.778
49.83
75.04
79.79
68.38
77.86
67.4
69.72
2.804
8,4,8 2 2]
57.42
77.19
81.87
75.42
81.01
67.8
73.45
2.468
6.018
8,4, 2]
57.51
77.48
82.32
75.88
66.61
73.48
8,4,2,8 4:
56.4
78.03
75.14
67.4
73.4
8,4,8 2]
55.63
75.88
80.12
74.01
80.36
67.96
72.33
2.549
3.029
8, 8,4.2.8 2]
54.35
76.85
79.33
74.6
80.47
67.4
72.17
2.543
8,4, 4 4:2]
55.2
76.98
82.45
T3.59
80.41
08.43
72.84
2.58
Table 30. Results comparing Single Precision MatQuant and Single Precision Extra Precision MatQuant when applied to Gemma- 2 2B, 9B and Mistral 7B with OmniQuant as the base algorithm.
int2
Gemma-2 2B
Gemma-2 9B
Mistral 7B
Method
Task Avg.
log pplx.
Task Avg.
log pplx.
Task Avg.
log pplx.
OnniQuant
51.33
3.835
60.24
3.292
59.74
3.931
S.P. MatQuant
53.42
3.631
64.02
3.171
63.58
2.976
MatQuant
52.37
3.800
63.35
3.187
62.75
3.153
MatQuant
3.185
68.58
2.857
67.36
2.464
E.P. MatQuant
55.71
3.292
68.52
2.809
65.99
2.569
23


Reviews:
[{"review_id": "thibr2JvMW", "reviewer": "Reviewer_ESHi", "summary": "The paper introduces Matryoshka Quantization (MatQuant), a novel multi-scale quantization technique designed to jointly learn various bitness representations in one single training, and improve low-bit precision quantization for LLMs. \nThe method is evaluated on Gemma-2 and Mistral models, and the results demonstrate better perplexity and downstream performance compared to the standard QAT baseline. The paper also demonstrates some variations of MatQuant, including Single Precision and quantizing both FFN and Attention.\n\n**Update after rebuttal**: My latest reply reflected my final update.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful review. Below are our responses to the comments/questions:\n\n**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. \n\n**Training curves:** Our experiments indicate that MatQuant requires less samples/iterations to converge to the same training perplexity. Specifically, \n1. 2-bit QAT on Gemma2 9B: MatQuant and S.P. Matquant achieves same training perplexity with 20M tokens as the baseline with 100M tokens\n2. 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.\n\n**Single Precision (S. P.) MatQuant vs Baseline:**\n1. 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. \n2. 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.\n\n**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. \n\nWe 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.", "claims_and_evidence": "[Claims are supported]\n\n* MatQuant maintains a single quantized model but serves it with different precision demands.\n* Int2 precision models of MatQuant achieve bigger improvement.\n* MatQuant can interpolate int3 and int6 representations without explicitly training on them.", "methods_and_evaluation_criteria": "Yes. The models (Gemma, Mistral), evaluation tasks (C4, ARC, BoolQ, HellaSwag,...), and baselines (QAT, OmniQuant) make sense.", "theoretical_claims": "No theoretical claims.", "experimental_designs_or_analyses": "Yes.\n\n* The main results on int2 are validated.\n* There are experiments to ablate the hyperparameter $\\lambda$.\n* Single Precision experiment shows better performance.\n\n[Weakness]\n\n* The paper doesn't show the figure of performance vs training samples (or training time) for MatQuant. MatQuant might require more (or maybe fewer) samples or time to converge. \n* Isn't Single Precision MatQuant the same as standard QAT because it set $\\lambda$ for 4 and 8 bits to 0? Why does it still show performance improvement?\n* Maybe just me, but I don't understand the config and datatype of Table 4. Why there are two 8 in \"[8, 4, 2, 8] -> [4;2]\", and what does 4; 2 mean?", "relation_to_broader_scientific_literature": "1. The paper is related to Matryoshka-style training.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PcP7zdxHwZ", "reviewer": "Reviewer_66kQ", "summary": "This paper introduces Matryoshka Quantization, a novel multi-scale quantization technique that enables single-model multi-bitwidth operation. The low-precision models extracted by MatQuant can have significant improvement compared with standard quantization.\n\n## update after rebuttal\nThe authors address most of my concerns, and I keep my original rating to lean toward accepting the paper.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful review. Below are our responses to the comments/questions:\n\n**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. \n\n**Training curves:** Our experiments indicate that MatQuant requires less samples/iterations to converge to the same training perplexity. Specifically, \n1. 2-bit QAT on Gemma2 9B: MatQuant and S.P. Matquant achieves same training perplexity with 20M tokens as the baseline with 100M tokens\n2. 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.\n\n**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.\n\n**Plugging MatQuant on the latest methods:** we thank the reviewer for referring the latest QAT style techniques. \nMost 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. \nMatQuant 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. \n\n**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.\n\nTable: 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.\n\n| DType | Method      | ARC-c | ARC-e | BoolQ | HellaSwag | PIQA  | Winogrande | Avg. | log pplx. |\n|-----------|-------------|-------|-------|-------|-----------|-------|------------|---------|-----------|\n| bf16  |             | 60.49 | 78.96 | 83.18 | 82.24     | 84.60 | 75.69      | 77.53   | 2.199     |\n| int8      | Baseline    | 60.49 | 79.12 | 80.34 | 83.15     | 84.60 | 76.01      | 77.28   | 2.169     |\n|      | MatQuant    | 60.41 | 79.88 | 78.41 | 82.54     | 85.04 | 75.77      | 77.01   | 2.141     |\n| int4      | Sliced int8 | 59.30 | 77.90 | 83.94 | 81.69     | 83.68 | 73.88      | 76.73   | 2.232     |\n|      | Baseline    | 59.39 | 79.38 | 83.79 | 82.45     | 84.44 | 75.30      | 77.46   | 2.250     |\n|       | MatQuant    | 59.73 | 78.96 | 76.06 | 81.68     | 84.49 | 74.82      | 75.96   | 2.201     |\n| int2      | Sliced int8 | 25.85 | 27.99 | 55.72 | 25.07     | 51.25 | 49.41      | 39.21   | 15.601    |\n|       | Baseline    | 32.25 | 55.56 | 67.58 | 55.50     | 69.59 | 59.04      | 56.59   | 3.061     |\n|       | MatQuant    | 44.8  | 71.25 | 70.31 | 67.61     | 76.33 | 64.64      | **65.82**   | **2.674**     |\n| int6      | Sliced int8 | 60.15 | 79.12 | 81.31 | 83.13     | 84.60 | 76.64      | 77.49   | 2.111     |\n|      | Baseline    | 60.41 | 79.84 | 80.18 | 82.77     | 84.60 | 76.48      | 77.38   | 2.170     |\n|       | MatQuant    | 59.98 | 79.21 | 75.57 | 82.4      | 84.49 | 75.45      | 76.18   | 2.146     |\n| int3     | Sliced int8 | 52.05 | 72.69 | 63.73 | 73.42     | 79.98 | 69.46      | 68.55   | 2.797     |\n|       | Baseline    | 57.17 | 79.38 | 73.61 | 78.91     | 83.24 | 72.85      | 74.19   | 2.373     |\n|       | MatQuant    | 54.95 | 76.18 | 63.18 | 77.36     | 82.37 | 72.30      | 71.06   | 2.494     |\n\nWe 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.", "claims_and_evidence": "The claims made in the submission are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The method makes sense for the problem.", "theoretical_claims": "No theory result in this paper.", "experimental_designs_or_analyses": "Lack of inference runtime benchmarks, which are critical for deployment. Also, does MatQuant need more GPU memory and time during training? It's better to provide the comparison results.", "relation_to_broader_scientific_literature": "Unlike existing works that need to optimize for each precision, MatQuant can train and maintain a single quantized model but serve it with the precision demanded by the deployment.", "essential_references_not_discussed": "Though QAT and OmniQuant are effective tools for quantization, there are some works like that show better performance than those. Considering the results of MatQuant+novel QAT methods can strengthen the paper's contribution.\n\nBondarenko Y, Del Chiaro R, Nagel M. Low-rank quantization-aware training for llms.\n\nAlso, as mentioned in the previous section. It's unclear whether MatQuant introduces more GPU memory usage or training time. There are several works that focus on the efficiency of QAT.\n\n\nChen M, Shao W, Xu P, Wang J, Gao P, Zhang K, Luo P. Efficientqat: Efficient quantization-aware training for large language models.", "other_strengths_and_weaknesses": "There is one more weakness: the main results are on LM <10 B. I have a scalability concern about how MatQuant performs on large models.", "other_comments_or_suggestions": "No extra comments.", "questions_for_authors": "1. What is the runtime overhead of MatQuant vs. standard QAT methods?\n2. How does MatQuant perform on very large-scale models?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9PUL5W10r4", "reviewer": "Reviewer_sotV", "summary": "The paper presents a method for multi-scale quantization of large language models across multiple precisions (int8, int4, int2). By utilizing the nested structure of integer data types, the proposed technique allows different precision levels to be nested within one another. The resulting quantized model can then be served at different precisions based on deployment requirements.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thorough feedback. Below are our responses to the comments/questions:\n\n**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.\n \nFor 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. \n\nAQLM 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).\n\nTable: 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.\n\n| Config.                                                  | Method                              | ARC-c | ARC-e | BoolQ | HellaSwag | PIQA  | Winogrande | Average |\n| :------------------------------------------------------- | :---------------------------------- | ----: | ----: | ----: | --------: | ----: | ---------: | ------: |\n| 222222444444444444444444444444444444222222 = 3.43 avg bits | Mixnmatch using Baseline 2, 4 bit | 30.03 | 48.02 | 62.26 |     44.47 | 67.30 |      59.75 |   51.97 |\n| 222222444444444444444444444444444444222222 = 3.43 avg bits | Mixnmatch using MatQuant 2, 4 bit   | 52.56 | 79.04 | 78.99 |     75.66 | 80.20 |      69.38 |   **72.64** |\n| 222444444444444444444444444444444444444222 = 3.71 avg bits | Mixnmatch using Baseline 2, 4 bit | 31.48 | 50.34 | 62.14 |     43.97 | 67.19 |      60.30 |   52.57 |\n| 222444444444444444444444444444444444444222 = 3.71 avg bits | Mixnmatch using MatQuant 2, 4 bit   | 56.83 | 81.06 | 80.73 |     76.34 | 81.07 |      67.88 |   **73.98** |\n| int4                                                     | Baseline                            | 58.79 | 78.37 | 83.55 |     76.71 | 81.45 |      67.09 |   74.33 |\n|                                                     | MatQuant                            | 57.25 | 77.36 | 84.86 |     75.52 | 81.50 |      66.77 |   73.88 |\n| int2                                                     | Baseline                            | 39.16 | 63.43 | 72.11 |     52.24 | 72.63 |      61.88 |   60.24 |\n|                                                     | MatQuant                            | 48.72 | 72.18 | 79.20 |     68.11 | 76.17 |      66.77 |   68.52 |\n\n**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.\n\n**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.\n\nWe 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.", "claims_and_evidence": "The paper correctly claims comparable accuracy of the proposed method and existing int4 and int8 quantization schemes, and improved int2 performance, although the experiments were restricted to the C4 dataset.", "methods_and_evaluation_criteria": "Use of the C4 dataset is meaningful but limiting; including more comprehensive results on e.g. the Pile or at least WikiText-2 is strongly recommended.", "theoretical_claims": "This is an empirical paper with no major theoretical results.", "experimental_designs_or_analyses": "The experiments are appropriately constructed and conducted.", "relation_to_broader_scientific_literature": "The proposed methods belongs to the category of Learning-Based Quantization Methods and can work with other such (existing) techniques, including QAT and OmniQuant.", "essential_references_not_discussed": "The proposed MatQuant outperforms existing methods only at the extreme int2 quantization levels. It is therefore critically important to compare and contrast MatQuant against other schemes specifically targeting such regimes. The work that comes to mind is Egiazarian et al, \"Extreme Compression of Large Language Models via Additive Quantization\", ICML '24, which reports strong performance in 2-bit (vector) quantization settings.", "other_strengths_and_weaknesses": "On the one hand, when considering quantization levels higher than int2, the proposed method does not perform as strongly as the existing (quantization level-specific) techniques. On the other hand, at int2 MatQuant does work better than the considered competing (scalar quantization) alternatives but the performance is much weaker than that of higher quantization level schemes. To assert the utility of MatQuant, the authors should compare and contrast it with SOTA methods for low-bitwidth quantization schemes (e.g., AQLM in the above mentioned reference).", "other_comments_or_suggestions": "None.", "questions_for_authors": "A clarifying question: It is stated that upon quantization by MatQuant, the weights shift to higher values. Does the choice of quantization bins agree with information-theoretically optimal approach to scalar quantization utilized by e.g. Normal Float Quantization method?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "None."}}]

Author rebuttal:

Meta-review: The paper introduces Matryoshka Quantization (MatQuant), a novel multi-scale quantization technique designed to jointly learn various bit-width quantized network representations in one single training, and improve low-bit precision quantization for LLMs. The method is evaluated on Gemma-2 and Mistral models, and the results demonstrate better perplexity and downstream performance compared to the standard QAT baseline. 

Two of the three reviewers were mildly positive with one reviewer slightly negative. 

The negative reviewer only expressed one issue, however, that of the proposed MatQuant outperforming existing methods only at the extreme int2 quantization level, and asked for comparing MatQuant against other schemes specifically targeting such regimes. The authors correctly respond 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. The 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.

When training multiple bit-widths a single MatQuant training uses up to 30% less GPU hours, and is more effective than having separate training runs for different bit-widths. Overall, the data provided in the rebuttal answered the reviewers' concerns about training time. For example, experiments indicate that MatQuant requires less samples/iterations to converge to the same training perplexity. Similarly, rebuttal data indicate that the method scales to even larger models than mentioned in the original submission.
Decision: Accept (poster)