BAM! これで完了：Mixture of Expertsのためのシンプルで効率的なパラメータ再利用

BAM! Just Like That: Simple and Efficient Parameter Upcycling for Mixture of Experts

August 15, 2024

著者: Qizhen Zhang, Nikolas Gritsch, Dwaraknath Gnaneshwar, Simon Guo, David Cairuz, Bharat Venkitesh, Jakob Foerster, Phil Blunsom, Sebastian Ruder, Ahmet Ustun, Acyr Locatelli

cs.AI

要旨

Mixture of Experts（MoE）フレームワークは、密なモデルを上回る優れた性能から、大規模言語モデルのアーキテクチャとして広く採用されています。しかし、大規模な環境でMoEをゼロから訓練することは非常にコストがかかります。既存の手法では、複数の密なエキスパートモデルを独立して事前訓練し、それらをMoEの初期化に利用することでこの問題を緩和しています。具体的には、エキスパートのフィードフォワードネットワーク（FFN）をMoEのエキスパート初期化に使用し、他のパラメータを統合します。しかし、この方法では密なモデルのパラメータの再利用がFFN層に限定されるため、これらのモデルをMoEに「アップサイクル」する際の利点が制約されます。本論文では、この欠点を解決するシンプルかつ効果的な手法であるBAM（Branch-Attend-Mix）を提案します。BAMは、FFNをMoE層の初期化に使用するだけでなく、エキスパートのアテンションパラメータを完全に活用し、それらをMixture of Attention（MoA）層のソフトバリアントとして初期化することで、専門化された密なモデルを最大限に活用します。アテンションパラメータのアップサイクルには、2つの方法を検討します：1）最適なモデル性能を得るために、密なモデルからすべてのアテンションパラメータを含む個別のアテンションエキスパートを初期化する方法、および2）推論効率を向上させるために、すべてのエキスパート間でキーとバリューパラメータを共有する方法です。さらに効率を向上させるため、MoEに並列アテンショントランスフォーマーアーキテクチャを採用し、アテンションエキスパートとFFNエキスパートを並列に計算できるようにします。590Mから2Bパラメータのシードモデルを用いた実験では、BAMが同じ計算量とデータ制約の下で、パープレキシティと下流タスクの性能の両方においてベースラインを上回ることを示しています。

English

The Mixture of Experts (MoE) framework has become a popular architecture for large language models due to its superior performance over dense models. However, training MoEs from scratch in a large-scale regime is prohibitively expensive. Existing methods mitigate this by pre-training multiple dense expert models independently and using them to initialize an MoE. This is done by using experts' feed-forward network (FFN) to initialize the MoE's experts while merging other parameters. However, this method limits the reuse of dense model parameters to only the FFN layers, thereby constraining the advantages when "upcycling" these models into MoEs. We propose BAM (Branch-Attend-Mix), a simple yet effective method that addresses this shortcoming. BAM makes full use of specialized dense models by not only using their FFN to initialize the MoE layers but also leveraging experts' attention parameters fully by initializing them into a soft-variant of Mixture of Attention (MoA) layers. We explore two methods for upcycling attention parameters: 1) initializing separate attention experts from dense models including all attention parameters for the best model performance; and 2) sharing key and value parameters across all experts to facilitate for better inference efficiency. To further improve efficiency, we adopt a parallel attention transformer architecture to MoEs, which allows the attention experts and FFN experts to be computed concurrently. Our experiments on seed models ranging from 590 million to 2 billion parameters demonstrate that BAM surpasses baselines in both perplexity and downstream task performance, within the same computational and data constraints.