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

要約

Mixture of Experts (MoE) フレームワークは、高密度モデルよりも優れたパフォーマンスを備えているため、大規模言語モデルの一般的なアーキテクチャになっています。
しかし、大規模な体制で教育省をゼロから訓練するには、法外な費用がかかります。
既存の方法では、複数の高密度エキスパート モデルを個別に事前トレーニングし、それらを使用して MoE を初期化することでこれを軽減します。
これは、他のパラメータを結合しながら、専門家のフィードフォワード ネットワーク (FFN) を使用して MoE の専門家を初期化することによって行われます。
ただし、この方法では、高密度モデル パラメータの再利用が FFN 層のみに制限されるため、これらのモデルを MoE に「アップサイクル」するときの利点が制限されます。
私たちは、この欠点を解決するシンプルかつ効果的な方法である BAM (Branch-Attend-Mix) を提案します。
BAM は、FFN を使用して MoE レイヤーを初期化するだけでなく、専門家の注意パラメーターを Mixture of Attendant (MoA) レイヤーのソフト バリアントに初期化することで完全に活用することにより、特殊な高密度モデルを最大限に活用します。
注意パラメータをアップサイクルするための 2 つの方法を検討します。1) 最高のモデル パフォーマンスを実現するために、すべての注意パラメータを含む密なモデルから個別の注意エキスパートを初期化します。
2) 推論効率の向上を促進するために、すべての専門家間でキーと値のパラメーターを共有します。
効率をさらに向上させるために、並列アテンション トランスフォーマー アーキテクチャを MoE に採用し、アテンション エキスパートと FFN エキスパートを同時に計算できるようにします。
5 億 9,000 万から 20 億のパラメーターにわたるシード モデルでの実験では、BAM が同じ計算とデータの制約内で、複雑さと下流タスクのパフォーマンスの両方でベースラインを上回っていることが実証されました。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google