Efficient Quantization of Mixture-of-Experts with Theoretical Generalization Guarantees

TL;DR

This paper introduces a theoretically grounded mixed-precision quantization method for Mixture-of-Experts models, improving accuracy and reducing inference costs by adaptively assigning bit-widths based on expert sensitivity.

Contribution

It proposes a novel expert-wise mixed precision strategy grounded in theory, considering experts' training sensitivity and variance for better quantization.

Findings

Achieves higher accuracy than existing quantization methods.

Reduces inference cost with negligible overhead.

Effective on large-scale MoE models like Switch Transformer and Mixtral.

Abstract

Sparse Mixture-of-Experts (MoE) allows scaling of language and vision models efficiently by activating only a small subset of experts per input. While this reduces computation, the large number of parameters still incurs substantial memory overhead during inference. Post-training quantization has been explored to address this issue. Because uniform quantization suffers from significant accuracy loss at low bit-widths, mixed-precision methods have been recently explored; however, they often require substantial computation for bit-width allocation and overlook the varying sensitivity of model performance to the quantization of different experts. We propose a theoretically grounded expert-wise mixed precision strategy that assigns bit-width to each expert primarily based on their change in routers l2 norm during training. Experts with smaller changes are shown to capture less frequent but critical features, and model performance is more sensitive to the quantization of these experts, thus requiring higher precision. Furthermore, to avoid allocating experts to lower precision that inject high quantization noise, experts with large maximum intra-neuron variance are also allocated higher precision. Experiments on large-scale MoE models, including Switch Transformer and Mixtral, show that our method achieves higher accuracy than existing approaches, while also reducing inference cost and incurring only negligible overhead for bit-width assignment.