Matryoshka Quantization (2502.06786v3)

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 regularization, int2 precision models extracted by MatQuant outperform standard int2 quantization by up to 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 gives an additional 6% improvement with OmniQuant as the base algorithm.

Matryoshka Quantization: A Multi-Scale Approach to Efficient Model Deployment

The paper "Matryoshka Quantization" addresses a critical challenge in deploying large machine learning models: optimizing models for different quantization levels without compromising accuracy. Traditionally, quantizing models to low precisions like int4 or int2 results in a significant loss in model quality. This paper introduces a novel multi-scale quantization approach leveraging the inherent nested (Matryoshka) structure of integer data types. The proposed technique, referred to as Matryoshka Quantization, trains a single robust model capable of operating effectively across multiple precision levels.

Background and Motivation

Quantization reduces model size and inference costs by converting high-precision weights into lower-precision representations. For instance, representing model weights as integers (int8, int4, int2) instead of floating point significantly minimizes the data transfer and storage requirements. However, lower precision can degrade model accuracy, particularly in large models like LLMs.

Current quantization methods usually treat each precision level independently, resulting in different models for each target precision. This approach is inefficient as it demands maintaining and serving multiple model versions. Matryoshka Quantization leverages the nested nature of integer types — akin to Russian Matryoshka dolls — where smaller integer widths are nested within larger ones. This enables a model trained with Matryoshka Quantization to be used at different precision levels by slicing the most significant bits.

Methodology

The core of Matryoshka Quantization is simultaneous optimization for multiple bit-widths via a multi-scale training method. This involves:

• Bit-Width Nesting: Training a model where quantized weights are represented at various precision levels using shared most significant bits (MSBs). Smaller bit-width models can be obtained from a larger model by right-shifting to extract the MSBs.
• Loss Optimization: The training framework jointly optimizes the quantization loss across target bit-widths, balancing accuracy across all specified precision levels.

The technique is general-purpose and adaptable to quantization schemes like Quantization Aware Training (QAT) and OmniQuant, effectively making it suitable for most learning-based quantization methods.

Results

Experimental results demonstrate the effectiveness of Matryoshka Quantization through applications on various transformer-based LLMs (e.g., Gemma-2 2B, 9B, and Mistral 7B models). Key findings include:

• Accuracy Preservation: The int8 and int4 models maintain comparable accuracy to independently trained baselines. Notably, the int2 models exhibit up to a 10% improvement in accuracy compared to existing quantization approaches.
• Interpolative Capabilities: Beyond targeted precision levels, the models exhibit strong interpolative behaviour for bit-widths like int6 and int3, performing similarly to explicitly trained baselines.
• Elastic Models through Mix'n'Match: The method allows for layer-wise precision adjustments without additional training, supporting a dense accuracy-cost trade-off and efficient deployment settings.

Implications and Future Directions

The introduction of a single model capable of functioning effectively across multiple precision levels has notable implications:

• Practical Deployment: The ability to switch precision dynamically could lead to more adaptable AI systems, especially in resource-constrained environments.
• Hardware Co-Design Opportunities: As the need for elastic quantization grows, the research opens pathways for hardware optimizations to support varying precision on-the-fly.
• Floating-Point Extensions: While the current method leverages the integer nature, extending Matryoshka-like optimization to floating points could further enhance training and deployment efficiency.

This work potentially transforms approaches to quantization by proposing a flexible framework that minimizes the need for maintaining multiple models, greatly enhancing the efficiency of deploying large-scale models. The ongoing exploration of co-distillation strategies and the possibility of extending to floating-point representations underscore the paper's contribution to advancing quantization strategies for modern AI systems.