Lory Fully Differentiable Mixture-of-Experts for Autoregressive Language Model Pre-training - 20240912 (1)

DEEP LEARNING JP [DL Papers] Lory: Fully Differentiable Mixture-of-Experts for Autoregressive Language Model Pre-training Kim Yongmin, The University of Tokyo http://deeplearning.jp/ 1

• http://deeplearning.jp/

書誌情報 書誌情報 ▪ タイトル： Lory: Fully Differentiable Mixture-of-Experts for Autoregressive Language Model Pre-training ▪ 著者：Zexuan Zhong†, Mengzhou Xia†, Danqi Chen†, Mike Lewis‡ ▪ 所属：†Princeton University, ‡Meta AI ▪ 出願情報：May 06, 2024, COLM2024 ▪ 概要：最初にAuto regressiveモデルのsoft mixture of expertを提案した研究 ▪ 選定理由: Mixture of expertの新動向であるsoft mixture of expertをLLMに導入したため ▪ 実装: 現時点ではなし（本文中にpseudo codeあり） 2

概要 ▪ 背景 目次 ▪ Sparse MoEは学習の不安定など様々な問題が存在 ▪ Soft MoEが提案されたが，auto regressiveモデルに適用時，計算量が莫大 ▪ 手法 ▪ Soft MoEを causal segment routingという手法で効率的にAuto regressiveモデルへ適用 ▪ 結果 ▪ 既存のSparse MoE手法 [Zhou+ 2022]より，効率的かつ精度がいい (特にout-of-distributionにおいて，精度がいい） ▪ Expertの数が増加するほど精度が向上（32個のExpertsまで実験） ▪ Sparse MoEのexpert unbalanceの問題を緩和 3

背景 目次 ▪ Spare MoE VS Soft MoE Sparse expert MoE(左)とsoft expert MoE(右)の例 [Puigcerver+ 2023] ▪ Sparse MoE（上位のexpertやtokenを選択） ▪ Router networkからtop-kのexpertを選択 [Expert Choice] [Zhou+ 2022] (もしくはtop-kのトークンをそれぞれのexpertに割り当て) [Token Choice] ▪ Soft MoE（全てのExpertとTokenを選択） ▪ Router networkの出力値から，各expertsの重みをMergeして，使用 [Muqeeth + 2023] ▪ トークンをweight averageして，各expertに割当 [Puigcerver + 2023] 4

背景 背景 目次 ▪ Spare MoEには様々な問題が存在 ▪ 学習不安定 ▪ Discrete Optimization Problem b/w experts ▪ Token dropping ▪ Token choiceにおいて，選択されないトークンが存在 ▪ Expert unbalance ▪ Expert choiceにおいて，選択されないexpertが存在 ▪ Expert数によって，計算コストが増加 ▪ Expert数の増加からTop-Kの(Expert, Token)を計算に負担が発生 ▪ Finetuningに対して、低性能 ▪ MoEは学習データにOverfittingされやすいため，fine-tuningする場合，性能が低下 5

背景 目次 ▪ 既存のSoft MoE [Muqeeth+ 2023, Puigcerver+ 2023]はAuto Regressiveモデルに適用不可 ▪ 両方ともEncoder-DecoderモデルからのClassificationタスクに適用 Soft token MoE [Puigcerver + 2023] Soft expert MoE [Muqeeth + 2023] ▪ 本研究では[Muqeeth + 2023]のSoft expert MoEをauto regressiveモデルに拡張 6

背景 ▪ Fully Differentiable MoE Architectures via Expert Merging [Muqeeth+ 2023] Sparse Expert MoE(左)とSoft Expert MoE(右)の例 [Muqeeth + 2023] ▪ Sparse MoE (Top-k experts) [Zhou+ 2022] 𝑜𝑥 = ∑𝐸𝑖=1 𝑒𝑖 ⋅ FFN ℎ𝑥 ; 𝜃𝑖 , where 𝑒𝑖 = Top − 𝑘 Softmax 𝑅 ℎ𝑥 ▪ Soft MoE [Muqeeth+ 2023] 𝑜𝑥 = FFN ℎ𝑥 ; ∑𝐸𝑖=1 𝑒𝑖 ⋅ 𝜃𝑖 , where 𝑒𝑖 = Softmax 𝑅 ℎ𝑥 𝑖 ▪ Auto Regressiveモデルでは各トークン入力ごと，モデル融合の計算が必要 → Expertの数が大きいと計算量が莫大なので，ほぼ不可能 7

手法 ▪ 手法 提案手法 1. Similarity based data batching 1. Similarity-based data batching 2. Casual segment routing 提案手法の全体図 ▪ バッチに似たようなデータが入るようにデータセットを構成 2. Causal segment routing ▪ Routingを個別トークンではなくて，input sequenceの複数トークンであるsegmentでrouting 8

手法 目次 Data Batching 1. Similarity-based ▪ Pre-train用の学習データは一般に，ランダムに連結したデータ ▪ 隣接したsegmentが互いに関係がない可能性が高い → Specializationが必要なrouterやexpertの学習が不安定 ▪ バッチに類似した文が入るようにバッチング Similarity-based data batchingの概要図 9

手法 2. Causal segment routing ▪ L個のinput sequenceからN個のsegment𝑆𝑘 を作成（segment size N = ⌈𝐿/𝑇⌉) ▪ Routerの入力はトークンではなくて，segment 𝑆を入力 ▪ 𝑘番目MoEを計算する例 3) 𝑘番目MoE (Merged FFN) 2) 𝑖番目のExpertのweight 𝑒𝑖 の計算 𝑅 ∗ : Router 𝜃᪄𝑘 = ∑𝑖 𝑒𝑖 ⋅ 𝜃𝑖 𝑒𝑖 = Softmax 𝑅 ℎ᪄ 𝑘−1 ℎ᪄ 𝑘−1 = 𝜃𝑖 : 𝑖番目のExpert 1 ∑ ℎ 𝑇 𝑥∈𝑆𝑘−1 𝑥 𝑘番目の Segmentation 1) Segment 𝑆𝑘−1にある 潜在表現の平均 提案手法における，𝑘番目のMoEを計算する例 10

