動機

RAG（Retrieval-Augmented Generation）において、検索された文書をLLMに入力する前に関連性の高い文章だけを選別することは、性能とコストの両面で重要な課題です。

一般的には別途のRerankerモデル（Cross-Encoderなど）を使用しますが、一つ興味深い視点があります：

LLMがテキストを処理する際にすでに各トークンに対する「注意度（Attention）」を計算しているなら、このシグナルを直接活用できないだろうか？

このアイデアを実験で検証しました。

核心アイデア

TransformerベースのLLMは入力トークン間のAttention Scoreを計算します。この際、最後に生成するトークンがどの入力トークンに注目しているかを観察すれば、そのクエリに対してどの文章がより関連性が高いかを判断できます。

全体フロー


flowchart LR
    C[Context] --> LLM
    Q[Query] --> LLM
    A[Answer Prefix] --> LLM
    LLM -->|Attention| R[Reranking]

Attentionの観察方法

一般的なLLM推論では、最後のトークンのAttentionに基づいて次のトークンを生成します。この実験では次のトークンを生成する代わりに、そのAttention分布自体を文章関連性スコアとして活用します。


flowchart TB
    subgraph Context
        S1[文章 1]
        S2[文章 2]
        S3[文章 3]
    end
    Anchor[アンカートークン] -->|high| S1
    Anchor -->|low| S2
    Anchor -->|mid| S3
    S1 --> R1[Rank 1]
    S3 --> R2[Rank 2]
    S2 --> R3[Rank 3]

例えば、Contextに「美咲は東京大学の学生です」、「太郎は25歳です」、「美咲は太郎の隣に住んでいます」があり、Queryが「どこに行きましょうか？」の場合、Answer Prefixである「美咲：」（アンカートークン）のAttentionが「東京大学の学生」の文章に最も高く現れます。

核心はanswer_hint_prefix（上記の例では「美咲：」）がアンカートークンの役割を果たすという点です。このトークンのAttentionがContext中のどこに集中しているかを測定し、現在の状況に最も関連のある文章を見つけ出します。

関連論文

AttentionRAG (arXiv:2503.10720)

この実験の基盤となった論文です。主な貢献：

  • クエリをNext-Token Prediction形式に変換：「Where is Daniel?」→「Daniel is in the ____」
  • アンカートークン：空欄位置のトークンが意味的焦点を単一トークンに集中
  • 全レイヤー集約：浅い層（構文情報）＋深い層（意味情報）を全て合算
  • 結果：LLMLingua比約10%性能向上、最大6.3倍コンテキスト圧縮

In-Context Re-ranking (ICLR 2025)

  • LLMの生成（generation）なしにAttentionパターンの変化のみで文書を再ランキング
  • 較正（Calibration）：意味のないクエリ（「N/A」）で基準線を測定し位置バイアスを除去
  • O(1)順伝播パスでRankGPT比60%以上のレイテンシ削減

Contrastive Retrieval Heads (arXiv:2510.02219)

  • 全てのAttention Headが同等ではないという観察
  • 全ヘッドの1%未満のみで最先端リランカー性能を達成
  • 有用なヘッドは中間レイヤーに集中分布

実験設計

2つのモデルで同一の実験を実施しました：

モデルレイヤー数パラメータ
Gemma 3 4B IT344B
Qwen3 Reranker 4B364B

入力構成

messages = [
    {"role": "user", "content": f"Context: {context}\n\nQuestion: {question}"},
    {"role": "assistant", "content": answer_hint_prefix},  # 例：「美咲：」
]

answer_hint_prefixの最後のトークンがアンカーの役割を果たします。このトークンがContextのどの部分にAttentionを集中させるかを測定します。

処理手順

  1. 特殊トークン除去：チャットテンプレートの<end_of_turn>、<|im_end|>等を除去
  2. Attention Sink回避：Context領域のトークンのみを対象にスコア計算
  3. ノイズトークン除外：ピリオド（.）、改行（\n\n）等、意味なく高スコアを得るトークンを除去
  4. 文章別平均スコア計算：ピリオドを基準に文章を区分し、各文章の平均Attentionスコアを算出

