Matryoshka Quantization

TL;DR

Matryoshka Quantization introduces a multi-scale quantization method that enables training a single model to be served at various precisions, significantly improving low-bit quantization quality and flexibility.

Contribution

The paper proposes Matryoshka Quantization, a novel nested quantization technique that enhances low-precision model performance and allows flexible deployment at different precisions from one trained model.

Findings

Int2 models with MatQuant outperform standard int2 quantization by up to 7%.

Using an extra bit for outliers yields a 6% improvement at 2.05-bit precision.

MatQuant enables a single model to serve multiple precisions effectively.

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.