Table of Contents
Fetching ...

MiLo: Efficient Quantized MoE Inference with Mixture of Low-Rank Compensators

Beichen Huang, Yueming Yuan, Zelei Shao, Minjia Zhang

TL;DR

MiLo tackles the heavy memory and accuracy penalties of extreme quantization in MoE models by introducing a calibration-free Quantize-Then-Compensate pipeline that augments INT3 MoEs with a mixture of adaptive low-rank compensators. The method jointly optimizes quantized weights and residuals via an iterative framework, using sub-problems for quantization error and low-rank approximation, and employs adaptive rank policies informed by MoE structure and data distribution. It further delivers hardware-friendly INT3 kernels with zero-bit-waste packing and efficient dequantization to realize real latency gains. Across state-of-the-art MoEs like Mixtral-8×7B and DeepSeek-MoE, MiLo achieves substantial accuracy recovery (often >87% of lost accuracy) with modest memory overhead and up to 3× speedups compared with baselines. The work demonstrates practical viability of highly compressed MoEs on real hardware, enabling scalable deployment without calibration data and with Tensor Core-friendly kernels.

Abstract

A critical approach for efficiently deploying Mixture-of-Experts (MoE) models with massive parameters is quantization. However, state-of-the-art MoE models suffer from non-negligible accuracy loss with extreme quantization, such as under 4 bits. To address this, we introduce MiLo, a novel method that augments highly quantized MoEs with a mixture of low-rank compensators. These compensators consume only a small amount of additional memory but significantly recover accuracy loss from extreme quantization. MiLo also identifies that MoEmodels exhibit distinctive characteristics across weights due to their hybrid dense-sparse architectures, and employs adaptive rank selection policies along with iterative optimizations to close the accuracy gap. MiLo does not rely on calibration data, allowing it to generalize to different MoE models and datasets without overfitting to a calibration set. To avoid the hardware inefficiencies of extreme quantization, such as 3-bit, MiLo develops Tensor Core-friendly 3-bit kernels, enabling measured latency speedups on 3-bit quantized MoE models. Our evaluation shows that MiLo outperforms existing methods on SoTA MoE models across various tasks.

MiLo: Efficient Quantized MoE Inference with Mixture of Low-Rank Compensators

TL;DR

MiLo tackles the heavy memory and accuracy penalties of extreme quantization in MoE models by introducing a calibration-free Quantize-Then-Compensate pipeline that augments INT3 MoEs with a mixture of adaptive low-rank compensators. The method jointly optimizes quantized weights and residuals via an iterative framework, using sub-problems for quantization error and low-rank approximation, and employs adaptive rank policies informed by MoE structure and data distribution. It further delivers hardware-friendly INT3 kernels with zero-bit-waste packing and efficient dequantization to realize real latency gains. Across state-of-the-art MoEs like Mixtral-8×7B and DeepSeek-MoE, MiLo achieves substantial accuracy recovery (often >87% of lost accuracy) with modest memory overhead and up to 3× speedups compared with baselines. The work demonstrates practical viability of highly compressed MoEs on real hardware, enabling scalable deployment without calibration data and with Tensor Core-friendly kernels.

Abstract

A critical approach for efficiently deploying Mixture-of-Experts (MoE) models with massive parameters is quantization. However, state-of-the-art MoE models suffer from non-negligible accuracy loss with extreme quantization, such as under 4 bits. To address this, we introduce MiLo, a novel method that augments highly quantized MoEs with a mixture of low-rank compensators. These compensators consume only a small amount of additional memory but significantly recover accuracy loss from extreme quantization. MiLo also identifies that MoEmodels exhibit distinctive characteristics across weights due to their hybrid dense-sparse architectures, and employs adaptive rank selection policies along with iterative optimizations to close the accuracy gap. MiLo does not rely on calibration data, allowing it to generalize to different MoE models and datasets without overfitting to a calibration set. To avoid the hardware inefficiencies of extreme quantization, such as 3-bit, MiLo develops Tensor Core-friendly 3-bit kernels, enabling measured latency speedups on 3-bit quantized MoE models. Our evaluation shows that MiLo outperforms existing methods on SoTA MoE models across various tasks.

Paper Structure

This paper contains 51 sections, 11 equations, 11 figures, 8 tables, 1 algorithm.

Figures (11)

  • Figure 1: Overview of our MiLo approach. At a high level, MiLo employs a Quantize-Then-Compensate approach, which augments low-bit quantized MoEs with a mixture of low-rank compensators, whose ranks are adaptively decided based on the distinctive characteristics of MoE weights. To minimize the accuracy loss, MiLo introduces an iterative optimization algorithm that jointly optimizes quantized MoEs and the mixture of low-rank compensators. MiLo includes a set of hardware-friendly INT3 kernels to achieve high measured speedups.
  • Figure 2: Mixtral-8x7B's (a) weight sampling from attention projection and (b) weight sampling from expert.
  • Figure 3: Heatmap of expert activation frequency in Mixtral-8$\times$7B and DeepSeek-MoE on the WikiText-2 task. The vertical axis from top to bottom represents the layer depth, and the horizontal axis represents expert indices.
  • Figure 4: The overlapping region of quantized and half-precision distribution in each figure is shown in green. (a) Information loss analysis for attention layer Mixtral-8$\times$7B. Left: INT3 weight quantization captures the outliers adequately but has large information loss at relatively insignificant weight values. Middle: INT4 is able to close some of the information gap but not completely. Right: INT3 together with low-rank matrices manage to close the information loss gap. (b) Information loss for expert layer at same $|w_i|$ range.
  • Figure 5: The correlation between relative Frobenius norm vs. Kurtosis. Each dot represents a weight matrix in layer 1 of DeepSeek-MoE.
  • ...and 6 more figures