Research / 基盤技術

Switch Transformerを読む 専門家を一つ選ぶ疎なモデル

多数の専門家層から入力ごとに一つを選び、計算量を抑えながらモデルを大きくできるか。Switch Transformer論文の設計、事前学習の速度、運用上の制約を読む。

論文初公開(arXiv)
参照版
v3
TRAIL公開
  • Switch Transformer論文は、多数の専門家層から各トークンに一つだけ選ぶ設計を調べた。
  • T5を基にした事前学習で、同じ計算資源を使う比較において最大7倍の速度向上を報告した。
  • 速度は論文の学習条件での値であり、推論の費用や業務の品質を直接示さない。

何を確かめた研究か

大きな言語モデルではパラメータを増やすと、通常は使う計算量も増える。「Switch Transformers」は、混合専門家モデル(Mixture of Experts、MoE)でこの関係を変えられるか調べた。MoEは複数の「専門家」と呼ぶニューラルネットワークを用意し、入力に応じて一部だけ使う方式だ。ここでの専門家は人間の専門家でも、特定分野の知識が保証された部品でもない。

従来のMoEは振り分けの複雑さ、機器間の通信、学習の不安定さが課題だった。著者らは各トークンを一つの専門家だけに送るSwitch層を提案する。使わない専門家の演算を省く一方、全専門家のパラメータを保存し、選択した専門家へデータを運ぶ必要がある。

どう検証したか

論文はT5-BaseとT5-Largeを土台に、通常の密なモデルとSwitch層を含むモデルを比較した。主な測定は事前学習の損失がどれだけ速く下がるかで、計算資源をそろえた条件を使う。多言語の評価ではmT5-Baseを比較対象とし、101言語での結果を調べた。より大きな構成では、約1.6兆のパラメータのモデルを事前学習し、T5-XXLとの速度を比べている。

専門家への振り分けは無制限ではない。一つの専門家が受け取るトークン数には容量があり、超過したトークンはその層の専門家処理を通らず次に進む。著者らは、負荷を均す補助損失、初期化や精度の扱いなども組み合わせた。低精度のbfloat16で大きな疎モデルを学習できるかも検証対象にした。

何が分かったか

論文の要約は、T5-BaseとT5-Largeに基づく設定で、同じ計算資源を使った事前学習速度が最大7倍になったと報告する。これは全設定の平均値ではなく、対象モデルと損失の到達点を含む実験条件での結果だ。多言語の設定でもmT5-Baseに対する改善を報告したが、言語ごとに同じ改善幅が出るという意味ではない。

約1.6兆のパラメータを持つ実験では、T5-XXLとの比較で4倍の速度向上を報告する。総パラメータ数を増やしながら、各トークンで活性化する部分を限ったことが設計の中心にある。ただし、総パラメータ数と一回の計算に使う量は別の尺度だ。パラメータ数だけで、応答品質や必要な機器を判断できない。

限界は何か

疎な計算でも専門家の重みを保持するメモリと、分散した機器間でトークンを交換する通信が要る。容量を超えたトークンの扱いも、入力や設定によって変わる。論文が示す事前学習の速度は、低負荷の短文応答、長文生成、実サービスの同時接続を測った値ではない。モデルの規模と学習速度だけで配備しやすさは決まらない。

仕事との接点

MoEモデルを候補にするときは、総パラメータ数だけでなく、実際に使う専門家数、機器間通信、メモリ、応答時間を確認したい。品質は自社の入力と言語で評価する。これはTRAIL編集部の確認項目であり、論文が特定の業務での費用削減を実証したものではない。

01

出典

Research の一覧へ戻る