ソフトマックスを放棄した初の大規模線形アテンショントランスフォーマーモデル: 1750億のパラメータ、より優れた速度と精度

最近、上海人工知能研究所とOpenNLPLabの研究チームが、ソフトマックスベースの注意メカニズムを完全に放棄し、代わりに新しく提案された線形注意を使用した新しい大規模言語モデルTransNormerLLMを提案しました。 TransNormerLLM は、線形注意に基づく最初の大規模言語モデル (LLM) であり、精度と効率の点で従来の Softmax 注意ベースのモデルよりも優れていることが報告されています。研究者らは、事前トレーニング済みモデルのオープンソース版もリリースする予定だ。

論文: https://arxiv.org/abs/2307.14995
モデル: https://github.com/OpenNLPLab/TransnormerLLM

大規模言語モデルは、自然言語処理 (NLP) の分野に革命をもたらしました。これらは、さまざまな種類のタスクの処理に優れており、計算フレームワーク内で人間の言語を理解、生成、対話する能力を向上させます。言語モデリングにおけるこれまでの開発は主に Transformer アーキテクチャを中心に行われており、主力モデルには基本的な Transformer、GPT シリーズ、BERT、BART など、画期的なモデルが含まれています。 Transformer アーキテクチャの成功は、データ駆動型モデルソリューションで入力トークン間の依存関係を識別できるソフトマックスアテンションメカニズムに基づいています。また、グローバル位置を認識することもできるため、モデルは自然言語の長距離ダイナミクスを効果的に処理できます。

それでも、従来のトランスフォーマーには依然として限界があります。まず第一に、シーケンスの長さに対して 2 次的な時間計算量があるため、スケーラビリティが制限され、トレーニングフェーズと推論フェーズの両方で計算リソースと時間効率が低下します。この二次の時間計算量を線形計算量に単純化するために、多くの研究者がさまざまなシーケンスモデリング手法を提案してきました。しかし、これらの方法は、次の 2 つの理由から LLM に使用するのが困難です。1) 言語モデリングタスクでのパフォーマンスが低い場合が多い。2) 実際のシナリオでは速度上の利点が見られない。

この論文で提案されている TransNormerLLM は、線形アテンションに基づく最初の LLM であり、精度と効率の両方で従来のソフトマックスアテンションを上回ります。 TransNormerLLM は、以前の線形アテンションアーキテクチャ TransNormer をベースに構築されており、パフォーマンスを向上させるためにいくつかの変更も行われています。 TransNormerLLM の主な改善点には、位置埋め込み、線形注意加速、ゲーティングメカニズム、テンソル正規化、推論加速などがあります。

特に注目に値する改善点の 1 つは、TransNormer の DiagAttention を線形アテンションに置き換えたことです。これにより、全体的なインタラクティブパフォーマンスが向上します。研究者らは希釈問題を解決するために指数関数的減衰を伴う LRPE も導入しました。さらに、研究者らは、トレーニング中の線形注意の速度を2倍にし、IOを感知することでメモリ使用量を4分の1に削減できるという、Lightning Attentionと呼ばれる新しい技術も紹介した。それだけでなく、GLU と正規化の方法も簡素化され、後者によって全体の速度が 20% 向上しました。また、彼らは、異なるシーケンス長でも数値安定性と一定の推論速度を確保できる堅牢な推論アルゴリズムを提案し、それによってトレーニング段階と推論段階の両方でモデルの効率を向上させました。

TransNormerLLMの有効性を検証するために、研究者らは6TBを超えるサイズと2兆を超えるトークンを持つ大規模なコーパスを慎重に収集しました。データの品質を保証するために、収集されたコーパスをフィルタリングするためのセルフクリーニング戦略も開発されました。表 1 に示すように、研究者は元の TransNormer モデルを拡張し、3 億 8,500 万から 1,750 億の範囲のパラメータを持つ複数の TransNormerLLM モデルを取得しました。その後、研究者らは新しい大規模コーパスに基づいて包括的な実験と制御変数の研究を実施し、その結果、新しい方法はソフトマックス・アテンション・ベースの方法よりも優れており、トレーニングと推論の速度が速いことが示されました。

表1: TransNormerLLMモデルのさまざまなバリエーション

LLM 分野の研究を促進するため、上海人工知能研究所と OpenNLPLab の研究者らは、事前トレーニング済みモデルもオープンソース化する予定です。研究者らによると、これは「研究者や実務者が私たちの研究結果に基づいてアプリケーションを構築し、LLM の効率的なトランスフォーマー構造を探求できるようにする」ためです。

トランスノーマーLLM

アーキテクチャの改善

以下は、TransNormerLLM のさまざまなモジュールと研究者によって提案されたいくつかの改善策の簡単な紹介です。

改善1: 位置エンコーディング

