【DL輪読会】From Sparse to Soft Mixtures of Experts

DEEP LEARNING JP From Sparse to Soft Mixtures of Experts スパース MoE から完全微分可能な ソフト MoE へ [DL Papers] Takayuki Yamamoto（LY Corporation, Waseda Univ. Kawahara Lab. D1） http://deeplearning.jp/ 1

• http://deeplearning.jp/

書誌情報 タイトル： From Sparse to Soft Mixtures of Experts スパースMoEから、完全微分可能なソフトMoEへ https://arxiv.org/abs/2308.00951 2023/9/2 著者： 概要： 完全微分可能なMoE手法の提案 選定理由： • • GPT-4に Mixture of Experts (=MoE) が使用されているというリーク情報による興味 MoEは単一モデルの大規模化を超えたメリットの可能性がある 実装： https://github.com/google-research/vmoe ※出典記載の無い図表は本論文からの引用 2

• https://arxiv.org/abs/2308.00951
• https://github.com/google-research/vmoe

既存研究の課題  Transformer を大きくすると、性能は向上するものの計算コストが増加する  それを改善する為の既存手法スパース  Expert MoE には次の課題がある 決定手法の離散最適化問題  訓練の不安定性 3

当研究が解決した課題と成果概要  完全微分可能な Soft-MoE というアーキテクチャを提案  計算量あたりの性能を、既存手法に比べて改善  小規模モデルで、大規模モデル同等の性能  小規模モデルでも( MoE 化で) パラメータ数は多いが、推論速度は逆に早い 4

関連研究：Meta の MoE の論文  Scaling Expert Language Models with Unsupervised Domain Discovery, 2023  入力コンテキストに近いクラスタを担当する Experts を使ってアンサンブル ※当論文は、輪読論文で Reference されていないが、図が分かりやすいため輪読者が掲載 5

関連研究：Tokens Choice 方式の MoE  Scaling Vision with Sparse Mixture of Experts, 2021  各Expert に送られる Tokens を選択する方式 6

関連研究：Tokens Choice 方式の MoE  Scaling Vision with Sparse Mixture of Experts, 2021 FFN(MLP)を置き換えている 7

関連研究：Expert Choice 方式の MoE  Mixture-of-Experts with Expert Choice Routing, 2022  各 Token をどの Expert に送信するかを選択する方式 8

当論文の MoE 既存手法との違い  既存手法は選択した Token を 各 Expert の Slot にアサイン  提案手法は全 Token 毎に重みがあり、各 Expert は全ての Token × 重みを Slots で受け取る 既存手法 提案手法 9

当論文の MoE  最初のMoE層のSlotとの線形結合Weightを可視化 Soft MoE S/16 の 128の Experts から8つを選択し抜粋。Token 毎に重みが異なる事が確認できる 10

手法 全体図 11

手法 数式 次ページから 図と共に説明 12

手法 入力側 抜粋 𝑑：次元 𝑑：次元 𝑝：Slot数/Expert 𝑚 𝑛×𝑝 𝑛  行列の次元数を[行,列]で表す  入力Token：X=[m,d]  slot数=n×p  学習パラメータ：φ=[d,slot]  Xφ=[m,slot]  D=softmax(Xφ) [m,slot] (m方向にsoftmax)  iがm方向、jがslot方向  ෨ 𝑋＝Input側Slot [slot,m]  𝑋෨ =𝐷 𝑇 X = [slot,m]×[m,d]=[slot,d]  𝑋෨ は各slotの値 [slot,d] 13

手法 Expert 処理部 抜粋 𝑑：次元 𝑑：次元 𝑝：Slot数/Expert  𝑋෨ はInput側の各slotの値 [slot,d]  各Input側Slotの値は、対応するExpertに入力され 処理され、Output側Slot値 𝑌෨ [slot,d]となる  各Expertは、対応するSlotからd次元 の入力を受 けてd次元の出力を返す構造  当論文でのExpertの構造はMLPと記載があるが、 具体的な層数や中間層の次元数などの記載はない。 推定であるが中間層が1層で次元数はdと思われる。 𝑛×𝑝 𝑛 𝑖：各Expert番号 𝑖：各slot番号 𝑝：ExpertあたりSlot数 14

