Mixture-of-Expertsにおける最適なエキスパート・アテンション割り当て: 動的モデル設計のスケーラブルな法則

要旨: 本論文は、Mixture-of-Experts（MoE）モデルへのニューラルスケーリング則の新たな拡張を提示し、エキスパート層とアテンション層のサブレイヤ間の計算資源の最適割り当てに焦点を当てています。MoEアーキテクチャは、計算量を比例的に増やすことなくモデルの容量を拡張する効率的な手法として浮上してきたため、最適なエキスパート・アテンションの計算比率を決定することが重要になります。比率 r を、トークンあたりの総 FLOPs のうち、エキスパート層に割り当てられる部分の割合として定義し、この比率が全体の計算予算およびモデルのスパース性とどう相互作用するかを調べます。GPT風のMoEトランスフォーマーを用いた大規模な実験を通じて、最適比率 r^* は総計算量とべき乗則に従い、スパーシティによって変化することを経験的に見出しました。私たちの分析は r^* の明示的な式を導出し、エキスパート・アテンション間の計算割り当てを正確に制御できるようにします。このアーキテクチャパラメータを取り入れることで、Chinchillaのスケーリング則を一般化し、サイズやデータを超えたMoEモデルの調整の新しい枠組みを提供します。私たちの知見は、固定された計算予算を遵守しつつ性能を最大化する、効率的なMoEモデルの設計に実践的なガイドラインを提供します。