7K Views
October 06, 23
スライド概要
DL輪読会資料
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
書誌情報 タイトル: 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
既存研究の課題 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合計が1になるように 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