純粋な MLP は下流のタスクには適していませんか? Meta AIらは、トランスフォーマーを上回るスパースMLPを提案した。

注意ベースのモデルの代替として、純粋な MLP アーキテクチャがますます注目を集めています。 NLP では、gMLP などの最近の研究により、純粋な MLP は言語モデリングにおいてトランスフォーマーと同等のパフォーマンスを達成できるものの、下流のタスクではトランスフォーマーよりも弱いことが示されています。

Meta AI とニューヨーク州立大学バッファロー校の研究者は、表現力の面での MLP の限界を分析し、特徴と入力 (トークン) の両方の次元で専門家の混合 (MoE) を備えたスパースにアクティブ化された MLP を提案しました。このスパースな純粋な MLP は、計算を変更せずにモデルの容量と表現力を大幅に向上させます。この研究では、条件付き計算と 2 つのルーティング戦略を組み合わせるという重要な課題に取り組んでいます。

論文アドレス: https://arxiv.org/pdf/2203.06850.pdf

トランスフォーマーベースのMoE、高密度トランスフォーマー、純粋MLPと比較して、本研究で提案されたスパース純粋MLP（sMLP）は、言語モデリングの難解さを改善し、トレーニング効率を最大2倍向上させます。最後に、6 つのダウンストリーム タスクでスパース純粋 MLP のゼロ ショット コンテキスト学習パフォーマンスを評価し、トランスフォーマー ベースの MoE および高密度トランスフォーマーよりも優れていることがわかりました。

方法

スパース活性化を伴う純粋なMLP

sMLP の全体的なアーキテクチャを図 2 に示します。このアーキテクチャには、N_1 個の密なブロックと N_2 個の疎なブロックが含まれています。 N_1 と N_2 はどちらもハイパーパラメータです。各スパース ブロックには 2 つのモジュールが含まれます。

• tMoEモジュール：この研究では、Base LayersのMoE（Lewis et al.、2021）を使用して、dense transformer（Vaswani et al.、2017b）のFFNモジュールを置き換えます。
• sMoEモジュール：この研究では、トランスフォーマーの自己注意モジュール（Vaswani et al.、2017b）とgMLPの空間ゲーティングユニット（SGU、Liu et al.、2021a）を置き換えるsMoEモジュールを設計しました。

tMoE モジュールと sMoE モジュールはどちらも 2 つの要素で構成されています。

• エキスパートモジュール

エキスパート モジュールは入力の処理を担当します。 tMoE モジュールの場合、上の図 2 に示すように、各エキスパートには FFN が含まれます。 sMoE モジュールの場合、各エキスパートには、図 6 (右) に示すように空間ゲーティング ユニットが含まれています。

• ゲーティング機能

このモジュールは、入力の各部分をどのエキスパートが処理するかを決定するため、MoE 構造を機能次元に拡張するための独自のルーティング方法を設計する必要があります。

図 3 (左) は、既存のトランスフォーマーベースの MoE (Lepikhin et al., 2020; Fedus et al., 2021; Lewis et al., 2021; Roller et al., 2021) のゲーティング関数の例を示しています。 x_ij は、i_番目のトークン内の j_番目の隠し次元の値を表します。

式（３）に示すように：

tMoEは、式(3)のパラメータ化によって記述された学習されたゲーティング関数を使用して、これらの4つのトークンをFFN層の3人のエキスパートに送信します。既存のいくつかの MoE とは異なり、この研究では、スパースな純粋な MLP アーキテクチャにおいて、図 3 (右) に示すように、隠し次元に沿って隠し表現をブロックし、ブロック ベクトルを異なるエキスパートに送信することを提案しています。

これらの既存の MoE とは異なり、本研究では、スパースなフル MLP アーキテクチャにおいて、潜在次元に沿って潜在表現をチャンク化し、チャンク化されたベクトルをさまざまなエキスパートに送信することを提案しています (図 3 (右) を参照)。

フィーチャ空間でのルーティング

ルーティング トークンと比較すると、ルーティングの隠し次元は、将来のトークンを単純に予測すると情報が漏洩するため、自己回帰モデルでは独特の課題に直面します。さらに、自己注意を備えたTransformerベースのMoEとは異なり、情報漏洩を防ぐために適切なマスキングを直接適用することはできないため、TransformerベースのMoEの言語モデリングには既存のルーティング方法を採用することはできません。この研究では、決定論的ルーティングと部分予測という 2 つのソリューションを比較しました。

実験と結果

トークン操作の比較

この研究では、sMLP モデルを 2 つの高密度モデル、Transformer (Vaswani et al., 2017b) と gMLP (Liu et al., 2021a) と比較しました。フル MLP ベース モデルとトランスフォーマー ベース モデルの主な違いは、トークン操作にあります。この研究では、Transformers の自己注意モジュール、gMLP の空間ゲーティング ユニット、sMLP モデルの sMoE モジュールという 3 つのトークン単位の操作を比較しました。表 3 は、3 つのトークン操作とそれぞれのヘッダー メカニズムを比較したものです。

下の図 4 は、このモデルと、ヘッドの数が異なる密なモデルを比較したものです。 Transformer モデルは、マルチヘッド メカニズムの恩恵を大いに受けます。ただし、gMLP モデルではパラメータの数が増えても、マルチヘッド メカニズムによってパフォーマンスは向上しません。 sMLP モデルは gMLP のマルチヘッド ソリューションとして見ることもできます。これにより、MLP ベースのモデルのパフォーマンスが大幅に向上し、トランスフォーマー モデルよりも優れたパフォーマンスを発揮します。

スパースMLPの結果

下の図 5 は、トレーニング ステップ数 (上) とトレーニング時間 (下) で測定された品質 (有効な困惑度) とトレーニング効率を示しています。研究者らは、2 つのルーティング戦略バリアントを備えた sMLP が、ほぼ同じ FLOP 数を持つ最先端のトランスフォーマーベースの MoE モデルよりも優れていることを発見しました。

以下の表4は、主な実験における詳細な比較結果をまとめたものです。研究者らは、すべてのモデルの FLOP を約 0.8T に制御しました。モデル層の数の違いに加えて、埋め込み次元は 1024、隠し次元は 4096 です。ご覧のとおり、sMLP モデルは 25,000 のトレーニング ステップで最高の一般化を達成しながら、最高のトレーニング速度を達成しています。 HASH レイヤーは、すべての Transformer ベースラインの中で最もパフォーマンスが高く、所要時間も最も短くなります。

拡張機能

モデルのスケーラビリティをテストするために、この研究ではモデル サイズのトレーニングを 2.0 TFLOP に増やしました。表4（下）に結果をまとめます。

表 4 (上部) のモデルと比較して、この研究では、表 5 に示すように、埋め込みを 1024 から 2048 に変更し、潜在次元を 4096 から 8192 に調整することで、すべてのモデルをスケールアップしました。この研究では、表 2 に示すように、事前トレーニング データのサイズも増加しました。

​​