核心コード

def aggregate_attention_scores(inputs, layer_numbers):
    with torch.no_grad():
        outputs = model(**inputs)
        attentions = outputs.attentions

    target_index = inputs['input_ids'].shape[1] - 1  # 最後のトークン（アンカー）

    per_layer_attentions = []
    for layer_num in layer_numbers:
        attention_matrix = attentions[layer_num].squeeze(0).mean(dim=0).cpu().float().numpy()
        focused_attention = attention_matrix[target_index, :]  # アンカーが注目する対象
        per_layer_attentions.append(focused_attention)

    return per_layer_attentions

実験シナリオ

以下のContextを固定し、様々な状況の質問を投げかけ、Attentionベースのrerankingが適切な文章を上位にランクするか確認しました：

美咲は23歳です。
美咲は東京大学の学生です。
太郎は地方大学の学生です。
太郎は25歳です。
美咲は健太と付き合っています。
美咲は太郎の隣に住んでいます。
健太は太郎と中学校の同級生です。
健太は大学生です。
美咲は健太に暴行されました。
太郎はアルバイトをして最近給料を受け取りました。
美咲は太郎の親戚です。
美咲は毎日運動する習慣があります。
美咲はお金を返さない習慣があります。

シナリオ1：「美咲が泣いています。なぜ泣いているの？」

→「美咲は健太に暴行されました」が上位にランクされるか確認

シナリオ2：「タクシーに乗りました。どこに行きましょうか？」

→「美咲は東京大学の学生です」が上位にランクされるか確認

シナリオ3：「お金を貸してほしいと言っています」

→「美咲はお金を返さない習慣があります」、「太郎は最近給料を受け取りました」が上位にランクされるか確認

観察結果

効果的だった点

  • 状況に合った文章が上位にランク：ほとんどのシナリオで直感的に関連性の高い文章が高スコアを獲得
  • 追加学習なしで動作：別途のRerankerモデル学習なしに、既存LLMのAttentionのみで動作
  • 全レイヤー合算が効果的：特定レイヤーよりも全レイヤーを合算した方がより安定した結果

限界と発見

  • 中間レイヤーの方が良い場合がある：CoRe-R論文の発見と一致し、全レイヤーを使用するよりも中間レイヤーのみを使用する方が良い場合がある
  • モデルサイズによる性能差：小さなモデルではAttentionシグナルの品質が低下
  • Attention Sink：最初のトークンや特殊トークンに異常に高いAttentionが集中する現象 → Context領域のみを対象にする必要あり
  • 改行トークンの高スコア：\n\n等が意味と無関係に高スコアを獲得 → 除外が必要
  • 文章区分：PoCレベルでピリオド基準で区分したが、多言語対応には<sep>トークンの活用が適切

モデル別の違い

特性Gemma 3 4BQwen3 Reranker 4B
特殊トークン処理<end_of_turn>除去<|im_end|>除去
改行トークン\n\nが1トークンとして扱われる\n単位
レイヤー数3436

Qwen3-Rerankerは名前から分かるようにRerankingに特化したモデルであり、Attentionシグナルの品質がより良いと期待しましたが、両モデルとも同様のレベルの結果を示しました。

結論と今後の方向

LLMのAttention Mapは別途の学習なしでも文書/文章の関連性を判断するのに活用できる有用なシグナルを含んでいます。

活用可能な方向：

  • RAGコンテキスト圧縮：長い検索結果から関連文章のみを抽出してLLMに渡す
  • ロールプレイ/対話システム：キャラクター設定（Context）中、現在の状況に合った情報のみを選別
  • 軽量Reranker：別途モデルなしに推論中に得られるAttentionを活用

改善課題：

  • 最適レイヤー組み合わせをモデル別に見つける方法（CoRe-R方式の対照的ヘッド選別）
  • ICR論文の較正（Calibration）技法を適用して位置バイアスを除去
  • 文章区分をピリオドではなくトークナイザーの特殊トークンで処理

参考資料