モジュラー大型モデルが登場！ IBMがWatsonXコアアーキテクチャの技術的詳細を公開

大規模言語モデル (LLM) は強力なパフォーマンスを備えていますが、既存のモデルのトレーニングと展開にはコストがかかります。また、以前の知識を忘れずに新しい知識を習得するためにそれらを拡張することも困難です。特定のタスク用の軽量モデルを抽出することも困難です。

最近、MIT-IBM Watson AI Lab、清華大学、Milaの研究者らが共同で、モジュール性を利用して大規模言語モデルの効率性と柔軟性を大幅に向上させる新しいニューラルネットワークアーキテクチャModuleFormerを提案しました。

ModuleFormer は、 Sparse Mixture of Experts (SMoE) に基づくモジュール アーキテクチャであり、新しい stick-breaking attention エキスパート モジュールと従来の MLP エキスパート モジュールという 2 種類のモジュールが含まれています。トレーニングと推論中、入力ベクトルに応じてさまざまなモジュールがスパースにアクティブ化されます。

以前の SMoE ベースのモジュール言語モデルとは異なり、ModuleFormer は、新しい負荷バランスと負荷集中損失関数を通じて、フィルタリングされていないデータからモジュール性を誘導できます。

実験では、チームはモジュール式アーキテクチャにより、大規模な事前トレーニング済み言語モデルに次の 3 つの重要な機能を持たせることができることを発見しました。

1) 効率性

ModuleFormer は入力ごとにモジュールのごく一部だけをアクティブにするため、従来の言語モデルと同等のパフォーマンスを 2 倍以上のスループットで実現できます。

2) スケーラビリティ

実験では、ModuleFormer は微調整フェーズでモジュールの一部のみを更新する必要があるため、従来の言語モデルよりも壊滅的な忘却が発生しにくく、新しいモジュールを通じて簡単に拡張して、トレーニング データに含まれていない新しい知識を学習できることが示されています。

3) モジュールの専門化とスクリーニング

微調整段階では、新しく提案された負荷集中損失機能である ModuleFormer によって、一部のモジュールを自動的に選別して対象タスクに集中できるようにし、タスクに関係のないモジュールを直接破棄して軽量な展開を実現します。

導入

現代の大規模言語モデル (LLM) は、いくつかのタスクで人間のパフォーマンスを上回る驚くべき結果を達成していますが、その効率性と柔軟性は依然として低いままです。

ほとんどの LLM (Llama、Falcon など) は推論とトレーニング中にすべてのパラメータを使用するため、これらのモデルを密モデルと呼びます。

しかし、これまでの研究では、特定のタスクでは、ニューラル モデル内のほとんどのパラメーターをモデルのパフォーマンスに影響を与えることなく削除できることが示されています。

さらに、LLM は一度トレーニングされると「時間的に固定」されますが、実際の多くのアプリケーションでは、LLM が最新の知識を持っていることが求められます。したがって、モデルの学習を継続する能力も非常に重要です。

モデルのサイズが大きくなるにつれて、ドメイン適応や継続的な学習のためにモデル全体を微調整することはコストがかかり、計算リソースが制限されるため、計算予算が少ないユーザーにとっては実行不可能になります。同時に、すべてのパラメータを更新すると、モデルが壊滅的な忘却の影響を受けやすくなります。

このため、元のパラメータ修正のごく一部だけを更新する LoRA のような軽量な適応方法が普及しつつあります。

しかし、実験により、このアプローチでは依然として壊滅的な忘却が発生する可能性があることが示されており、LoRA は、モデルに新しい言語を学習させるなど、モデルに多くの新しい知識を学習させる必要があるシナリオには適していません。

この記事の著者は、モジュール化が前述の問題を解決する良い方法であると考えています。モジュラー モデルにはいくつかの利点があります。

• モデルは、入力またはタスクに応じて条件付きでモジュールのセットをアクティブ化できるため、モデル全体を密にアクティブ化するよりも計算量が少なくて済みます。
• ドメインまたはタスクが指定されると、ドメイン/タスク関連のモジュールのセットを組み立てて、新しい軽量モデルを形成できます。
• このモデルは、ドメイン適応や継続的な学習のための新しいモジュールを簡単に追加できます。
• モデルの微調整中に入力に関連するモジュールのみが更新されるため、モデルが壊滅的な忘却に陥る可能性が低くなる可能性があります。

