Matryoshka Quantization | Summary

This article explains the key points of Matryoshka Quantization.

• 2025-02-10 (arXiv)
• Nair, Pranav, Datta, Puranjay, Dean, Jeff, Jain, Prateek, Kusupati, Aditya.
• Google DeepMind
• Paper

Read this article in Korean

Summary

• Quantization, representing model weights with fewer discrete values, reduces storage requirements but can lower accuracy and require separate models for different deployment budgets. With only feed-forward network (FFN) weights quantized, Matryoshka Quantization (MatQuant) improves average downstream accuracy over independently trained 2-bit baselines by 1.04–3.11 percentage points with OmniQuant, which learns quantization parameters while keeping model weights fixed, and 4.46–7.02 percentage points with Quantization Aware Training (QAT), which updates weights using quantized forward passes.
• MatQuant jointly optimizes int8, int4, and int2, integer representations using 8, 4, and 2 bits, through shared most significant bits, the bits encoding the larger place values. The trained representation also supports extraction of int3 and int6 models without additional training and Mix’n’Match, the assignment of different precisions to different layers.
• The principal experiments quantize the FFN weights, the parameters in each Transformer layer’s feed-forward block, in Gemma-2 2B, Gemma-2 9B, and Mistral 7B. Evaluation averages zero-shot accuracy, performance without task-specific examples, across six downstream tasks and reports log perplexity on C4, a text dataset; lower log perplexity indicates better prediction of the evaluation text.
• The Errata distinguishes strict 2-bit MatQuant from Extra Precison MatQuant, which retains an additional quantization value and requires approximately 2.05 bits per quantized weight on average. On Gemma-2 9B with OmniQuant, the latter reaches 70.17 average accuracy with co-distillation, training a lower-precision model against higher-precision outputs, but requires custom handling of the additional bits; the paper reports no measured deployment speedups.

1. Introduction

The introduction motivates quantization through the communication cost of transferring model weights during autoregressive decoding, where a model generates tokens sequentially. Conventional optimization treats each precision independently, whereas MatQuant asks whether one shared representation can provide several accurate deployment options.

• Lower-precision models are extracted from a common int8 representation.
• The main experiments quantize FFN weights; quantizing attention weights, which govern how tokens exchange information, is examined separately.

The nested-bit diagram shows how several precision losses act on shared integer weights, while the histograms show that joint training changes quantized bucket occupancy. The caption’s greater-than-8% int2 claim does not match the corrected strict-int2 comparison in Table 1 and should be interpreted with the extra-precision distinction in the Errata.

2. Related Work

The related work distinguishes learning-free quantization, which calibrates representations without gradient descent, from learning-based methods that optimize model weights or auxiliary quantization parameters. It connects MatQuant to multi-scale training, where nested models of different capacities share parameters and are trained jointly.

• Prior methods reduce reconstruction error, the difference between original and quantized outputs, or handle outliers, weights or activations with unusually large values. MatQuant adds joint optimization across nested precisions.
• The paper acknowledges earlier any-precision neural networks and extends nested precision to large language models, including extraction at untrained intermediate precisions.

3. Matryoshka Quantization

MatQuant adds multi-precision training to a learning-based quantizer while retaining its underlying quantization rule. It combines losses from several bit-widths and preserves the base method’s choice of trainable parameters and learning objective.

• The demonstrated integrations are QAT and OmniQuant.
• The stated applicability requires model or auxiliary parameters to be optimized with gradient descent; the experiments test these two quantizers.

3.1. Preliminaries

QAT minimizes end-to-end cross-entropy, a loss measuring prediction error against target tokens, using quantized weights in the forward pass and a straight through estimator (STE), an approximation that propagates gradients through quantization. Its MinMax quantizer maps the observed weight range into a fixed number of integer values using a scale and zero point, the offset locating zero in that representation.

• OmniQuant freezes the original model weights and learns clipping, scaling, and shifting parameters by minimizing squared differences between full-precision and quantized layer outputs.
• OmniQuant optimizes each Transformer layer independently with limited calibration data, while QAT updates model weights against the end-to-end prediction objective.

3.2. MatQuant

For a c-bit quantized vector qᶜ and 0 < r < c, MatQuant extracts an r-bit representation with S(qᶜ,r) = clamp(round(qᶜ/2^(c−r)), 0, 2^r−1) × 2^(c−r). Rounding selects the lower-precision value, clamping restricts it to the r-bit range, and rescaling places the result on the original quantized scale.

• The training objective sums the base quantizer’s loss over examples and target bit-widths, weighting each precision by λᵣ; the default target set is R = {8,4,2}.
• Interpolative Behavior includes slicing at untrained int3 and int6 precisions and Mix’n’Match across layers.
• Joint training shifts occupancy toward higher-valued quantized buckets, the discrete values available to the quantizer. The authors hypothesize that the int8 representation provides flexibility to accommodate the more restrictive int2 representation.