手法 出力側 抜粋 参考：前ページD 𝑑：次元 𝑚 𝑛×𝑝  Output側Slot値 𝑌෨ [slot,d]  Combine Weights：C=softmax(Xφ) [m.slot]  先程のDとの違いは、Σのjをn･pで、softmaxの方向が違う n×p方向なので、全Slots合計が１になるように  output Token：Y= C 𝑌෨ [m,slot]×[slot,d]⇒[m,d] 15

提案手法の特徴  完全微分可能  Token の取りこぼしや Expert のアンバランスがない  高速：Sparse  Sparse 手法(既存)に必要な低速な Sort や Top-k が不要 手法(既存)は Expert 数がスケールできず大抵の研究が数十個  提案手法は、Expert 数がスケール可能で数千まで可能 16

学習･評価方法   実験タスク：Vision-Transformer による画像認識  JFT-4Bタスク ( google内部データセットで4B以上の画像、29k以上のクラスラベル)  ImageNetタスク プレトレーニング    データセット：JFT-4Bで学習 validation  JFT-4B validation precision@1とImageNet 10-shot accuracyで実施  ImageNetはモデル重み凍結し、分類ヘッドを付け各クラス10枚の画像で学習し評価 評価  ImageNet-1kの train セット(130万枚)でファインチューニングし、val セットで評価 17

実験① パレートモデル実験 結果  提案手法は学習計算量に対して精度が高い 提案手法：青 パレートモデル 18

実験②   計算量固定で最適構成探索と改善理由解明 実験結果 同じ計算量だと、Expert 数を増やし、Slots / Expert を減らした方が提案手法は精度が高い 既存手法は、Expert 数を増やすと学習時間が大幅に増加するが、提案手法はほぼ横ばい Expert 数による性能 (総計算量固定) 学習時間 提案手法：青 * 最後の2層に MoE を持つ ViT-S/16 で実験 19

結論 (ディスカッション)  Sparse 手法(既存)は、計算量の課題に直面する可能性があり、普及が 遅れているのではないか  Sparse 手法(既存)は、元々巨大なモデルから解き放つ為に考案された ので、分散される傾向があり、通信コストがかかる  提案手法( Soft-MoE )にも上記は当てはまるが、最新の高密度モデルは 十分大きく分散も可能である為、この観点での課題は減少している 20

輪読者考察 Attention との類似性 Self-Attention数式 𝑑：次元 𝑑：次元 𝑛×𝑝 D は𝑠𝑜𝑓𝑡𝑚𝑎𝑥(  ただし、 𝑑𝑘の分母に相当するものはない  𝑋෨ = 𝐷 𝑇 𝑋 は、 Vの部分が X 即ち入力になっているの で類似  という事は次のようなAttentionと考えられる 𝑝：Slot数/Expert 𝑚 𝑄𝐾𝑇 ) と類似 𝑑𝑘  𝑛   QとV → X(入力)  K →φ(ルーティングテーブル） Expert が Transformer でいう FFN 部だとすると、 FFN部 が マルチヘッド的な Expert になっていて、 そこへのルーティング重みがAttention 構造で行わ れている。という事(?) 21

輪読者 感想  Transformerの層の重ね方やFFNのあり方に対する議論なのでは  ViTの後半の層に MoE という事は、前半が表現学習の Transformer、 後半が意味合い変換の Soft-MoE という層構成が最適という事か  画像の場合はこの構造でも良いが、系列が重要な LLM に応用する場 合は、各 Expert が系列を処理できる構造にする必要がありそう  LLM系タスクの場合は、ExpertがMLPではなく、ここも Transformer構造が必要なのでは？ 22

Thank you. 23