この論文では、新しいモジュール アーキテクチャである ModuleFormer (図 1a) と、その中でモジュール操作を実行する方法を提案します。 ModuleFormer の各レイヤーには、フィードフォワード ニューラル ネットワーク (MLP) エキスパート モジュールのグループと、新しく提案された Stickbreaking 注意エキスパート モジュールのグループが含まれています。

トレーニング中にさまざまなモジュールの負荷を分散させるために、著者らは新しい相互情報損失関数を提案しました。さらに、この記事では、ModuleFormer に新しいモジュールを挿入する方法 (図 1b) と、モジュールのプルーニングを実行する方法 (図 1c) も示します。

モジュールのプルーニングを実現するために、この記事では、新しい負荷集中損失関数を紹介します。この関数は、微調整中に特定のタスクを実行するのに最適なモジュールを自動的に選択し、微調整を通じてこれらのモジュールのタスク実行能力をさらに強化するために使用されます。

実験結果によると、ModuleFormer は密なモデルに比べて以下の点で大幅な改善が見られます。

• スパース モジュール アクティベーション メカニズム (図 1a) により、レイテンシ (50%) が低く、メモリ フットプリントが小さい状態で、高密度 LLM と同じパフォーマンスを実現できるため、ModuleFormer は高密度モデルの 2 倍のスループットを実現できます。
• 新しいドメインで微調整された後の壊滅的な忘却に悩まされることが少なくなり、新しいモジュールを追加することで新しい言語や知識を簡単に学習することもできます。 （図1b）。
• 下流のタスクで微調整して、モジュールの一部をタスク処理専用のモジュールに特化することができ、未使用のモジュールはタスクでのモデルのパフォーマンスを犠牲にすることなく破棄できます (図 1c)。

ModuleFormerモデルの詳細説明

スパースミックスオブエキスパート（SMoE）の基本構造

SMoE は、Shazeer の論文「Outrageously large neural networks: The sparsely-gated mixing-of-experts layer」で初めて提案されました。

各 SMoE レイヤーは、入力を処理して出力を生成するエキスパート モジュール m_1、m_2、…、m_n のセットと、エキスパートを選択するためのルーティング関数 g で構成されます。この記事では、ルーティング機能を単層ニューラル ネットワークでモデル化します。

入力ベクトル x が与えられると、ルーティング関数 g はエキスパート モジュール g(m|x) の確率分布を計算し、モデルは分布内の上位 k 個のエキスパート モジュールを自動的に選択して入力 x を処理します。 SMoE レイヤーの出力はエキスパート モジュールの出力であり、対応するルーティング確率を合計することによって得られます。

topk に存在しないエキスパート モジュールは計算されず、対応する確率 g は 0 に置き換えられます。

ModuleFormer では、著者は 2 つの異なるエキスパート モジュールを使用して、従来のトランスフォーマー デコーダーの自己注意層とフィードフォワード層 (MLP) をそれぞれ構築します。フィードフォワード層のエキスパート モジュールは、単一の隠し層を持つ完全に接続されたニューラル ネットワークである一般的な SMoE エキスパート モジュールと一致しています。自己注意層のエキスパート モジュールは、新しく提案されたスティック ブレーク注意モジュールです。

スティックブレイクアテンションモジュール

Stick-Breaking 自己注意は、Transformer デコーダーの自己注意層を置き換えるように設計されており、標準の自己注意層のソフトマックス関数の代わりに、ディリクレ過程の Stick-Breaking プロセスを使用して、各文字xt と前の文字 x <t 間の注意分布をモデル化します。

Stick-Breaking 自己注意の利点は、位置埋め込みや相対位置バイアスなどの追加の位置情報を導入することなく、最も近い関連文字に自動的に焦点を合わせることです。

