言語モデルのメモリ管理の進展

大規模言語モデル（LLM）は、より複雑な情報を扱うのが上手くなってきた。最近の研究では、非常に長い入力シーケンスを効率的かつコスト効果的に処理する方法に焦点が当てられている。この一つの方法は、以前の入力の重要な部分を保存するために、先入れ先出し（FIFO）メモリを使うこと。これにより、新しい入力を処理しながら過去の関連情報に集中できる。ただ、この方法はかなりのメモリを必要とするし、異なるモデル構造に合わせて調整しなきゃいけない。もう一つの課題は、いくつかの方法ではモデルが過去の入力しか見られず、双方向で情報を分析する特定の言語モデルには使いにくいこと。

これらの問題に対処するために、古い情報をいつ削除するかを決めるポリシーを提案することで、メモリ使用量を抑えつつ、さまざまなモデルに適応できるようにしたい。さらに、Attendreレイヤーと呼ばれる新しいレイヤーを紹介することで、情報が追い出された後でも過去の入力のメモリを使えるようにした。このアイデアを、複数選択肢の質問を含む読解タスクで試して、うまくいくことを示した。

#背景

LLMが進化して、構造化された文書や混合メディアなど、より複雑な入力を扱えるようになった。これには、長い入力を扱うための賢い方法が必要となる。従来のモデルはドットプロダクトアテンションを使用していて、入力の長さが増えると時間コストが急速に増える。さまざまな方法がこのコストを下げようとしてきたが、多くの場合、全入力を一度に読む必要があるため、シーケンスの長さに制限がある。

さらに長い入力を扱うために、研究者たちは入力を小さく管理しやすい部分、つまりチャンクに分けて一つずつ処理する方法を提案している。一部の方法は、これらのチャンク間で関連情報を追跡するためのメモリの形を利用している。Memorizing Transformerと呼ばれるアプローチは、以前の入力から重要な詳細を保存して、今後の処理に役立てている。ただし、全シーケンスを処理するモデルと同じ性能を得るには、かなりのメモリが必要だ。

StreamingLLMと呼ばれる別の方法は、メモリ使用量を減らせるが、特定の入力タイプに対して最適に機能する。多くの既存の方法は、入力チャンクを順番に読むため、未来の情報にアクセスできず、応用が制限される。これは、過去と未来の文脈の両方を頼りにしてパフォーマンスを向上させる言語モデルにとって特に問題だ。

#私たちのアプローチ

私たちの目標は、メモリの要件を最小限にしながら、さまざまなモデルに柔軟性を持たせるメモリ設計を作ることだ。さらに、モデルが決定を下すときに過去と未来の入力の両方を考慮できるようにしたい。どの情報を保持するかを決める新しい追い出しポリシーを導入することで、重要度の低い情報を削除する。

最初に紹介するポリシーはLRA、または最近最も注目されたもの。これは、特定の情報がどれだけ重要かを示すアテンションスコアを考慮して、保持すべき情報と削除すべき情報を決める。2つ目のポリシーはLFA、または最も頻繁に注目されたもので、時間を通じてその情報がどれだけ使われてきたかを考慮する。

これらのポリシーはメモリ管理プロセスを導き、モデルが最も関連性の高い情報を保持しつつ、重要度の低い詳細を追い出すことを可能にする。さらに、Attendreレイヤーを発展させ、過去のクエリに基づいて追い出されたメモリ情報にアクセスできるようにして、モデルの長い文脈を扱う能力と理解力を向上させる。

#メモリタイプと管理

私たちのモデル設計では、2つのタイプのメモリに注目している。最初は、将来の使用のためにデータのグループを保存するシンプルなメモリで、一部の言語モデルが後のクエリのために重要な情報を覚える方法と似ている。2つ目のメモリタイプは、検索可能なキーを追加して、クエリに基づいて特定の情報を取得できるようにする。