4. Experiments

The experiments use Gemma-2 2B, Gemma-2 9B, and Mistral 7B, primarily with weight-only quantization of FFN blocks. OmniQuant uses 128 C4 examples of sequence length 2048, batch size 4, and 10M training tokens, except that the independently trained int2 baseline receives 20M tokens; QAT uses 100M C4 tokens, batch size 16, sequence length 8192, and one epoch.

• The principal baselines are models trained independently for each target precision and models obtained by slicing a conventionally trained int8 representation.
• MatQuant’s int3 and int6 models receive no additional training, whereas their independent baselines are explicitly trained at those precisions.
• The six evaluation tasks are ARC-c and ARC-e, science question answering; BoolQ, yes/no question answering; HellaSwag, sentence completion; PIQA, physical commonsense reasoning; and Winogrande, pronoun resolution. The prose identifies C4’s test set for perplexity, while table captions identify its validation set.

4.1. MatQuant with OmniQuant

With OmniQuant, strict int2 MatQuant increases average downstream accuracy from 51.33 to 52.37 on Gemma-2 2B, 60.24 to 63.35 on Gemma-2 9B, and 59.74 to 62.75 on Mistral 7B. These are gains of 1.04, 3.11, and 3.01 percentage points over independently trained int2 models.

• The int8 and int4 results remain close to their independently trained counterparts, although Table 1 includes gaps of 0.54 and 0.56 percentage points, slightly exceeding the prose’s stated 0.5% bound.
• Untrained int3 extraction yields 64.47, 72.87, and 71.16 average accuracy, compared with 64.37, 73.23, and 71.68 for independently trained int3 baselines.
• Slicing a conventional int8 model performs much worse at int2: 39.78, 38.11, and 37.29 average accuracy.

The independently trained, jointly trained, and sliced-int8 rows compare explicit nested optimization with separate training and post-training extraction. Strict int2 improves on all three models, while int3 and int6 extraction remains close to independent training without additional optimization.

4.2. MatQuant with QAT

With QAT, strict int2 MatQuant improves average downstream accuracy from 47.74 to 52.20 on Gemma-2 2B, 56.02 to 62.29 on Gemma-2 9B, and 54.95 to 61.97 on Mistral 7B. The resulting gains are 4.46, 6.27, and 7.02 percentage points.

• Intermediate-precision extraction does not uniformly match independent training: Mistral 7B int3 reaches 67.16 versus 68.82 for its baseline. At int8, Mistral 7B reaches 72.58 versus 73.48.
• The authors interpret lower C4 log perplexity alongside weaker downstream accuracy than OmniQuant as evidence of overfitting to the C4 training subset. This explanation is inferred from metric patterns, and those patterns have exceptions across the reported settings.
• Table 2 reports Gemma-2 9B int2 MatQuant log perplexity as 2.265, whereas Tables 4, 5, 24, and 30 report 2.660 for the corresponding non-co-distilled result.

Strict int2 accuracy gains are 4.46, 6.27, and 7.02 percentage points across the three models. Intermediate-precision rows reveal exceptions to equivalent quality, including the 1.66-point Mistral 7B int3 gap; the Gemma-2 9B int2 log-perplexity entry also differs from later tables.

4.3. Layerwise Mix’n’Match

Layerwise Mix’n’Match constructs models by assigning int2, int4, or int8 to different layers without retraining those combinations. For Gemma-2 9B with OmniQuant, selected mixed configurations fill the accuracy-versus-effective-bits curve, and a configuration below 4 effective bits per FFN parameter outperforms the homogeneous int4 slice.

• Among the explored allocation strategies, Pyramid uses higher precision in middle layers and lower precision near the beginning and end.
• Effective bits per FFN parameter measures the average precision assigned to those weights. The figure reports selected configurations and does not measure runtime latency or establish optimality over every layer assignment.

Selected mixed-precision configurations fill gaps between homogeneous precision choices, including configurations below 4 effective bits that exceed the int4 slice’s accuracy. Effective bits per FFN parameter represents a weight-storage budget and does not directly measure inference latency.

5. Ablations and Discussion

The ablations examine precision-loss weighting, supervision between nested models, and optimization of a single extracted precision. The discussion also examines FFN plus attention weight quantization, hardware support, and possible extensions to floating-point representations.

5.1. Weightings (𝜆𝑟) for MatQuant

Changing λᵣ adjusts the training priority of each precision and creates trade-offs among the extracted models. The results generally favor a relatively large int2 loss weight, but the response varies by model and is non-monotonic.

