Distribution-Consistent Inference for Dynamic Sparse Mixture-of-Experts

2026-09-08

AI Summary

本文针对MoE模型在减少专家数量时产生的分布偏移问题，提出了一种名为LDA的轻量级推理时间校正方法来对齐表示，从而恢复性能。

Abstract

Mixture-of-Experts (MoE) architectures have emerged as a powerful paradigm for scaling model capacity while preserving efficient inference in large foundation models. However, most MoE models use a fixed top-$k$ expert selection policy, assigning the same expert budget to every token even when fewer experts may be sufficient. Inference-time dynamic top-$k$ routing can reduce computation without retraining, but existing methods often overlook the distributional shift caused by deviating from the training-time routing configuration. We show that reducing the number of activated experts consistently increases the RMS scale and variance of SMoE outputs, inducing a representation mismatch that contributes to downstream performance degradation in addition to the loss of expert capacity. To address this correctable component, we propose Layer-wise Distribution Alignment (LDA), a lightweight inference-time correction that uses layer-wise calibration statistics to align reduced-routing representations with the default configuration. Across multiple SMoE LLMs, benchmarks, and routing strategies, LDA recovers much of the performance lost induced by the distributional shift under reduced routing while preserving sparse-inference efficiency with negligible overhead.

Problem

• Mixture-of-Experts
• dynamic routing
• distribution shift
• performance degradation

Innovation

• Layer-wise Distribution Alignment
• sparse inference
• dynamic top-k routing
• representation mismatch