LLaMA2コンテキストを10万に拡張し、MITと香港中文大学はLongLoRAメソッドを開発

一般的に、大規模なモデルを事前トレーニングする場合、テキストの長さは固定されます。より長いテキストをサポートする場合は、モデルを微調整する必要があります。ただし、長いコンテキストで LLM をトレーニングするには計算コストがかかり、多くのトレーニング時間と GPU リソースが必要になります。

たとえば、8192 長のコンテキストでモデルをトレーニングするには、2048 長のコンテキストでモデルをトレーニングする場合よりも 16 倍の計算リソースが必要です。それでも、コンテキストの長さは、LLM が応答時にコンテキスト全体を明確に理解する能力を表すため、モデルのパフォーマンスにとって非常に重要です。

最近、MITと香港中文大学が共同で研究を行い、LongLoRAを提案しました。これは、計算コストを制限しながら、事前トレーニング済みの大規模言語モデルのコンテキスト サイズを拡張する効果的な微調整方法です。

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

プロジェクトアドレス: https://github.com/dvlab-research/LongLoRA

この論文は、2つの側面からLLMのコンテキスト拡張を加速します。

一方、推論中には密なグローバル注意が必要ですが、疎なローカル注意を通じてモデルを効果的かつ効率的に微調整できます。本論文で提案されているシフトショートアテンションは、コンテキスト拡張を効果的に実現し、計算量を大幅に節約し、バニラアテンションを使用した微調整と同様のパフォーマンスを発揮します。

一方、コンテキスト拡張の​​ための LoRA は、トレーニング可能な埋め込みと正規化でうまく機能します。 LongLoRA は、7B/13B から 70B までの LLaMA2 モデルのさまざまなタスクで良好な結果を示しています。単一の 8x A100 デバイスでは、LongLoRA は LLaMA2 7B を 4k コンテキストから 100k に拡張し、LLaMA2 70B を 32k に拡張します。 LongLoRA は、元のアーキテクチャを維持しながらモデルのコンテキストを拡張し、FlashAttention-2 などの既存のほとんどの技術と互換性があります。 LongLoRA を実用的にするために、研究者は教師あり微調整用のデータセット LongQA を収集しました。データセットには、3,000 を超える長いコンテキストの質問と回答のペアが含まれています。

LongLoRA は、注意レベルと重みレベルの両方で、事前トレーニング済みの大規模言語モデルのコンテキスト拡張を加速できます。ハイライトは次のとおりです。

• Shift short attention は実装が簡単で、Flash-Attention と互換性があり、推論中に使用する必要がありません。
• LLaMA2-LongLoRA-7B-100k、LLaMA2-LongLoRA-13B-64k、LLaMA2-LongLoRA-70B-32k など、7B から 70B まで、コンテキスト長が 8k から 100k のすべてのモデルがリリースされています。
• 教師あり微調整のために、ロングコンテキスト QA データセット LongQA が確立されています。研究者らは、SFT の 13B および 70B 32k モデルである Llama-2-13b-chat-longlora-32k-sft と Llama-2-70b-chat-longlora-32k-sft をリリースしており、来月データセットをリリースする予定です。

LongLoRA 技術詳細

短い注意をシフトする

標準的な自己注意モデルの計算コストは​​ O(n^2) であるため、LLM メモリは高価になり、長いシーケンスでは遅くなります。トレーニング中にこの問題を回避するために、この論文では、下の図 2 に示すように、シフトショートアテンション (S^2-Attn) を提案します。

研究者らは、以下の表 1 に示すように、微調整の重要性を検証しました。微調整を行わないと、適切な位置埋め込みが装備されていても、コンテキストの長さが長くなるにつれてモデルのパフォーマンスが低下します。

既存の効率的な注意設計により、長いコンテキストの言語モデルの効率も向上します。以下の表 2 では、研究者らが S^2-Attn をいくつかの典型的な効率的なアテンションと比較し、前者は効率的な微調整を実現できるだけでなく、完全なアテンション テストもサポートできることを発見しました。

さらに、S^2-Attnは実装が簡単で、(1)ハーフアテンションヘッド内のトークンを変換し、(2)トークン次元からバッチ次元に特徴を移動するという2つのステップのみで済みます。このプロセスには数行のコードしか必要ありません。

ロングコンテキストLoRAの改善

