コンテキストが噛みつくとき：文書レベルのアテンション崩壊によるRAGポイズニング検出

D-SCANは、RAG（検索拡張生成）システムにおける文書ポイズニング攻撃を検出するための分析フレームワークです。生成中のLLM内部状態（トークン確率とアテンション重み）を収集し、多次元の特徴量を抽出して、クリーンな取得文書とポイズニングされた取得文書を判別する分類器を訓練します。

1. インストールと要件

前提条件

• Python >= 3.9
• CUDA互換GPU（8Bパラメータモデルの推論には24GB以上のVRAMを推奨）

モデルの準備

このプロジェクトはデフォルトでLlama-3.1-8B-Instructを使用します。モデルをローカルパスにダウンロードし、collect_inner_state.py内のMODEL_ID変数を更新してください：

MODEL_ID = "/your/path/to/Llama-3.1-8B-Instruct"

2. クイックスタート

ワークフロー全体は3つのステップで構成されています：内部状態の収集 → 特徴量の計算 → 分類器の訓練。 質問と関連する取得文書は https://huggingface.co/datasets/An998/D-SCAN で提供されています。

ステップ1：モデルの内部状態を収集

# Update configuration parameters in collect_inner_state.py, then run
python collect_inner_state.py

ステップ2：特徴量メトリクスの計算

# Analyze both attack and clean data, compute all features, and save results
python compute_feature.py \
    --attack_dir ./saved_reppl_weights_serial_query1_attack2_2wiki \
    --clean_dir ./saved_reppl_weights_serial_query1_pure1_5_2wiki \
    --output_dir ./analysis_results/2wiki_all \
    --max_attack_samples 3000 \
    --max_clean_samples 3000

ステップ3：分類器の訓練と評価

fit_D-SCAN.ipynbを開き、セルを順番に実行して以下の作業を行います：

1. ステップ2で生成された特徴量データを読み込む
2. ロジスティック回帰／ランダムフォレスト分類器を訓練する（5分割交差検証）
3. テストセットでの転移性能を評価する

3. パイプラインの詳細

3.1 データ収集（collect_inner_state.py）

各質問-文書入力に対して、LLMはマルチサンプル生成（デフォルト：10サンプル、temperature=1.0）を実行し、以下の内部状態を収集します：

主要な設定パラメータ

MODEL_ID = "/path/to/model"     # Model path
NUM_SAMPLES = 10                # Number of samples per question
MAX_NEW_TOKENS = 50             # Maximum generated tokens
TEMPERATURE = 1.0               # Sampling temperature
data_type = 'pure1_5'           # Data type: 'pure1_5' (clean) or 'attack2' (attack)
dataset_name = 'hotpotqa'       # Dataset: 'hotpotqa', '2wiki', 'musique'

出力はdata_{batch_id}_reppl.ptファイルとresults_stats_*.jsonl統計ファイルとして保存されます。

3.2 特徴量の計算（compute_feature.py）

収集された内部状態から10カテゴリ・100次元以上の特徴量を抽出します。分類器はこれらの特徴量を使用して、特定のクエリに対する取得文書の中にポイズニングされた文書が存在するかどうかを判定します。

メトリクス概要

メトリクスの詳細説明

1. PerplexityMetrics（生成確率統計）

outer_ppl_probs（各生成トークンの確率）から計算されます：

• ppl_mean_prob: 全生成トークンにわたる平均確率
• ppl_std_prob: 確率の標準偏差
• ppl_min_prob: 最小確率（極端な不確実性）
• ppl_low_prob_ratio: 低確率トークン（<0.1）の比率
• ppl_cross_sample_var: サンプル間の平均確率の分散
• ppl_coef_variation: 変動係数（CV = std/mean）
• ppl_skewness: 確率分布の歪度
• ppl_kurtosis: 確率分布の尖度

2. AttentionEntropyMetrics（アテンションエントロピー）

inner_ppl_matrix（各ステップのアテンション重み）から計算されます：

• attn_entropy_mean/std/max: アテンション分布エントロピーの平均、標準偏差、最大値
• attn_entropy_cv: アテンションエントロピーの変動係数

3. AttentionConcentrationMetrics（アテンション集中度）

• attn_top5/10/20_ratio_mean/std: Top-K%トークンが占めるアテンションの割合
• attn_gini_mean/std: アテンション分布のジニ係数（不平等度の指標）

4. DocumentAttentionDensityMetrics（文書アテンション密度）

• doc_attn_dens_std/range/max/min: 文書間のアテンション密度の標準偏差、範囲、最大値、最小値
• doc_attn_dens_entropy: 文書アテンション密度分布のエントロピー
• doc_attn_dens_temporal_var_mean: 文書アテンション密度の時間的分散の平均