TransNormer の下位層では、希釈問題を回避するために DiagAttention を使用します。ただし、これにより、トークン間のグローバルな相互作用機能が欠如することになります。この問題に対処するために、研究者らは TransNormerLLM に指数関数的減衰を伴う LRPE (線形相対位置エンコーディング) を使用しました。これにより、下位層で完全な注意を維持できます。研究者たちはこの方法をLRPE-dと名付けた。

改善2: ゲート機構

ゲーティングにより、モデルのパフォーマンスが向上し、トレーニングプロセスがスムーズになります。研究者らは、論文「Transformer quality in linear time」の Flash メソッドを TransNormerLLM に使用し、トークン混合にゲート線形アテンション (GLA) 構造を使用しました。

モデルの速度をさらに向上させるために、ゲート自体が非線形性を導入する可能性があるため、元の GLU 構造の活性化関数を削除する Simple GLU (SGLU) も提案されました。

改善3: テンソル正規化

研究者らは、TransNormer で導入された NormAttention を使用しました。 TransNormerLLM では、RMSNorm を新しい単純な正規化関数 SimpleRMSNorm (略称 SRMSNorm) に置き換えました。

全体構造

図1はTransNormerLLMの全体構造を示しています。

このアーキテクチャでは、入力 X の更新は 2 つの連続したステップで実行されます。まず、SRMSNorm を使用して正規化されたゲート線形アテンション (GLA) モジュールを通過します。次に、SRMSNorm を使用して再度正規化されたシンプルゲート線形ユニット (SGLU) モジュールに渡されます。この全体的なアーキテクチャは、モデルのパフォーマンスの向上に役立ちます。このプロセス全体の疑似コードを以下に示します。

トレーニングの最適化

雷注意

注意計算を高速化するために、研究者らは Lightning Attention アルゴリズムを導入しました。これにより、新たに提案された線形注意が IO (入出力) 処理により適したものになります。

アルゴリズム 1 は Lightning Attention のフォワードパスの実装の詳細を示し、アルゴリズム 2 はバックワードパスを示します。研究者らは、勾配をより速く計算できる実装も持っており、将来リリースされる予定だと述べている。

モデルの並列化

すべてのモデルパラメータ、勾配、およびオプティマイザ状態テンソルをコンピュータクラスターに分散するために、研究者は Fully Sharded Data Parallel (FSDP) を使用しました。この戦略的なパーティショニングアプローチにより、各 GPU のメモリフットプリントが削減され、メモリの使用率が最適化されます。さらに効率を向上させるために、彼らはアクティベーションチェックポイントを使用しました。これにより、後方パス中にメモリにキャッシュされるアクティベーションの数が減少します。代わりに、これらの勾配が計算されるときに、それらは削除され、再計算されます。このテクノロジーは、コンピューティング効率の向上とリソースの節約に役立ちます。さらに、GPU メモリの消費量を削減しながら計算を高速化するために、研究者らは自動混合精度 (AMP) も使用しました。

上記の結果に加えて、研究者らは、主に NVIDIA の Megatron-LM モデル並列化にヒントを得て、線形トランスフォーマーでモデル並列化を実行することで、システムエンジニアリングの最適化をさらに進めました。従来のトランスフォーマーモデルでは、各トランスフォーマー層に自己注意モジュールがあり、その後に 2 層の多層パーセプトロン (MLP) モジュールが続きます。 Megatron-LM モデルの並列処理を使用する場合、これら 2 つのモジュールで独立して使用されます。同様に、TransNormerLLM 構造も SGLU と GLA という 2 つの主要モジュールで構成されており、これら 2 つのモデルの並列化は個別に実行されます。

堅牢な推論

これにより、TransNormerLLM は RNN の形式で推論を実行できるようになります。アルゴリズム 3 にこのプロセスの詳細を示します。しかし、数値の精度には問題があります。

これらの問題を回避するために、研究者は堅牢な推論アルゴリズムを提案しました。アルゴリズム 4 を参照してください。

オリジナルの推論アルゴリズムとロバスト推論アルゴリズムによって得られる結果は同じです。

コーパス

研究者らはインターネットから公開されている大量のテキストを収集したが、その総サイズは700テラバイトを超える。収集されたデータは、図 2 に示すように、データ前処理手順を経て処理され、約 2 兆個のトークンを含む 6 TB のクリーンなコーパスが残りました。透明性を高め、ユーザーの理解を深めるために、データソースを分類しました。表 2 に具体的なカテゴリを示します。

図2: データ前処理プロセス

表2: コーパス統計

実験

研究者らは、Metaseq フレームワークで PyTorch と Trition を使用して TransNormerLLM を実装しました。モデルは Adam オプティマイザーを使用してトレーニングされ、FSDP も使用されてモデルが NVIDIA A100 80G クラスターに効率的に拡張されました。また、パフォーマンスを最適化するために、モデルの並列処理技術を適切に使用しました。