• For Gemma-2 2B, changing the loss weights from (0.1,0.1,1) to (0.4,0.4,1) raises int4 accuracy from 66.58 to 67.48 while reducing int2 accuracy from 52.37 to 51.69.
• For Gemma-2 9B, int2 accuracy reaches 64.1 at (0.3,0.3,1), compared with 63.35 at (0.1,0.1,1), so increasing int2’s relative priority does not invariably improve its accuracy.

The sweep shows competition among precision objectives: increasing higher-precision loss weights can reduce int2 accuracy. The best tested weighting varies by model, and the response is non-monotonic.

5.2. Co-distillation for MatQuant

Co-distillation uses outputs from the nested int8 model to supervise lower-precision models, either replacing their original targets or supplementing them with equal weighting. On Gemma-2 9B with OmniQuant, supervising both int4 and int2 raises strict int2 accuracy from 63.35 to 64.32, a 0.97-percentage-point gain.

• Supervising only int2 while retaining its original objective gives 62.91 with OmniQuant, below the non-co-distilled result.
• With QAT, the three tested co-distillation configurations give int2 accuracies of 62.31, 62.70, and 62.60, compared with 62.29 without co-distillation.

Co-distillation helps selectively. With OmniQuant, supervising both lower precisions produces 64.32 int2 accuracy, whereas supplementing only int2’s objective produces 62.91, below the 63.35 non-co-distilled result.

5.3. Single Precison MatQuant

Single Precison MatQuant removes the int4 and int8 loss terms while continuing to obtain int2 by slicing an int8 representation. With OmniQuant, it reaches 53.42, 64.02, and 63.58 average accuracy on the three models, exceeding multi-precision MatQuant by up to 1.05 percentage points.

• The authors hypothesize that the six additional underlying bits provide optimization flexibility even when only the extracted int2 model is supervised; the experiments do not isolate this mechanism.
• Higher-precision slices can deteriorate: on Gemma-2 9B with QAT, Single Precison MatQuant gives 72.71 at int8 and 71.67 at int4, versus 74.52 and 73.24 for multi-precision MatQuant.
• For FFN plus attention quantization, independent QAT is unstable at int2 and int3. MatQuant variants produce evaluable models, but strict int2 accuracy remains low: 44.19 for Gemma-2 9B and 38.36 for Mistral 7B with multi-precision MatQuant.

Single Precison MatQuant exceeds the independent int2 baseline in every displayed setting and exceeds multi-precision MatQuant in all three OmniQuant settings. Its advantage over multi-precision training is not universal under QAT.

Missing independent-QAT entries at int2 and int3 correspond to reported instability when both FFN and attention weights are quantized. MatQuant variants yield evaluable results, but their low int2 accuracies show that stability does not ensure preserved model quality.

5.4. Deployment Considerations

Deployment can select a homogeneous sliced model or a mixed-precision layer allocation according to memory limits and available integer kernels. The paper discusses native int8 and int4 support and custom CUDA kernels, GPU implementations of low-precision operations, for formats such as int2 and int3.

• A mixed int2/int4 configuration can approximate an int3 weight-storage budget when optimized int3 support is unavailable.
• The experiments establish accuracy at different weight-precision budgets; throughput, latency, and end-to-end memory savings are not benchmarked.

5.5. Extension to Floating Point

Extending MatQuant to FP8 and FP4, 8-bit and 4-bit floating-point representations, is proposed as future work. Their exponent bits encode powers of two, so slicing them does not preserve the constant spacing between representable values used by the integer construction.

• The integer example extracts values 0, 64, 128, and 192 from int8, with constant spacing 64.
• The paper provides no floating-point MatQuant experiments or validated floating-point slicing algorithm.

6. Conclusions

The conclusions emphasize one shared integer representation supporting several trained precisions, untrained intermediate precisions, and layerwise combinations. The largest demonstrated accuracy advantage over independently trained baselines is at strict int2, while int4 and int8 generally remain close to those baselines.

• Single Precison MatQuant improves some fixed-precision results, showing that nested optimization can help when deployment needs only one precision.
• The evidence covers the tested models and training data, primarily with FFN-only weight quantization.

7. Errata

The Errata explains that the first draft omitted clamping after rounding, permitting \(2^r+1\) quantized values rather than \(2^r\). This variant, named Extra Precison MatQuant, needs an additional bit for the small fraction of weights assigned to the extra value and has a different representation budget from strict r-bit MatQuant.

• For unsigned int8 value 234 and extraction to int2, rounding gives 4; strict MatQuant clamps this to 3 and rescales to 192, while the unclamped variant produces 256.
• Table 7 reports Gemma-2 9B accuracy rising from 63.35 at 2 bits to 68.25 at 2.050 average bits. Tables 8 and 27 instead report 68.52 for the non-co-distilled extra-precision model; the PDF does not reconcile these values.
• Table 8 reports 70.17 at 2.050 average bits when int8 supervises int2 alongside its original objective. Figure 3 applies mixed layer allocation to the extra-precision model, but sparse additions of the extra bits require a custom CUDA kernel.

