陳丹奇チームの新しい研究: Llama-2 コンテキストが 128k に拡張され、メモリが 1/6 でスループットが 10 倍に

Chen Danqi のチームは、新しい LLMコンテキスト ウィンドウ拡張メソッドをリリースしました。

8k トークン ドキュメントのみでトレーニングされ、Llama-2 ウィンドウを 128k まで拡張できます。

最も重要なのは、このプロセスでは元のメモリの 1/6しか必要とせず、モデルは10 倍のスループットを達成することです。

それに加えて、トレーニングコストも大幅に削減できます。

この方法を使用して 7B サイズの Alpaca 2 を変更するには、 A100 が 1 つだけ必要です。

チームは次のように述べた。

この方法が便利で使いやすく、将来の LLM に安価で効果的なロングコンテキスト機能を提供できることを願っています。

現在、モデルとコードはHuggingFaceとGitHubで公開されています。

2つのコンポーネントを追加するだけ

この方法はCEPEと呼ばれ、「Context Expansion with Parallel Encoding 」の略です。

軽量フレームワークであるため、事前にトレーニングされ微調整されたモデルのコンテキスト ウィンドウを拡張するために使用できます。

事前トレーニング済みのデコーダーのみの言語モデルの場合、CEPE は次の 2 つの小さなコンポーネントを追加してモデルを拡張します。

1 つは、長いコンテキストのブロック エンコード用の小さなエンコーダです。

1 つは、デコーダーの各レイヤーに挿入され、エンコーダーの表現に焦点を当てる、交差アテンション モジュールです。

完全なアーキテクチャは次のとおりです。

この図では、エンコーダー モデルが 3 つの追加のコンテキスト ブロックを並列にエンコードし、最終的な隠し表現と連結して、デコーダーの交差注意層への入力として機能します。

ここで、交差注意層は、デコーダー モデル内の自己注意層とフィードフォワード層の間のエンコーダー表現に焦点を当てています。

CEPE は、ラベル付けされていないトレーニング データを慎重に選択することで、モデルが長いコンテキストに対応できるようにし、ドキュメント検索に優れた機能を発揮できるようにします。

著者は、このような CEPE には 3 つの大きな利点があると紹介しています。

（１）長さは一般化できる

位置エンコーディングによって制約されないため、コンテキストはセグメントでエンコードされ、各セグメントには独自の位置エンコーディングがあります。

（２）高効率<br>小型エンコーダと並列エンコードを使用してコンテキストを処理することで、計算コストを削減できます。

同時に、クロスアテンションはエンコーダーの最後のレイヤーの表現にのみ焦点を当て、デコーダーのみを使用する言語モデルは各レイヤーの各トークンのキーと値のペアをキャッシュする必要があるため、比較すると CEPE ははるかに少ないメモリしか必要としません。
（３）研修コストの削減

完全な微調整方法とは異なり、CEPE は、大規模なデコーダー モデルを固定したまま、エンコーダーと交差アテンションのみを微調整します。

著者らは、7Bデコーダーを400Mエンコーダーとクロスアテンションレイヤー（合計14億パラメータ）を備えたモデルに拡張することで、80GBのA100 GPUで完了できることを紹介しています。

混乱は減り続けている

研究チームは CEPE を Llama-2 に適用し、200 億トークン(Llama-2 の事前トレーニング予算のわずか 1%)を使って RedPajama のフィルターバージョンでトレーニングしました。

まず、完全に微調整された 2 つのモデル (LLAMA2-32K と YARN-64K) と比較すると、CEPE はすべてのデータセットでより低いまたは同等のパープレキシティを達成しながら、メモリ使用量が少なく、スループットが高くなります。

コンテキストを 128k (トレーニングの長さ 8k を大幅に超える)に増やすと、CEPE の困惑度は低いメモリ状態を維持しながら減少し続けます。

対照的に、Llama-2-32K と YARN-64K は、トレーニングの長さを超えて一般化できないだけでなく、メモリ コストも大幅に増加します。

第二に、検索機能が強化されます。

次の表に示すように:

取得したコンテキストを使用することで、CEPE はモデルの複雑性を効果的に改善し、RePlug よりも優れたパフォーマンスを発揮できます。

注目すべきは、k=50 段落（トレーニング用に 60 段落）であっても、CEPE は難解さを改善し続けていることです。

これは、CEPE は検索強化設定にうまく移行できるが、フルコンテキスト デコーダー モデルではこの能力が低下することを示唆しています。

3 番目に、オープン ドメインの質問応答機能が大幅に向上しました。

下の図に示すように、CEPE はすべてのデータセットと段落 k パラメータで他のモデルを大幅に上回っており、他のモデルとは異なり、k 値が増加するとパフォーマンスが大幅に低下します。

これは、CEPE が大量の冗長または無関係な段落に対して敏感ではないことも示しています。

まとめると、CEPE は、はるかに低いメモリと計算コストで、上記のすべてのタスクにおいて他のほとんどのソリューションよりも優れたパフォーマンスを発揮します。

最後に、これらに基づいて、著者らは命令チューニングモデル専用のCEPE-Distilled (CEPED)を提案しました。

ラベルなしデータのみを使用してモデルのコンテキスト ウィンドウを拡張し、補助的な KL ダイバージェンス損失を通じて元の命令チューニング モデルの動作を新しいアーキテクチャに抽出し、高価な長いコンテキスト命令トレース データを管理する必要性を排除します。

最終的に、CEPED は Llama-2 のコンテキスト ウィンドウを拡張し、指示を理解する能力を維持しながらモデルの長いテキストのパフォーマンスを向上させることができます。

チームについて

CEPE には合計 3 人の著者がいます。

一人はプリンストン大学でコンピューターサイエンスの修士課程に在籍するハワード・イェン氏だ。

2人目は、清華大学を卒業し、同校の博士課程に在籍する高天宇さん。

彼らは全員、責任著者であるChen Danqi氏の学生です。

原著論文: https://arxiv.org/abs/2402.16617
参考リンク: https://twitter.com/HowardYen1/status/1762474556101661158