t 個の時間ステップを含む入力ベクトルのシーケンス x1、x2、…、xt が与えられた場合、各入力はキー ベクトルのシーケンス k1、k2、…、kt と値ベクトルのシーケンス v1、v2、…、vt に投影されます。時刻 t における注目度を計算するために、入力 x_t はクエリ ベクトル q_t = W_q x_t に投影されます。ここで、W_q はクエリ投影行列です。これまでのすべてのステップと現在のステップ i ≤ t について、時間ステップ i のキーが時間ステップ t のクエリと一致する確率を計算します。

このクエリ確率はシグモイド活性化関数を使用するため、正規化されていないことに注意してください。次に、クエリ確率はスティックブレイクプロセスを通じて正規化されます。

このようにして、瞬間 t に最も近く、クエリ確率が大きい瞬間に注意が自動的に割り当てられます。これにより、自己注意メカニズムは追加の位置情報なしで相対位置を効果的にモデル化できるようになります。最後に、自己注意モジュールの出力は、注意重みによって履歴値ベクトルを合計して投影することによって得られます。

ModuleFormer でのモジュール制御

事前トレーニングにおける負荷分散

SMoE が同じモジュールを繰り返し使用して他のモジュールの余分な容量を無駄にすることを避けるために、通常、負荷分散損失関数が採用され、各エキスパートの使用頻度が調整されます。以前の SMoE モデルとは異なり、チームは入力文字とモジュール間の相互情報量 (MI) を最大化することを目指しています。

簡単にするために、バッチX内のトークンの分布は均一であると仮定します。つまり、p(x) = 1/Xです。すべての定数成分を除去した後、相互情報損失（式6）はp(m)のエントロピーとp(m | x)の条件付きエントロピーの差に簡略化できます。

上記において、p(m) = sum_x(g(m|x)p(x)) であり、p(x) はバッチ内の各文字の確率、H(m) はモジュール分布の限界エントロピー、H(m|x) は入力文字 x が与えられたモジュールのエントロピー、|X| は入力文字の数です。長さ T およびサイズ B のミニバッチの場合、文字数は |X| = BT であり、文字の確率は p(x) = 1/|X| です。

直感的には、相互情報量損失は、入力 x が与えられた場合、モジュールの確率分布の限界エントロピーを最大化し、モジュールの条件付き分布の商を最小化します。各エキスパートの負荷をバッチ全体で分散させ（H(m)を最大化）、同時に各入力xがルーティング確率をより少ないモジュールに集中するように促します（H(m | x)を最小化）。

微調整における負荷集中

チームは事前トレーニング中に各エキスパートの能力を最大限に引き出したいと考えており、微調整中は少数のモジュールを下流のタスクに集中させたいと考えています。これにより、未使用のモジュールが削除され、微調整されたモデルのパラメータの数が削減されます。負荷をより少ないモジュールに集中させるために、チームはモジュールの限界エントロピーを最小限に抑える新しい負荷集中損失関数を導入しました。

これにより、モデルは下流のタスクを処理するために使用するモジュールの数が少なくなります。微調整後、トレーニング セットまたは検証セットで使用されるモジュール頻度 f_m を計算できます。 f_m はこのタスクにおけるモジュール m の重要度を表し、f_m が特定のしきい値未満のエキスパートを削除することでモデルのプルーニングを簡単に実現できます。

新しいモジュールを使用して新しい知識を学習する

モジュラー モデルの場合、新しいモジュールを挿入することは、モデル全体を微調整せずに新しい知識を学習するための簡単でパラメータ効率の高い方法です。各層にN_newのランダムに初期化されたモジュールを挿入する場合、ルータ内の層Aを埋め込むモジュール（式2のA）も、形状が（N_new、D_rtr）の新しい行列A'を含むように拡張する必要があります。したがって、

新しいルーティング関数は次のように記述できます。

微調整中は他のモジュール パラメータが固定されるため、新しいモジュールを使用した継続的な学習によって、壊滅的な忘却の問題を大幅に回避できます。

ただし、致命的な忘却はルーティング機能に影響を及ぼす可能性があります。新しいモジュールが新しいドメインでトレーニングされるとき、ルーティング機能が古いドメインからの入力を新しいエキスパートに誤ってルーティングすると、モデルが壊滅的な忘却に悩まされる可能性があります。

