Open3
LLMのアーキテクチャ、事前学習周りの論文メモ
atfujitaSwitch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
問題意識
- MoEは翻訳タスクで注目すべき成功をおさめたが複雑さや通信コスト、学習の不安定さなどで広い適用が妨げられている。本論文では従来よりも簡素化されたSwitch Transformerアーキテクチャを提案。
本論文のアプローチ
- 各トークンをルーティングして割り当てるExpertの数をこれまでの先行研究の2以上から1（Single Expert）へ変更（Switch layer）。ルーティングの計算コストを削減し、各Expertのバッチサイズを少なくとも半分にできる（Expert Capacityが上がる。Capacityを超えるとオーバーフローであふれたトークンは計算がスキップされるため重要）。
- Expertの計算の精度にbfloat16を使うことで通信コストを削減。ただ、bfloat16のような低精度でのルーター内のsoftmax演算は学習の不安定さを引き起こすため、ルーター内はfloat32で計算。これによって学習が安定。
- 適切な初期化が学習の成功には重要。デフォルトのTransformerのscale hyper-parameterを10分の1にするとSwitch Transformerの学習が安定した。これはモデルサイズを変えても当てはまった。
- MoEは過学習しやすくFine-tuningが難しいという課題があるが、シンプルにドロップアウトで対応。ただ、単純に全てのレイヤのドロップアウトレートを大きくしてもパフォーマンスが下がったため、Expertレイヤ以外は小さなレート、Expertレイヤには大きなレートを適用することでパフォーマンスが向上した。
評価とAblation Study
- 固定ステップ数での評価では、Expert数を増やすほどTest lossは小さくなり、サンプル効率が良くなった（少ないステップ数で収束する）。T5-baseと64 expertのSwitch Transformer-baseでは後者の方が同じ精度を達成するのにステップ数単位で7.5倍、時間で7倍高速だった。T5-largeと比較しても2.5倍高速だった。
- ダウンストリームタスクでもSwitch Transformer-base, largeはT5のbase, largeにほとんどのタスクで精度が上回る結果となった。
- 蒸留も実験。Non-Expertレイヤの重みで初期化したStudent-modelとTeacherのSwitch Transformer-base + ダウンストリームタスクの損失で蒸留したらT5-baseよりも精度が良かった（純粋な初期化からスタートするとT5-baseとほとんど変わらない結果だった）。
- 多言語モデルのmT5をSwitch Transformerにして比較。101の言語全てでmT5と比較して改善が見られた。
- Switch Transformer-XXL（395B） vs T5-XXL（11B）の比較では、前者の250kステップでの精度が後者の500kステップの精度を上回っていた。
- モデルのスケーリングはExpertパラレル（Expertごとにデバイスを分ける）だけで1.6Tパラメータ（Expert数は2048）まで大きくできた。大規模なスパースモデルは学習が不安定になることがあるが、1.6Tでは学習の不安定さはなかった。一方で、Switch Transformer-XXLでは不安定になることがあった。
atfujitaMoE-Mamba: Efficient Selective State Space Models with Mixture of Experts
問題意識
- SSM（State Space Model）ベースのアーキテクチャをスケーリングするためにMoEレイヤを導入。
本論文のアプローチ
- Mambaレイヤを一つおきにMoEレイヤ（Single ExpertにルーティングするSwitchレイヤ）へ置き換えた。
評価とAblation Study
- 素のTransformer, Transformer-MoE, Mamba, Mamba-MLP（N Expert=1）, MoE-Mambaで比較実験。English C4データセット（6.5Bトークン）で100kステップ学習（ハイパラは素のMamba向けに使われたものを全てのモデルに適用）。MoE-Mambaが最もTest lossが下がり、素のMambaよりも同じ精度を達成するために必要なステップ数が2.2倍短縮され、TransformerやTransformer-MoEよりも精度が良かった。一方で、Mamba-MLPは素のMambaに比べわずかに精度が劣る結果になった。
- Expert数を1, 4, 8, 16, 32へ変更して実験するとExpert数が大きいほど精度が良くなった。
- Mambaレイヤ内のOutput projectionやConv ProjectionをMoEに置き換えたが素のMambaと精度は変わらなかった。MambaとMoEを並列にしたが、素のMambaよりも精度が悪化した。
- Future WorkとLimitationとして
- この論文ではパラメータ数が1Bに満たない小さなモデルで実験されていること
- Mamba vs Mamba-MLPではMambaの方が優れていたため、MambaレイヤをMoE化してみるのはおもしろそうとのこと
- MoEにSwitchレイヤを適用したが、Mambaに対して最良のデザインなのかはわからないので色々と試す必要がある。
atfujitaMixtral of Experts
問題意識
- Mistral 7BをSparse Mixture of Experts (SMoE) でスケーリングさせたMixtral 8×7Bを提案。
本論文のアプローチ
- MistralのFFNを8 ExpertsのMoEレイヤへ置き換え。 ルーターは各トークンを2つのExpertへルーティングし、2つのExpertの出力は加重合計される。モデルのパラメータ数は47Bだが、推論時にアクティブなパラメータは13Bとなり、パラメータ数の割に推論は高速になる（一方でGPUメモリ上には47Bのパラメータがロードされることには注意が必要）。
- MoEレイヤはMegablocksのカーネルで大規模な疎行列の掛け算へとキャストされるのでGPU上で効率的に計算できる。また、モデルパラレルとエキスパートパラレルで複数のGPU上で分散学習できる。
- Mistral 7Bと同じくコンテキスト長は32kだが、事前学習に使用する多言語データの割合を増やした。
- InstructバージョンはSFTとDPOを実行した。
評価とAblation Study
- Llama2 7B/13B/70Bとベンチマークで比較したところ、ほとんどのタスクでLlama2 70Bを上回っており、特に数学やコード生成では優れた結果となった。
- 多言語タスクでの比較でも、Llama1 33BとLlama2 70Bを上回る結果となった。
- Llama2 70BとGPT3.5との比較でも7つのタスクの内、4つで最も優れた結果となった。
- 長いコンテキストでの性能をpasskey検索タスクで評価したが、シーケンスの長さや場所に限らず100%の精度で検索できた。
- モデルのバイアスの評価でもLlama2 70Bよりもバイアスが小さい結果となった。
- Instructバージョンは、LMSysリーダーボードで、GPT3.5-TurboやGemini Pro, Claude-2.1といった商用モデルよりも優れたパフォーマンスであった。