手法 ▪ Causal segment routingの問題 ▪ Segment 𝑆1 に関してはℎ1 を使用して，MoE (Merged FFN)を計算 ▪ しかし，情報リークが発生 → 𝑅 ℎത1 に対しては学習しない 11

手法 ▪ 推論時のprompt-only routing ▪ 推論時はsegmentationの作成が困難 → 単一トークンでMoEする必要あり ▪ しかし，計算コストが莫大 → 推論時はinput prompt，全体の潜在表現を平均したものをrouterの入力として使用 ▪ その後，routerの出力𝑒𝑖 は固定（出力においてはMoE固定） 入力Promptから，Expertの組み合わせを決めて、固定 12

実験設定 目次 Like Model (active parameter1 size: 0.3B, 1.5B) ▪ モデル: LLama2 ▪ 学習設定 ▪ Context window: 4096 token ▪ Segment length: 𝑇 = 256 ▪ Dataset: Commoncrawl dataset ▪ The number of total training token: 150B ▪ Similarity-based data batchingを適用 1各MoE Layerを融合した後のパレメータ数 13

結果 ▪ Expert数による学習損失と評価データでのPerplexity (Fluency) 0. 3B 1.5B 学 習 損 失 訓練トークン数 Expert数ごとの学習損失 (左)と評価データセットでのPerplexity (右) ▪ 左: Expert数が多いほど，学習損失が下がる ▪ 右: Expert数が多いほど評価データセットでのPerplexityが下がる 14

結果 ▪ Expert数における，downstraemタスクでの能力評価 ▪ 0.3B/32EをDenseモデルと比較した場合の平均的な性能の上昇率 ▪ Commonsense reasoning: +3.7% / Reading comprehension: +3.3% ▪ Question and answering: +1.5% / Text classification: +11.1% 15

結果 ▪ Casual segment routing ablation study ▪ Dense, prefix routingとの比較 * Prefix routing: 最初のsegment𝑆1 からのroutingだけ使用 → 全体のsegmentから学習する提案手法と比較 学 習 損 失 訓練トークン数 Prefix routingと提案手法との学習損失の比較 16

結果 ▪ Similarity-based data batchingの評価 ▪ Similarity-based data batchingの有無からの実験 学 習 損 失 訓練トークン数 Similarity-basedとrandom data batchingとの比較 学習損失の評価 (左)，評価データセットでの学習損失の評価 (右) ▪ Denseモデルより，MoEモデルの方が学習損失が低い ▪ Similarity-based data batchingの方が学習損失が低い ▪ 学習トークン数が大きくなるほど，学習損失の改善も大きくなる 17

結果 ▪ 既存のMoE手法との比較 ▪ Expert choice (EC) [Zhou+ 2022]と比較 ▪ segment-level routingとtoken-level routingで比較 学 習 損 失 訓練トークン数 ▪ 同じrouting設定ではLoryの方が学習損失が低い ▪ Token-levelにすると若干ECの方が低いが，学習コストが高い 18

結果 ▪ 既存のMoE手法 (EC)との比較 ▪ Out-of-Distributionデータからの評価 (Perplexity) Out-of-Distributionデータセットにおける 既存のMoE手法 (EC)と提案手法 (Lory)の比較 ▪ 学習データ(Common crawl)と類似したC4での評価はECのPerplexityが低い ▪ 学習データともっと異なるPythonでの評価はLoryのPerplexityが低い → Loryの方がSparse MoEのより一般性能がいい 19