これを回避するために、チームは壊滅的な忘却を回避するためにルーティング機能を正規化し、2 つのトレーニング戦略を提案しました。

1) 完全に微調整されたルーティング。式 9 では、A と B は事前トレーニング済みのパラメータを使用して初期化され、A' はランダムに初期化されます。この戦略は、トレーニング データに新しいデータと古いデータの両方が含まれている状況向けに設計されています。

2) A' のみをトレーニングします。この戦略は生涯学習の状況向けに設計されており、以前にトレーニングされたデータは使用しません。この状況では、新しいモジュールが頻繁に使用されすぎて、致命的な忘却につながる可能性があります。チームは、A' のノルムを制限するために正規化項を導入しました。

欠陥があると指摘されている継続学習の従来の正則化方法 (減衰や L2 損失など) とは異なり、ルーティング正則化は専門家の能力を制限するのではなく、新しい専門家を使用する傾向を制限するだけです。

評価する

研究者は、ModuleFormer に基づいて、Pile データセットでサイズと計算量の異なる 3 つの ModuleFormer 言語モデル (MoLM) 言語モデルを事前トレーニングしました。

基本的なパフォーマンス評価

チームは、言語モデル評価ハーネスを使用して、ゼロショット、少数ショット、言語モデリング タスクで言語モデルを評価しました。

ゼロショットタスクと少数ショットタスクの両方において、目標は、コンテキストに応じて、指定された一連のオプションから最も適切な完了を選択することです。コンテキストに応じて最も確率の高い最終的な完了が選択されます。

言語モデリングでは、Wikitext データセットでテストが実行されます。目標は、次のトークン予測の困惑を最小限に抑えることです。

コード生成のために、モデルは HumanEval データセットで評価されます。 HumanEval には 164 個の手書きの Python プログラミング問題が含まれています。モデルは、提供されているすべてのテスト ケースに合格できるように、タスクの説明プロンプトに従って関数を完了する必要があります。

表 2 と 3 は、常識的推論、クローズドブック質問応答、コード生成ベンチマークにおける MoLM とベースライン言語モデルのパフォーマンスを示しています。

全体的に、MoLM-4B-K2 モデルのパフォーマンスは約 13 億のパラメータを持つ密なモデルに匹敵し、MoLM-4B-K4 モデルと MoLM-8B-K2 モデルのパフォーマンスは約 27 億のパラメータを持つ密なモデルに匹敵します。

MoLM は計算構造が疎であるため、同等のパフォーマンスを持つ密なモデルとして、各文字の活性化パラメータ (および同等の計算労力) の約 25% のみを処理します。その結果、ピーク時のメモリ使用量を抑えながらレイテンシを 50% 削減し、GPU メモリが完全に占有されているときのスループットを 2 倍に向上させます。

モジュールを追加して新しい言語を学ぶ

このセクションでは、新しい言語を学習するモデルの能力をテストしました。継続的な共同事前トレーニングと継続的な生涯事前トレーニングという 2 つの主要な実験設定が研究されています。

両者の違いは英語のテキストの有無にあります。どちらの設定でも、CC-100 コーパスで言語モデリング タスクを実行して、ModuleFormer と GPT-Neo を継続的に事前トレーニングします。品質を評価するために、XGLM と mGPT によって導入された 0 ショット法の mLAMA ベンチマークを採用します。

継続的な共同事前トレーニング: この部分では、共同トレーニングされたモデルに対して継続的な事前トレーニングを実行します。具体的には、英語と新しい言語を組み合わせて新しいトレーニング コーパスを構築し、埋め込みレイヤーをトレーニング可能な状態に保ちます。共同トレーニング [Caruana、1997] は、古いタスクと新しいタスクの両方の習得を示す、よく知られたマルチタスク学習アプローチです。しかし、多くの場合、異なるタスク間でマイナスの干渉が生じます。

表4は、モデルを継続的にトレーニングして得られた結果を示しています。表から次のことがわかります。