Average-bit columns distinguish strict integer precision from the extra-value representation. The largest gains occur near 2 bits, including Gemma-2 9B increasing from 63.35 at 2 bits to the displayed 68.25 at 2.050 bits; later tables give 68.52 for the non-co-distilled extra-precision model.

The best displayed extra-precision int2 configuration reaches 70.17 at 2.050 average bits with int8-to-int2 supervision alongside the original objective. This condition differs from both strict int2 MatQuant and the non-co-distilled extra-precision model.

The stronger low-bit extra-precision model raises the displayed mixed-precision accuracy-versus-storage curve. Some selected configurations approach or exceed the bfloat16 reference, the model’s 16-bit floating-point representation, but the plot does not quantify the implementation cost of sparse extra-bit corrections.

Impact Statement

The Impact Statement describes potential benefits from reducing the resources needed to deploy large language models, including broader accessibility and more sustainable systems. These benefits are prospective: the paper does not measure energy consumption, environmental impact, or changes in access, and it acknowledges general risks associated with wider deployment of powerful models.

Appendix

• A. Particulars of the Slicing Operation. clarifies that extraction includes rounding rather than simply discarding lower bits. Its examples show why the next bit affects rounding and why values rounded outside the target integer range require clamping.
• B. Addition Training Details reports TPUv5e accelerator hardware, an OmniQuant learning rate of 1e−3, and a QAT learning rate warmed to 1e−5 followed by cosine decay. OmniQuant co-distillation uses 8.3M tokens, while FFN plus attention QAT uses 300M C4 tokens, batch size 16, and one epoch. The same appendix gives loss weights (λ₈,λ₄,λ₂) = (0.1,0.1,1.0) for Gemma models unless otherwise stated, (0.4,0.4,1.0) for Mistral 7B with OmniQuant, and (0.2,0.2,1.0) for Mistral 7B with QAT; Extra Precison MatQuant uses (1.0,1.0,1.0). It compares Pyramid, Reverse Pyramid, Increasing, and Decreasing layer allocations and reports that Pyramid consistently performs best among those tested strategies.
• C. Detailed Downstream Evaluations for OmniQuant and QAT provides task-level results in Tables 9–14. Gains are not uniform across tasks: Gemma-2 2B int2 with OmniQuant improves the six-task average from 51.33 to 52.37 while ARC-c decreases from 31.31 to 29.95.
• D. Detailed Downstream Evaluations for MatQuant Re-weighting expands the loss-weight ablations in Tables 15–17. Its task-level and intermediate-precision results show that changing precision priorities affects both directly supervised and interpolated models.
• E. Detailed Downstream Evaluations for Co-Distillation gives task accuracy and log perplexity for Gemma-2 9B in Tables 18 and 19. The strongest tested strict int2 OmniQuant average is 64.32 when int8 supervises both lower precisions.
• F. Detailed Evaluations for FFN + Attention Quantization provides Tables 20 and 21, both explicitly labeled as QAT experiments, although the appendix introduction mentions both OmniQuant and QAT. They expose remaining low-precision difficulty: Mistral 7B multi-precision MatQuant reaches 40.55 at int3, whereas Single Precison MatQuant reaches 67.59.
• G. Detailed Evaluation for Single Precison MatQuant provides Tables 22–25 and a weight-distribution plot comparing its quantized buckets with the baseline. The results document improvements at the supervised int2 precision and losses in higher-precision slices, showing the trade-off between optimizing one precision and preserving several.
• H. Detailed Evaluation for Extra Precison MatQuant provides Tables 26–29, including task-level extra-precision and co-distillation results. Table 30 additionally compares single-precision and extra-precision combinations: with OmniQuant, S.P. E.P. MatQuant, the single-precision extra-precision variant, reaches 57.38, 68.58, and 67.36 average accuracy, while extra precision does not consistently improve QAT results; these variants retain the extra-value qualification described in the Errata.

Brief Thoughts

The clearest result is improved strict int2 accuracy against independently trained models across all three tested models and both base quantizers. The OmniQuant improvement holds despite its int2 baseline receiving more training tokens; intermediate slicing and mixed layer allocation also provide additional accuracy-versus-weight-storage choices. The Errata changes the interpretation of the low-bit results: strict 2-bit and approximately 2.05-bit models use different representation budgets, and the extra-value variant provides additional gains with OmniQuant. Deployment value remains unmeasured. Runtime benchmarks would need to account for kernel support and extra-bit handling, while the low accuracies in some FFN plus attention settings and inconsistent table entries constrain broader claims.