容量感知推論：混合エクスパートにおけるストラゲラー効果の緩和

arXiv:2503.05066v4 発表タイプ：置換 要約: 混合エクスパート（MoE）は、スパースなエクスパート活性化を活用することで、性能と効率のバランスを保ちながら大規模言語モデルのスケールアップに有効なアーキテクチャである。しかし、エクスパート並列化の下では、トークンからエクスパートへの割り当ての不均衡により推論非効率性が生じ、負荷の少ないエクスパートが早期に計算を完了しすぎた場合、過負荷のエクスパートを待たされることでグローバルな遅延が生じる。我々はこれを「ストラゲラー効果」と呼び、最も負担の大きいエクスパートが全体の推論遅延を決定すると定義する。これを解決するために、まず過剰な負荷のエクスパートから余分なトークンを棄却することによってエクスパート容量の制限を強制する「容量感知トークンドロップ」を提案し、性能への最小限の影響（OLMoE で 30% のスピードアップと 0.9% の劣化のみ）で負荷バランスを改善した。次に、残存する負荷の低いエクスパートが容量閾値の下に十分に存在する状況において、「容量感知拡張ドロップ」を導入し、トークンが厳密なローカル容量制約を適用する前に追加のローカルエクスパートを含めることを許可することで、負荷バランスを改善し、未使用のエクスパートの利用効率を向上させた。言語モデルとマルチモーダル MoE モデルの両方で実施された大規模な実験により、我々のアプローチの有効性が示され、エクスパート利用率、モデル性能、推論効率の劇的な向上をもたらした。例えば、Expanded Drop を Mixtral-8×7B-Instruct に適用すると、平均性能向上が 0.2%、推論スピードアップが 1.85 倍となった。コードは以下の URL から公開されている：https://github.com/CASE-Lab-UMD/Capacity-Aware-MoE。

arXiv:2503.05066v4 Announce Type: replace Abstract: The Mixture of Experts (MoE) is an effective architecture for scaling large language models by leveraging sparse expert activation to balance performance and efficiency. However, under expert parallelism, MoE suffers from inference inefficiencies due to imbalanced token-to-expert assignment, where underloaded experts complete computations early but must wait for overloaded experts, leading to global delays. We define this phenomenon as the \textbf{\textit{Straggler Effect}}, as the most burdened experts dictate the overall inference latency. To address this, we first propose \textit{\textbf{Capacity-Aware Token Drop}}, which enforces expert capacity limits by discarding excess tokens from overloaded experts, effectively reducing load imbalance with minimal performance impact (e.g., $30\%$ speedup with only $0.9\%$ degradation on OLMoE). Next, given the presence of low-load experts remaining well below the capacity threshold, we introduce \textit{\textbf{Capacity-Aware Expanded Drop}}, which allows tokens to include additional local experts in their candidate set before enforcing strict local capacity constraints, thereby improving load balance and enhancing the utilization of underused experts. Extensive experiments on both language and multimodal MoE models demonstrate the effectiveness of our approach, yielding substantial gains in expert utilization, model performance, and inference efficiency, e.g., applying Expanded Drop to Mixtral-8$\times$7B-Instruct yields a {0.2\%} average performance improvement and a {1.85$\times$} inference speedup. The code is released at: https://github.com/CASE-Lab-UMD/Capacity-Aware-MoE.