TaskMoE：効率的に推論を行うためにタスク別のルーティングを学習(1/2)

１．TaskMoE：効率的に推論を行うためにタスク別のルーティングを学習(1/2)まとめ

・MoEは効果的だが多くのトークンを複数のexpertsに送ると多くの計算コストがかかる
・TokenMoEはTaskMoEと異なりタスク毎にサブネットワークを割り当てる手法
・これによりマルチタスクモデル学習後、タスク専用のサブネットワークを抽出して利用可能

２．TaskMoEとは？

大規模言語モデルの規模拡大は、自然言語理解(T5)、生成(GPT-3)、多言語ニューラル機械翻訳(M4)において大幅な品質向上をもたらしています。

より大きなモデルを構築するためのよくある手法の1つは、深さ(レイヤー数)と幅(レイヤーの次元数)を増やして、ネットワークの次元を単純に拡大することです。

これらの密モデル(dense model)は、トークンと呼ばれる小さな部品に分割した入力を受け取り、すべてのトークンをネットワーク全体に与え、すべての層とパラメータを活性化させます。このような大規模な密モデルは、複数の自然言語処理(NLP:Natural Language Processing)タスクで最先端の成果を上げていますが、その学習にかかるコストはモデルサイズに比例して増加します。

代替案には、MoE(Mixture of Expert)をベースにした、ネットワークの一部のみを活性化させる疎なモデル(例：GShard-M4やGLaM)を構築する手法があり、最近人気が高まっています。

この手法では、ネットワークに渡される各トークンは、後続のサブネットワークに渡されます。全てのサブネットワークに渡されるわけではないので、モデルパラメータのいくつかをスキップさせる事ができます。

入力トークンを各サブネットワーク(expertsと呼ばれます)にどのように分配するかの選択は、ネットワークの残りの部分と一緒に学習される小さなルーターネットワーク(router networks)によって決定されます。これにより、研究者は、学習コストに比例して増加させることなく、モデルサイズ(つまりは性能)を増加させることができます。

これは訓練時には効果的な戦略ですが、多くのトークンを複数のexpertsに送ると、expertsを多数のアクセラレータに分散させなければならないため、やはり推論に計算コストがかかります。例えば、1.2兆パラメータのGLaMモデルを扱うには、256個のTPU-v3チップが必要です。密なモデルと同様に、MoE モデルの処理に必要なプロセッサーの数も、モデルサイズに対して線形に規模が拡大するため、計算要件が増加する一方で、通信オーバーヘッドが大きくなり、エンジニアリング作業の複雑さが増します。

論文「Beyond Distillation: Task-level Mixture-of-Experts for Efficient Inference」では、Task-level Mixture-of-Experts(TaskMoE)と呼ばれる、モデルの規模を拡大する事による品質向上メリットを利用しつつ、効率的にサービスを提供する方法を紹介します。

この手法は、大規模なマルチタスクモデルを学習し、そこから推論に適した小規模で独立したタスクごとのサブネットワークを抽出することで、モデルの品質を損なうことなく、推論の応答時間を大幅に削減するものです。

私たちは、この方法が多言語のニューラル機械翻訳(NMT:Neural Machine Translation)に有効であることを、他のMoEモデルや、知識蒸留手法を用いて圧縮したモデルと比較して実証します。

タスク情報を用いた大規模な疎活性化モデルの学習

私たちはルーターネットワークを、各タスクに固有の入力トークンを、タスクに関連する同じサブネットワークに送るよう学習します。例えば、多言語NMTの場合、ある言語に関する全てのトークンは同一のサブネットワークにルーティングされます。

これは、ルーターネットワークが入力中の異なるトークンをタスクに依存させずに異なるサブネットワークに送ることを学習する最近のアプローチ、例えばTokenMoEのような、疎ゲートエキスパート混合モデルとは異なっています。

サブネットワークを抽出して蒸留作業を回避する推論

このようにTaskMoEとTokenMoEではモデルの学習方法が異なるため、推論へのアプローチも異なります。

TokenMoEは学習時、推論時ともに同じタスクのトークンを多くのexpertsに分配するという手法をとっているため、推論時の計算量が多くなってしまうのです。

TaskMoEでは、学習時と推論時に、より小さなサブネットワークに1つのタスクを割り当てます。推論実行時には、各タスクの未使用のexpertsを破棄してサブネットワークを抽出します。

TaskMoEとその派生型は、1つの大きなマルチタスクネットワークを訓練し、推論時に各タスクに対して個別のサブネットワークを使用することで、訓練後に追加の圧縮手法を使用することなくサイズの圧縮を実現できます。

以下では、TaskMoEネットワークを学習し、推論時にタスクごとのサブネットワークを抽出するプロセスを説明します。

学習時、タスクベースのMoEでは言語情報(翻訳元言語、翻訳先言語、またはその両方)に基づき、同じ言語のトークンが同じexpertsにルーティングされます。その後、推論時に各タスクのサブネットワークを抽出し、未使用のexpertsを破棄します。

３．TaskMoE：効率的に推論を行うためにタスク別のルーティングを学習(1/2)関連リンク

１）ai.googleblog.com
Learning to Route by Task for Efficient Inference

２）arxiv.org
Beyond Distillation: Task-level Mixture-of-Experts for Efficient Inference