5. SampleConsistencyMetrics（マルチサンプル一貫性）

• sample_attn_consistency/std: トークンレベルアテンションのサンプル間コサイン類似度
• sample_doc_consistency: 文書レベルアテンションのサンプル間コサイン類似度
• sample_doc_js_divergence: 文書レベルアテンションのサンプル間JSダイバージェンス

6. AttentionDynamicsMetrics（アテンション動態）

• attn_doc_switch_mean/max: 優位文書の切り替え回数
• attn_entropy_change_mean/std: ステップごとのアテンションエントロピー変化

7. TokenLevelAttentionMetrics（トークンレベルのアテンション変動）

• tla_std_mean/std/max/median/p90/p99/high_ratio/cv: 生成ステップ全体にわたるトークンごとのアテンション標準偏差の統計量
• tla_ent_mean/std/max/median/p90/p99/high_ratio/cv: 生成ステップ全体にわたるトークンごとのアテンションエントロピーの統計量

8. AnswerProbabilityMetrics（回答確率の詳細統計）

• aprob_p10/p25/p50/p75/p90/iqr: 確率の分位数
• aprob_high_ratio_05/08: 高確率トークンの比率
• aprob_low_ratio_01/001: 低確率トークンの比率
• aprob_log_mean/std/min: 対数確率の統計量
• aprob_ppl_mean/std/max: シーケンスレベルのパープレキシティ
• aprob_geometric_mean: 確率の幾何平均
• aprob_distribution_entropy: 確率ヒストグラムの情報エントロピー

9. ProbabilityDynamicsMetrics（確率動態）

• pdyn_diff_mean/abs_diff_mean/abs_diff_std/abs_diff_max: 確率差分の統計量
• pdyn_max_drop/max_jump: 1ステップあたりの最大下落／跳躍
• pdyn_volatility_mean/std: ボラティリティ
• pdyn_trend_slope_mean/std: 線形トレンドの傾き
• pdyn_autocorr_mean/std: 自己相関係数
• pdyn_spike_ratio_01/03: スパイク（変異点）の比率

10. CrossSampleProbabilityConsistencyMetrics（サンプル間確率一貫性）

• cspc_mean_prob_std/cv/range: サンプル間の平均確率の一貫性
• cspc_ppl_std/cv/range: サンプル間のパープレキシティの一貫性
• cspc_min_prob_std/range: 最小確率の一貫性
• cspc_seq_cosine_mean/std: 確率シーケンスのサンプル間コサイン類似度
• cspc_seq_pearson_mean: 確率シーケンスのサンプル間ピアソン相関
• cspc_seq_mse_mean: 確率シーケンスのサンプル間MSE
• cspc_divergence_index: サンプル間ダイバージェンス指標

compute_feature.pyのコマンドライン引数

python compute_feature.py \
    --attack_dir <attack_data_directory> \
    --clean_dir <clean_data_directory> \
    --output_dir <output_directory> \
    --max_attack_samples 3000 \
    --max_clean_samples 3000 \
    --min_correct_count 0 \
    --min_attack_target_count 0 \
    --num_use_samples 10 \
    --model_path /path/to/model    # Optional: load tokenizer for accuracy calculation

出力ファイル

3.3 分類器の訓練（fit_D-SCAN.ipynb）

ノートブックのワークフロー：

1. データの読み込み：compute_feature.pyの出力であるfull_analysis_results.jsonを読み込む
2. 特徴量の選択：すべての特徴量を使用するか、カテゴリでフィルタリングする（例：アテンション関連のみ）
3. 分類器の訓練：
   • ロジスティック回帰（L2正則化）：線形分類器。特徴量が少ない場合に適する
   • ランダムフォレスト（100本の木、max_depth=5）：非線形分類器。特徴量間の相互作用を捉える
4. 5分割交差検証：AUC、Accuracy、Precision、Recall、F1を報告する
5. 特徴量重要度の分析：Top-10の重要特徴量を出力する
6. テストセット評価：訓練済みスケーラーと分類器を使用して、独立したテストセットでの転移性能を評価する
7. 可視化：ROC曲線、特徴量重要度の棒グラフ、学習 vs テストの比較

分類器のための特徴量選択

ノートブックのuse_features変数を使用すると、分類器が使用する特徴量サブセットを柔軟に制御できます：

# Use all features
use_features = [f for f in feature_names if f not in exclude_cols]

# Use only attention-related features
use_features = [f for f in use_features if 'attn' in f]

# Combine by metric category (example)
use_features = [f for f in feature_names if f.startswith(('ppl_', 'doc_', 'sample_'))]

カテゴリ別の分類器性能もcompute_feature.pyの実行中に表示されます：