Retrieval Head Mechanistically Explains Long-Context Factuality

要約

ロングコンテキスト言語モデルの最近の進歩にも関わらず、トランスフォーマーベースのモデルがどのようにしてロングコンテキスト内の任意の場所から関連情報を取得する機能を発揮するのかは依然として解明されていません。
この文書はこの疑問に答えることを目的としています。
幅広いモデルにわたる私たちの体系的な調査により、特別なタイプのアテンションヘッドが情報の検索に主に関与していることが明らかになり、これを検索ヘッドと呼んでいます。
我々は、検索ヘッドの興味深い特性を特定します。(1) ユニバーサル: ロングコンテキスト機能を備えた調査済みのすべてのモデルには、一連の検索ヘッドがあります。
(2) スパース: アテンション ヘッドのごく一部 (5\% 未満) のみが取得されます。
(3) 組み込み: 短いコンテキストで事前トレーニングされたモデルには検索ヘッドがすでに存在します。
継続的な事前トレーニングによってコンテキストの長さを拡張しても、情報取得を実行するのは依然として同じヘッドのセットです。
(4) 動的にアクティブ化: Llama-2 7B を例にとると、コンテキストがどのように変更されたとしても、12 個の検索ヘッドが常に必要な情報に注目します。
残りの検索ヘッドは、さまざまなコンテキストでアクティブ化されます。
(5) 因果関係: 検索ヘッドを完全に枝刈りすると、関連情報の検索に失敗し、幻覚が起こりますが、ランダムな非検索ヘッドを枝刈りしてもモデルの検索能力には影響しません。
さらに、検索ヘッドが思考連鎖 (CoT) 推論に強く影響し、モデルが質問と以前に生成されたコンテキストを頻繁に参照する必要があることを示します。
逆に、モデルがその固有の知識を使用して答えを直接生成するタスクは、検索ヘッドのマスクアウトによる影響が少なくなります。
これらの観察結果は、モデルの内部のどの部分が入力トークンから情報を求めているかをまとめて説明します。
私たちの洞察は、幻覚の軽減、推論の改善、KV キャッシュの圧縮に関する将来の研究を促進すると信じています。

要約(オリジナル)

Despite the recent progress in long-context language models, it remains elusive how transformer-based models exhibit the capability to retrieve relevant information from arbitrary locations within the long context. This paper aims to address this question. Our systematic investigation across a wide spectrum of models reveals that a special type of attention heads are largely responsible for retrieving information, which we dub retrieval heads. We identify intriguing properties of retrieval heads:(1) universal: all the explored models with long-context capability have a set of retrieval heads; (2) sparse: only a small portion (less than 5\%) of the attention heads are retrieval. (3) intrinsic: retrieval heads already exist in models pretrained with short context. When extending the context length by continual pretraining, it is still the same set of heads that perform information retrieval. (4) dynamically activated: take Llama-2 7B for example, 12 retrieval heads always attend to the required information no matter how the context is changed. The rest of the retrieval heads are activated in different contexts. (5) causal: completely pruning retrieval heads leads to failure in retrieving relevant information and results in hallucination, while pruning random non-retrieval heads does not affect the model’s retrieval ability. We further show that retrieval heads strongly influence chain-of-thought (CoT) reasoning, where the model needs to frequently refer back the question and previously-generated context. Conversely, tasks where the model directly generates the answer using its intrinsic knowledge are less impacted by masking out retrieval heads. These observations collectively explain which internal part of the model seeks information from the input tokens. We believe our insights will foster future research on reducing hallucination, improving reasoning, and compressing the KV cache.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google