1) 完全に調整されたケースではスパース モデルの方が干渉が少なくなり、最終的に最高のパフォーマンスが達成されたことがわかりました。
2) ModuleFormer は、モジュールを追加することで (新しいエキスパートを挿入)、以前の LoRA 方式よりも少数のパラメータを調整する能力 (パラメータ効率) が優れています。これらの結果は、スパースなアーキテクチャが忘却に対するより強い抵抗につながることを示唆しています。

継続的な生涯事前トレーニング: この実験設定では、モデルは新しい言語テキストのみでトレーニングされます。アブラハムとロビンズ[2005]は、モデルが直面する困難な課題を説明する安定性と可塑性のジレンマを提唱しました。
1) モデルは新しい言語を学習するための高い可塑性を持つ必要がある。
2) モデルは、多くのトレーニング反復中に英語のトークンにさらされることがないことを考慮して、優れた安定性を備えている必要があります。

表 5 は、異なるルーティング正規化損失重みでの LoRA ベースラインと当社の方法の結果を示しています。当社の ModuleFormer は、ルーティング正規化損失の助けを借りて、安定性と可塑性のバランスをとる強力な能力を発揮します。

損失重みを増やすことで新しいエキスパートの使用を制限すると、モデルは安定性を獲得しますが、可塑性は失われます。対照的に、LoRA を使用して GPT-Neo を微調整すると、安定性と可塑性の両方で遅れをとります。

1億3,300万のトレーニング可能なパラメータを持つ高ランク LoRA と比較すると、低ランク LoRA (トレーニングパラメータを2,400万に削減) と基本的な正規化では安定性が向上しません。

モデルの微調整と圧縮

このセクションでは、ModuleFormer のモジュールをすばやく削除して、パフォーマンスを損なうことなく、サイズが小さいタスク固有のモデルを作成できることを示します。

まず、Python コードのみを含む GitHub-code-clean データセットから 150 億文字のサブセットを作成しました。次に、重み 0.001 の負荷集中損失関数を使用して、このデータセットで MoLM-4B-K2 モデルを微調整します。

微調整後、微調整データセットからランダムにサンプリングされた小さな評価セットで各エキスパートのアクティブ化頻度を計算し、各レイヤーをレイヤー内の最大頻度で割って正規化します。その後、閾値τを設定し、正規化された頻度がこの閾値より低いすべてのモジュールを削除しました。

私たちは、プルーニングされた MoLM-4B-K2 モデルを HumanEval データセットでテストしました。

図 2a は、pass@k メトリックと残りのパラメータの割合の相関関係を示しています。図 2b は、残りのパラメータ比と閾値の相関関係を示しています。私たちは次のことを観察しました:

1) 不要なモジュールを削除しても結果にはほとんど影響がありません。パフォーマンスを犠牲にすることなく、パラメータの 40% ～ 50% を削減できます。逆に、適切な剪定 (33%) により、微調整されたモデルがタスクでより優れたパフォーマンスを発揮できるようになります。

2) モジュールの分布には大きな違いがあり、モジュールの約半分は最も頻繁に使用されるエキスパートの 0.3% 未満でアクティブ化されています。この結果は、荷重集中損失関数の有効性を示しています。

要約する

本稿では、新しいモジュールアーキテクチャであるModuleFormerと、関連するモジュール操作方法を提案します。

ModuleFormer には、新しい Stickbreaking 注意メカニズム、事前トレーニング用の新しい相互情報負荷分散損失関数、微調整用の新しい負荷集中損失関数など、いくつかの新しいコンポーネントが含まれています。

ModuleFormer に基づいて、新しい言語モデル MoLM を事前トレーニングしました。私たちの実験結果では、MoLM は高密度 LLM と比較していくつかの新しい機能を備えていることがわかりました。

1) 低レイテンシ (50%) とより小さなメモリフットプリントで高密度 LLM と同じパフォーマンスを実現し、スループットを 2 倍以上向上します。

2) モデル全体が新しいドメインに適応するように微調整された後、モデルは壊滅的な忘却に対して堅牢であり、新しい言語や知識を学習するように簡単に拡張することもできます。

3) 下流のタスクを微調整できるため、一部のモジュールがタスクに集中し、タスクで使用されていないモジュールはパフォーマンスに影響を与えずに削減できます。

論文アドレス: https://arxiv.org/abs/2306.04640