Robust Heterogeneous Analog-Digital Computing for Mixture-of-Experts Models with Theoretical Generalization Guarantees

TL;DR

This paper introduces a noise-robust, retraining-free heterogeneous analog-digital computing framework for large Mixture-of-Experts models, combining digital and analog hardware to improve efficiency while maintaining accuracy.

Contribution

It proposes a novel heterogeneous computation method that identifies and computes noise-sensitive experts digitally, reducing retraining needs and hardware nonideality impacts.

Findings

Maintains model accuracy under analog hardware nonidealities.

Effective for large language models like DeepSeekMoE and OLMoE.

Reduces energy and memory usage during inference.

Abstract

Sparse Mixture-of-Experts (MoE) models enable efficient scalability by activating only a small sub-set of experts per input, yet their massive parameter counts lead to substantial memory and energy inefficiency during inference. Analog in-memory computing (AIMC) offers a promising solution by eliminating frequent data movement between memory and compute units. However, mitigating hardware nonidealities of AIMC typically requires noise-aware retraining, which is infeasible for large MoE models. In this paper, we propose a retraining-free heterogeneous computation framework in which noise-sensitive experts, which are provably identifiable by their maximum neuron norm, are computed digitally while the majority of the experts are executed on AIMC hardware. We further assign densely activated modules, such as attention layers, to digital computation due to their high noise sensitivity despite comprising a small fraction of parameters. Extensive experiments on large MoE language models, including DeepSeekMoE and OLMoE, across multiple benchmark tasks validate the robustness of our approach in maintaining accuracy under analog nonidealities.