フォーカストランスフォーマー：コンテキストスケーリングのための対照的トレーニング

Focused Transformer: Contrastive Training for Context Scaling

July 6, 2023

著者: Szymon Tworkowski, Konrad Staniszewski, Mikołaj Pacek, Yuhuai Wu, Henryk Michalewski, Piotr Miłoś

cs.AI

要旨

大規模言語モデルは、新しい情報を文脈に応じて取り込む優れた能力を有しています。しかし、このアプローチの真の潜在能力は、有効な文脈長の制約によってしばしば制限されています。この問題に対する一つの解決策は、外部メモリにアクセスできるアテンションレイヤーを導入することです。この外部メモリは（キー、値）ペアで構成されています。しかし、ドキュメントの数が増えるにつれて、関連するキーと無関係なキーの比率が低下し、モデルが無関係なキーに注目しやすくなります。私たちは、異なる意味的価値に関連するキーが重複し、それらを区別することが難しくなるという「注意散漫問題」と呼ばれる重要な課題を特定しました。この問題に対処するため、私たちはFocused Transformer（FoT）という手法を導入しました。この手法は、コントラスティブ学習にインスパイアされたトレーニングプロセスを採用し、（キー、値）空間の構造を強化することで、文脈長の拡張を可能にします。私たちの方法は、既存の大規模モデルを微調整して、その有効な文脈を延長することを可能にします。これは、3Bおよび7BのOpenLLaMAチェックポイントを微調整することで実証されています。その結果得られたモデルをLongLLaMAと名付け、長い文脈を必要とするタスクにおいて進展を示しています。さらに、私たちのLongLLaMAモデルが256kの文脈長でのパスキー検索を巧みに管理できることを示しています。

English

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.