Prime Intellect Team, Sami Jaghouar, Justus Mattern, Jack Min Ong, Jannik Straube, Manveer Basra, Aaron Pazdera, Kushal Thaman, Matthew Di Ferrante, Felix Gabriel, Fares Obeid, Kemal Erdem, Michael Keiblinger, Johannes Hagemann
D. Sculley, Will Cukierski, Phil Culliton, Sohier Dane, Maggie Demkin, Ryan Holbrook, Addison Howard, Paul Mooney, Walter Reade, Megan Risdal, Nate Keating
スパースなMixture of Experts(MoE)アーキテクチャは、Transformerモデルのスケーリングにおいて有望なアプローチとして登場した。初期の研究では主にフィードフォワードネットワーク(FFN)層にMoEを組み込んでいたが、最近の研究ではモデルの性能向上を目指し、MoEパラダイムをアテンション層に拡張することが検討されている。しかし、既存のアテンションベースのMoE層は専用の実装を必要とし、FFNベースのMoE層と比較して最適な性能を発揮していない。本論文では、アテンションメカニズムの新たな再定式化を導入し、アテンションモジュール内に潜在するFFNに類似した構造を明らかにすることで、アテンション層とFFN層におけるMoE設計を統一することを目指す。提案するアーキテクチャであるUMoEは、アテンションベースのMoE層を通じて優れた性能を達成しつつ、FFNとアテンションコンポーネント間の効率的なパラメータ共有を可能にする。