[論文レビュー] Focused Transformer: Contrastive Training for Context Scaling

Szymon Tworkowski, Konrad Staniszewski+4 more|arXiv (Cornell University)|Jul 6, 2023

Topic Modeling被引用数 13

ひとこと要約

本論文は Focused Transformer（FoT）を紹介します。FoTはメモリ注意機構を備えた対比学習法で文脈長を拡張し、微調整済みの OpenLLaMA モデル（LongLLaMA）が最大256kトークンを処理できるようにし、長い文脈タスクを改善します。

ABSTRACT

Large language models have an exceptional capability to incorporate new information in a contextual manner. However, the full potential of such an approach is often restrained due to a limitation in the effective context length. One solution to this issue is to endow an attention layer with access to an external memory, which comprises of (key, value) pairs. Yet, as the number of documents increases, the proportion of relevant keys to irrelevant ones decreases, leading the model to focus more on the irrelevant keys. We identify a significant challenge, dubbed the distraction issue, where keys linked to different semantic values might overlap, making them hard to distinguish. To tackle this problem, we introduce the Focused Transformer (FoT), a technique that employs a training process inspired by contrastive learning. This novel approach enhances the structure of the (key, value) space, enabling an extension of the context length. Our method allows for fine-tuning pre-existing, large-scale models to lengthen their effective context. This is demonstrated by our fine-tuning of $3B$ and $7B$ OpenLLaMA checkpoints. The resulting models, which we name LongLLaMA, exhibit advancements in tasks requiring a long context. We further illustrate that our LongLLaMA models adeptly manage a $256 k$ context length for passkey retrieval.

研究の動機と目的

• LLMのマルチドキュメント注意における文脈スケーリングを制限する分散注意（distraction）問題を特定する。
• 鍵・値空間を形成することで、スクラッチからの再訓練を行わずに文脈長を拡張するFoTを開発する。
• 既存モデル（OpenLLaMA）を現実的に微調整してLongLLaMAを作成し、より長い文脈を実現することを示す。
• パスキー検索、QA、Few-shotタスクで長い文脈能力を経験的に評価し、訓練長を超えた外挿を示す。

提案手法

• 推論時にkNNを介して追加の（key, value）文脈を取得するメモリ注意層を追加する。
• 対照学習に触発されたクロスバッチ手順で訓練し、メモリ注意を正例（現在/前の文脈）と負例（他の文書）にさらす。
• メモリ層の位置エンコーディングを（LongLLaMAを除き）削除して外挿を可能にする；厳密なkNN検索にはFAISSを用いる。
• FoTでOpenLLaMA-3BおよびOpenLLaMA-7Bを微調整してLongLLaMAモデルを作成する；パスキー検索、Qasper、TREC、WebQS、および長文脈データセット（PG-19、arXiv、GitHub、Isabelle）で評価する。
• FoTが鍵値空間の構造を改善し、全スケールの再訓練なしで拡張文脈への注意を可能にすることを示す。
• FoTベースのモデルが訓練長を超えて文脈長を外挿できることを示す（例：パスキー検索で256kトークン）。

実験結果

リサーチクエスチョン

• RQ1 FoTは多数の文書にわたるトランスフォーマの文脈拡張時の分散注意の問題を緩和できますか？
• RQ2 FoTはそれらの長さでの訓練なしで、はるかに長い文脈への外挿を可能にしますか？
• RQ3 FoTで訓練したモデルは、長い文脈で困惑度を維持または改善し、Few-shot性能はどうですか？
• RQ4 パスキー検索や文書QAなどの長文脈タスクでの実用的な利得は何ですか？
• RQ5 FoTは異なるモデルサイズやデータセット（PG-19、arXiv、GitHub、Isabelle、Qasper、TREC、WebQS）でどのように機能しますか？

主な発見

• FoTはクロスバッチ対照学習を介して鍵値空間を形成することで、分散注意の問題を緩和します。
• LongLLaMA-3BおよびLongLLaMA-7Bは256kトークンまで外挿し、パスキー検索で100kで94.5%、256kで73%の精度を達成します。
• 長い文脈を用いた場合、LongLLaMAモデルはTRECとWebQSでベースラインと比較して精度が向上します。
• FoTは標準的なファインチューニングベースラインとは異なり、訓練文脈長を超えた外挿を可能にします。
• PG-19、arXiv、GitHub、IsabelleでFoTは長い文脈で困惑度が改善され、ベースラインと比べて最大64kの文脈利益を得ます。
• FoTは、アーキテクチャを変更することなく、既存モデルを微調整して文脈を拡張する、シンプルなプラグアンドプレイ拡張です。