LoRA は、LLM を他のデータセットに適応させるための効果的で一般的な方法です。完全な微調整と比較すると、トレーニング可能なパラメータとメモリコストが大幅に節約されます。しかし、LLM を短いコンテキスト長から長いコンテキスト長に適応させることは容易ではありません。研究者たちは、LoRA と完全な微調整の間に明らかなギャップがあることを観察しました。下の表 3 に示すように、ターゲット コンテキストの長さが長くなるにつれて、LoRA と完全な微調整の間のギャップも大きくなります。

このギャップを埋めるために、研究者たちはトレーニングのために埋め込み層と正規化層を有効にしました。表 3 に示すように、それらは限られたパラメータを占めますが、長いコンテキストの適応に影響を与えます。特に、正規化レイヤーは、LLaMA2 7B 全体のパラメータのわずか 0.004% を占めます。実験では、研究者はこの改良された LoRA を LoRA+ と呼びました。

実験と結果

研究者らは、事前トレーニング済みの 7B、13B、および 70B LLaMA2 モデルを拡張しました。拡張コンテキスト ウィンドウの最大サイズは、7B モデルの場合は 100k、13B モデルの場合は 65536、70B モデルの場合は 32768 です。

研究者は、位置補間のトレーニング ハイパーパラメータのほとんどを使用しましたが、1 つの 8×A100 GPU デバイスが一部のケースでのみ使用されたため、バッチ サイズは小さくなりました。すべてのモデルは、次のトークン予測目標に合わせて微調整されます。研究者らは、β_1 = 0.9、β_2 = 0.95 の AdamW を使用しました。7B モデルと 13B モデルの学習率は 2 × 10^−5 に設定され、70B モデルの学習率は 10^−5 に設定されました。

彼らは線形学習率ウォームアップも使用しました。重量はゼロまで減少します。デバイスあたりのバッチ サイズは 1 に設定され、勾配累積ステップは 8 に設定されています。つまり、8 つの GPU が使用され、グローバル バッチ サイズは 64 になります。モデルは 1000 ステップにわたってトレーニングされました。

研究者らはトレーニングに Redpajama データセットを使用し、教師あり微調整用にロングコンテキスト QA データセット LongQA を構築しました。 Redpajama によって微調整されたモデルは優れた困惑度を示していますが、チャット機能は限られています。研究者たちは、技術論文、SF小説、その他の書籍などの資料に関連する3,000以上の質問と回答のペアを収集しました。設計された質問には、概要、関係性、登場人物などが含まれます。

下の表 4 から、トレーニング コンテキストと評価コンテキストの長さが同じ場合、コンテキスト サイズが大きくなるにつれて困惑度が減少することがわかります。

以下の表 5 では、研究者らは、単一の 8×A100 デバイスで微調整できる最大コンテキスト長をさらに調べています。これらは、LLaMA2 7B、13B、70B をそれぞれ 100k、65536、32768 のコンテキスト長に拡張します。 LongLoRA は、これらの超大規模な設定で有望な結果を達成します。さらに、実験により、拡張モデルではコンテキストの長さが短くなると、ある程度の困惑度の低下が見られることも判明しました。

以下の表 6 では、研究者が LongChat で導入されたトピック検索タスクについて、このモデルを他のオープン LLM と比較しています。このタスクは、会話の長さが 3k、6k、10k、13k、16k に及ぶ非常に長い会話からターゲット トピックを取得することです。

アブレーション実験

以下の表 7 では、研究者らが LLaMA2 7B をさまざまな種類の層に細分化しています。彼らは FLOP を分析しました。完全な注意については、コンテキストの長さが長くなるにつれて、Attn の割合も劇的に増加します。たとえば、コンテキスト長が 8192 の場合、Attn は合計 FLOP の 24.5% を占めますが、コンテキスト長が 65536 の場合は 72.2% に増加します。 S^2 -Attnを使用すると39.4%に低下します。

以下の表 8 は、PG19 検証セットで 8192 コンテキスト長に拡張された場合の LLaMA2 7B モデルの複雑さと微調整ステップ数の関係を示しています。微調整を行わないと、ステップ 0 ではモデルの長期コンテキスト機能が制限されることがわかります。完全な微調整は、低レベルのトレーニングよりも速く収束します。 200歩ほど歩いたところで二人は徐々に接近し、最後には大きな差はなくなりました。

以下の表 2 は、微調整中のさまざまな注意モードの効果を示しています。

エフェクト表示

モデルは、「ハリー・ポッター」の内容を読んだ後、なぜスネイプがハリーを嫌っているように見えるのかを説明し、登場人物同士の関係を要約することもできます。

それだけでなく、紙に印刷することで、関連情報をすぐに理解できるようになります。

詳細は原文をご参照ください。