MoEとは|47Bのモデルを13B分の計算で動かす

MoEとは|47Bのモデルを13B分の計算で動かす

モデルの紹介文に「8x7B」や「総パラメータ47B・活性13B」といった書き方を見かけるようになりました。

数字が2つ並んでいると、どちらがそのモデルの大きさなのか分からなくなります

この2つの数字を生んでいる仕組みがMoE、Mixture of Expertsです。

今回は、MoEが何を節約していて、代わりに何を節約できないのかを整理します。

MoEは「全部は使わない」ことで速くする

通常のモデルは、1トークンを処理するたびに全パラメータを通す作りです。

MoEは層の中に複数の小さなネットワーク(エキスパート)を並べ、そのうち一部だけを動かす仕組み

残りのエキスパートは、そのトークンについては計算に参加しません。

モデルの知識量は増やしたまま、1回あたりの計算量だけを据え置く、というのが狙いです。

置き換えられるのは主にフィードフォワード層で、注意機構などはそのまま共有されます。

ルーターがトークンごとに担当を選ぶ

どのエキスパートを使うかを決めるのが、層ごとに置かれたルーターです。

Mixtral 8x7Bを例にすると、各層に8つのエキスパートがあり、論文(arXiv:2401.04088)によればルーターはトークンごとに2つを選びます。

選ばれる2つは層ごとに変わり、次のトークンでも同じ組み合わせとは限りません。

文章の途中で担当が切り替わるので、ひとつのモデルの中に得意分野の違う小型モデルを何本も抱えているような形になります

「エキスパート」という名前から、日本語担当やコード担当のように役割が分かれていると想像しがちなところ。

実際にはルーターが学習で勝手に決めた分担なので、人間が読んで意味の分かる区切りになっている保証はありません

総パラメータ数と活性パラメータ数は別物

冒頭の2つの数字は、ここで決まります。

Mixtral 8x7Bは総パラメータが47B、推論時に実際に使われるのは13Bというのが論文の記述です。

つまり計算コストは13B級、知識量は47B級という配分。

同じ論文では、この構成でLlama 2 70BやGPT-3.5と同等以上のスコアを出したと報告されています。

ちなみに7Bのエキスパートが8つでも56Bにならないのは、エキスパート化されていない部分が共有されているからです。

名前の掛け算をそのまま信じると、サイズを2割ほど多く見積もることになります。

得をした分は、どこかで払っている

いいことばかりに見えて、実は扱いにくさもあるのです。

Hugging Face の解説記事Mixture of Experts Explainedは、要点のひとつにファインチューニングの難しさを挙げています。

もうひとつ、学習中に立ちはだかるのがルーターの偏りです。

同記事の説明によると、放っておくとルーターは一部の人気エキスパートばかりを選ぶ方向へ収束し、それが自己強化されてしまうそう。

選ばれ続けるエキスパートだけが賢くなり、残りは育たないまま、という状態です。

だから実装側では、わざとノイズを入れるなどして担当をばらけさせる工夫が入っています。

使う側から見ると、同じ総パラメータ数でも密なモデルとMoEでは得意な使い方が違うと考えたほうが実態に近いはず。

速くはなるが、メモリは減らない

ここがMoEでいちばん誤解しやすいところ。

計算に使うのは一部でも、どのエキスパートが選ばれるかは事前に分からないので、全エキスパートをメモリに載せておく必要があります

同じ記事も、要点として「全エキスパートがメモリに載るため高いVRAMを要求する」と記しているのです。

速くなるのは計算であって、必要なメモリではない

手元のGPUで動くかどうかを判断するときは、活性13Bではなく総47Bのほうを見ることになります。

だから量子化でメモリを削る話とは、目的がきれいに補い合う関係です。

量子化とは?AIモデルを軽くする仕組みを整理

ローカルで触るときに見るところ

モデル名に掛け算が入っていたら、まず総パラメータ数を確認する。

次に、活性パラメータ数から生成の速さを見積もる。

量子化後のファイルサイズが手元のVRAMに収まるかを最後に確かめれば、動くかどうかの判断はだいたい付きます。

この順番で見ると、同じ「8x7B」でも配布形式によって扱いやすさが違う理由が分かりやすくなるはず。

実際に手元で動かす手順は、こちらにまとめてあります。

WindowsでOllamaを動かす|ローカルLLM入門

最後に

MoEは、モデルを小さくする技術ではありませんでした。

全部を並べておいて、毎回そのうち一部だけを起こす仕組み。

だから計算は軽くなり、メモリは軽くならない、という非対称が生まれます。

紹介文に数字が2つ並んでいたら、速さの話とメモリの話が同居していると思って読むと迷いません。

以上です。