結果 目次 ▪ Expert utilizationとspecialization ▪ Expert utilization: Expertを多様に使っているか ▪ Expert specialization: データの特定の特性に関して，特化したexpertができたか 重 み の 平 均 値 0.3B / 8EでRouting networkの重みの平均（0, 11, 23層で測定した結果） ▪ 全てのExpertが使用されている ▪ Domain levelのexpert specializationが観測（例: 11層の7 expert) ▪ arXivとPython (コードベース) の重みがBooksとWikipediaより類似 ▪ 層が高くなるほどDomainに特化したexpertがもっと存在 20

結果 ▪ Sparse MoEでのexpert utilizationとspecialization比較 ▪ Sparse MoEではtoken-levelのspecializationが観測 ▪ 例) punctuations, articlesに特化したexpert Mistral of Expert 8x7Bのrouter networkの重みの平均 (0，15，31層で観測した結果) [Jiang+ 2024] ▪ Sparse MoEで観測されたdomain levelのspecializationはsoft MoEより強度が低い 21

まとめ・所感 ▪ まとめ 目次 ▪ Sparse MoEは様々な問題が存在するため，soft MoEが必要 ▪ 既存のsoft MoE [Muqeeth+ 2023]をauto regressiveモデルに適用すれば計算量が莫大 ▪ 最初にauto regressiveモデルにおけるsoft MoEを効率的に適用（across experts) ▪ Sparse MoEよりout-of-distributionでの性能がいい ▪ Expert unbalance問題が解消 ▪ 各expertにおけるspecialization現象が発生しやすい ▪ 所感 ▪ MoEのlevelを低くする(segmentationの長さを減らす）ことはコストが高い ▪ 推論時のMoEはメモリコストが高すぎるため、考案が必要 22

宣伝 ▪ LLMの無害化に関する論文をCOLM2024に投稿 [Kim+ 2024] ▪ Decoupling Noise and Toxic Parameters for Large Language Model Detoxification by Task Vector Merging 提案手法の全体図 ▪ LLMにおいて，既存の無害化手法は大きな性能低下が発生 → モデルマージ手法を用いて，既存の手法により，性能低下を最大38.9%抑制 1 ▪ 提案手法による無害化したモデルが既存の手法からのモデルより，元のモデルともっと類似であることを検証 1GPT2-smallにおいて，task vector negation[Ilharco+ 2023]の無害化手法との比較結果 23

参考文献 ▪ [Zhou+ 2022] Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoc V Le, James Laudon, et al. Mixture-of-experts with expert choice routing. In Advances in Neural Information Processing Systems (NeurIPS), 2022. ▪ [Muqeeth+ 2023] M. Muqeeth, H. Liu, and C. Raffel, “Soft Merging of Experts with Adaptive Routing.” arXiv, Jun. 06, 2023. Accessed: May 13, 2024. [Online]. Available: http://arxiv.org/abs/2306.03745 ▪ [Puigcerver+ 2023] J. Puigcerver, C. Riquelme, B. Mustafa, and N. Houlsby, “From Sparse to Soft Mixtures of Experts.” arXiv, Aug. 02, 2023. Accessed: Apr. 16, 2024. [Online]. Available: http://arxiv.org/abs/2308.00951 ▪ [Jiang+ 2024] A. Q. Jiang et al., “Mixtral of Experts.” arXiv, Jan. 08, 2024. Accessed: May 13, 2024. [Online]. Available: http://arxiv.org/abs/2401.04088 ▪ [Tang+ 2024] A. Tang, L. Shen, Y. Luo, N. Yin, L. Zhang, and D. Tao, “Merging Multi-Task Models via Weight-Ensembling Mixture of Experts.” arXiv, Feb. 01, 2024. Accessed: Apr. 12, 2024. [Online]. Available: http://arxiv.org/abs/2402.00433 ▪ [Kim+ 2024] Y. Kim, T. Kojima, Y. Iwasawa, Y. Matsuo. Decoupling noise and toxic parameters for language model detoxification by task vector merging. in First conference on language modeling (2024). Available: https://openreview.net/forum?id=TBNYjdOazs#discussion ▪ [Ilharco+ 2023] G. Ilharco, M. T. Ribeiro, M. Wortsman, L. Schmidt, H. Hajishirzi, and A. Farhadi, “Editing models with task arithmetic,” in The eleventh international conference on learning representations, 2023. 24

• http://arxiv.org/abs/2306.03745
• http://arxiv.org/abs/2308.00951
• http://arxiv.org/abs/2401.04088
• http://arxiv.org/abs/2402.00433
• https://openreview.net/forum?id=TBNYjdOazs