[論文レビュー] LM-Infinite: Zero-Shot Extreme Length Generalization for Large Language Models

Chi Han, Qifan Wang+5 more|arXiv (Cornell University)|Aug 30, 2023

Topic Modeling被引用数 5

ひとこと要約

本稿では、大規模言語モデルにおけるゼロショット極端長スパン一般化のための単純でパラメータフリーな手法であるLM-Infiniteを提案する。推論時に$λ$-型のアテンションマスクと距離制限を適用することにより、微調整を一切行わずに、128kトークンにまで流暢で一貫性のあるテキスト生成が可能となり、最新のLLMと同等の性能を達成する。また、推論速度が2.72倍向上する。

ABSTRACT

Today's large language models (LLMs) typically train on short text segments (e.g., <4K tokens) due to the quadratic complexity of their Transformer architectures. As a result, their performance suffers drastically on inputs longer than those encountered during training, substantially limiting their applications in real-world tasks involving long contexts such as encoding scientific articles, code repositories, or long dialogues. Through theoretical analysis and empirical investigation, this work identifies three major factors contributing to this length generalization failure. Our theoretical analysis further reveals that commonly used techniques like truncating the attention window or relative positional encodings are inadequate to address them. Answering these challenges, we propose LM-Infinite, a simple and effective method for enhancing LLMs' capabilities of handling long contexts. LM-Infinite is highly flexible and can be used with most modern LLMs off-the-shelf. Without any parameter updates, it allows LLMs pre-trained with 2K or 4K-long segments to generalize to up to 200M length inputs while retaining perplexity. It also improves performance on downstream tasks such as Passkey Retrieval and Qasper in the zero-shot setting. LM-Infinite brings substantial efficiency improvements: it achieves 2.7x decoding speed up and 7.5x memory saving over the original model. Our codes are released at \url{https://github.com/Glaciohound/LM-Infinite}.

研究の動機と目的

• 訓練長さを超過するシーケンスを生成する際、大規模言語モデルにおける長さ一般化失敗という重要な課題に対処すること。
• 長文シーケンスで性能低下を引き起こす分布外（OOD）要因——未確認の距離、過剰なアテンション対象トークン数、および暗黙的に符号化された位置情報——を特定・診断すること。
• 既存の事前学習済みLLMがパラメータ更新や微調整なしに、極端な長さにまで流暢に一般化できる、プラグアンドプレイなソリューションを開発すること。
• 多様なLLMアーキテクチャおよび長文データセットにおいて、計算効率と一貫性のある生成品質を両立すること。

提案手法

• 各クエリがアテンションするトークン数を制限する$λ$-型アテンションマスクを導入し、計算負荷を低減するとともに、OODアテンションパターンの影響を軽減する。
• 相対位置埋め込みにおける最大相対距離を制限する距離制限を適用し、推論時に未確認の距離が発生するのを防ぐ。
• 事前学習モデルの再トレーニング、微調整、パラメータ更新を一切行わず、推論時のみに適用する。
• 長大なシーケンス全体に完全な自己アテンションを適用しないことで、$O(n)$の時間・空間計算量を維持し、効率的な推論を実現する。
• RoPE や Alibi などの既存の相対位置埋め込み方式とシームレスに統合され、多数のLLMに広く互換性を持つ。
• デコード中にリアルタイムで適用することで、アーキテクチャの変更なしに動的かつ柔軟に長文処理が可能になる。

Figure 1: Diagnosis of three OOD factors in LLMs.

実験結果

リサーチクエスチョン

• RQ1 大規模言語モデルが訓練長さを超えて生成する際、長さ一般化失敗を引き起こす主な分布外（OOD）要因は何か？
• RQ2 単純でパラメータフリーな手法が、微調整やモデル再トレーニングなしに、長文生成の流暢さと一貫性を向上させられるか？
• RQ3 どの程度まで$λ$-型アテンションマスクと距離制限が、128kトークンに達するシーケンスの生成品質を維持できるか？
• RQ4 本手法は、長文タスクにおける切り捨てベースラインや微調整済みモデルと比較して、効率性と性能でどの程度優れているか？
• RQ5 相対位置埋め込みを用いる多様なLLMアーキテクチャ（例：LLaMA、Llama-2、MPT、GPT-J）に、本手法は一般化可能か？

主な発見

• LM-Infiniteにより、LLaMAとLlama-2はArXivの32kトークンシーケンスにおいても流暢で一貫性のある生成を維持し、それぞれBLEUスコア19.7および18.4を達成。微調整済みモデルと同等の性能を示した。
• OpenWebText2では、Llama-2 + LM-Infiniteが16kトークンでBLEUスコア9.6を達成した一方、ベースラインLlama-2はゴミのような出力（BLEU = 0.0）を生成した。
• 32k長さのシーケンスにおいて、推論速度が2.72倍向上し、エンコーディング段階でも3.16倍の高速化を達成した。推論効率が顕著に向上した。
• 切り捨てベースラインと比較して、同程度の計算コストで約5のBLEUスコアの優位性を達成した。あるいは、同じBLEUスコアを達成するのに計算量の25％にまで削減可能だった。
• ArXivおよびOpenWebText2の両データセットで、128kトークンまで一貫性のある生成品質を維持した。ゼロショット一般化の強力な妥当性を示した。
• LLaMA、Llama-2、MPT-7B、GPT-Jなど、相対位置埋め込みを用いる多数の最先端LLMに普遍的に適用可能であり、モデル固有の調整を必要としない。

Figure 2: (a) \model is a plug-and-play solution for various LLMs, consisting of a $\Lambda$ -shaped mask and a distance constraint during attention. (b) We also provide a conceptual model for understanding how relative position encoding works.