どちらの場合でも、メモリサイズを管理しやすく保つことを目指している。大きなメモリはより良い文脈を提供できるが、高い時間と空間コストも伴う。異なる追い出しポリシーを探求して、メモリサイズと効率の最適なバランスを見つける。これらのポリシーを実装することで、処理を遅くする可能性のある不要な古いエントリを保持しないようにしている。

#Attendreレイヤー

Attendreレイヤーは、モデルがメモリとどのように相互作用するかを管理するように設計されている。このレイヤーは、クエリと関連するキー・バリューペアを別々のメモリストアに挿入することで機能する。新しいクエリがすぐに新しい情報にアクセスする従来の方法とは異なり、Attendreレイヤーでは、クエリが一時的に脇に置かれた前のエントリを参照できる。つまり、過去と未来の文脈が処理に影響を与えることができる。

入力シーケンスがモデルに供給されると、Attendreレイヤーはプロセス内でメモリの構造を維持するのを助ける。このセットアップにより、モデルは応答に深みを加える歴史的情報と関わることができ、長い文脈のより包括的な理解を提供する。

#実験：TriviaQAタスク

私たちは、関連文書に基づいて質問に答えるTriviaQA読解タスクで私たちの方法をテストした。異なるアプローチで訓練された2つの言語モデルを使用した。私たちの目的は、新しいメモリ方法がさらに微調整を必要とせずにどれだけ効果的に機能するかを確認することだった。

各質問では、モデルが特定の順序で配置されたさまざまなソースから情報を引き出す必要があった。質問を文脈の前または後に戦略的に配置することで、私たちのメモリポリシーがモデルの理解と応答能力をどれだけ支えるかを評価できた。

実験結果は、LRAおよびLFAポリシーを使用するモデルが従来のアプローチよりかなり良いパフォーマンスを示した。メモリサイズが小さくても、これらの新しい方法は強いパフォーマンスを維持できることが示され、長い文脈を扱う際の価値を証明した。

#PaLM 2-Sモデルの結果

PaLM 2-Sモデルを使用した試験では、入力の長さとメモリサイズを変化させた。メモリなしのモデルは長い入力シーケンスで苦労したが、より多くの文脈を処理させることでパフォーマンスが改善された。メモリ方法を導入したとき、モデルは大きな向上を示し、特にLRAおよびLFAポリシーを通じて効果的なメモリ使用を許可したときに顕著だった。

結果は、知的ポリシーを持った小さなメモリでも、大きな文脈を活用する大きなモデルと同等のパフォーマンスを発揮できることを示した。ポリシーにより、保持するメモリの選択が改善され、読解タスクにおけるパフォーマンスが向上した。

#FLAN-T5 XXLモデルの結果

FLAN-T5 XXLモデルへの評価を拡張した。これは長文処理を促進するように特別に訓練されていなかったにもかかわらず、モデルは長い入力を扱い、すでに私たちのメモリ戦略によって強化されたモデルの結果に競争力のある成果を達成した。

前のモデルと同様に、私たちの調査結果はLRAおよびLFAポリシーを使用することで出力スコアが顕著に改善されたことを示した。ここでも、メモリサイズと効率的なポリシーの相互作用が、メモリ制約の中でも強いパフォーマンスを生む結果につながった。

#結論

結論として、メモリモジュールを開発し、LRAやLFAのような追い出しポリシーを用いてメモリ使用を最適化し、さまざまなタイプの言語モデルに適応できるようにした。Attendreレイヤーは、モデルが過去と未来の文脈に関わることを可能にする重要なメカニズムであり、理解力と応答力を向上させる。

私たちの実験は、これらの新しいメモリ戦略が、長い入力シーケンスを理解するタスクでのパフォーマンスを大幅に改善することを示している。次のステップは、より広範なタスクでのテストや、メモリシステムの最大限の可能性を引き出すためのさらなる微調整アプローチの探求を行うことだ。また、メモリをさらに圧縮して、さまざまなアプリケーションで重要な文脈情報に効率的にアクセスする方法を検討する予定だ。