GRIN: GRadient-INformed MoE

要約

Mixture-of-Experts (MoE) モデルは、エキスパート ルーティングによるスパース計算により、高密度モデルよりも効果的に拡張し、エキスパート モジュールの小さなサブセットのみを選択的にアクティブにします。
ただし、離散エキスパート ルーティングは標準的な逆伝播を妨げ、したがって深層学習の基礎である勾配ベースの最適化を妨げるため、スパース計算は従来のトレーニング実践に課題をもたらします。
MoE のスケーリング能力をさらに追求するために、エキスパート ルーティングのためのスパース勾配推定を組み込み、トークン ドロップを回避するためにモデルの並列処理を構成する GRIN (GRAdient-INformed MoE training) を導入します。
GRIN を自己回帰言語モデリングに適用して、トップ 2 の 16$\times$3.8B の MoE モデルを開発します。
有効化されたパラメーターがわずか 6.6B の私たちのモデルは、7B の密なモデルを上回り、同じデータでトレーニングされた 14B の密なモデルのパフォーマンスに匹敵します。
多様なタスクにわたる広範な評価により、MOE の有効性を大幅に向上させる GRIN の可能性が実証され、MMLU で 79.4、HellaSwag で 83.7、HumanEval で 74.4、MATH で 58.9 を達成しました。

要約(オリジナル)

Mixture-of-Experts (MoE) models scale more effectively than dense models due to sparse computation through expert routing, selectively activating only a small subset of expert modules. However, sparse computation challenges traditional training practices, as discrete expert routing hinders standard backpropagation and thus gradient-based optimization, which are the cornerstone of deep learning. To better pursue the scaling power of MoE, we introduce GRIN (GRadient-INformed MoE training), which incorporates sparse gradient estimation for expert routing and configures model parallelism to avoid token dropping. Applying GRIN to autoregressive language modeling, we develop a top-2 16$\times$3.8B MoE model. Our model, with only 6.6B activated parameters, outperforms a 7B dense model and matches the performance of a 14B dense model trained on the same data. Extensive evaluations across diverse tasks demonstrate the potential of GRIN to significantly enhance MoE efficacy, achieving 79.4 on MMLU, 83.7 on HellaSwag, 74.4 on HumanEval, and 58.9 on MATH.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google