Moe model experts get better balance in efficient quantization

Colla-Q: Toward Collaborative Experts in MoE Quantization via Minimax Precision Balancing

Machine LearningComputation and Language

Summary

Reducing memory and speed needs in large AI models often involves cutting down the precision of calculations, but this can hurt how well the models work. The paper’s authors found that this problem is worse for AI models called Mixture-of-Experts, where small parts called experts work together. If one expert suffers from low precision, the whole model’s performance drops. They created a method called Colla-Q that smartly gives different precision levels to experts based on their importance, so all parts work well together even with lower precision. This helps the AI be more reliable and less sensitive to changes in test conditions.

What this means in practice

  • For machine learning engineers: Deploy MoE models with efficient quantization that balances expert precision to maintain better overall model accuracy and stability.
  • For hardware accelerator designers: Design quantization schemes that allocate different bit-widths across expert units to optimize performance in MoE architectures on limited precision hardware.

Authors

Eunju Shin, Jongbin Ryu

Abstract

In this paper, we present a Mixture-of-Experts (MoE) quantization method based on activation entropy. Although quantization reduces memory and computational costs, it can substantially degrade performance. In particular, performance decline is pronounced in quantized MoE models, where individual experts have a small number of parameters that are sensitive to low-bit representation. Considering that MoE operates as an ensemble model with collaborative contributions from routed experts, a significant performance decline of a particular expert due to quantization can harm model performance. Therefore, we propose Colla-Q, a bit-allocation framework to maintain balanced performance across experts through an activation-entropy-based bit-width allocation algorithm. This approach encourages each expert to operate collaboratively in the quantized model, thereby 1) improving the overall MoE performance and 2) reducing the dependence on the calibration dataset. Since uniformly adjusting each expert's performance facilitates robustness and stability of the MoE model, the proposed MoE quantization method can generalize more consistently across different calibration datasets. Our code is available at: https://github.com/mmai-laboratory/Colla_Q