当上下文反噬：基于文档级注意力坍缩的 RAG 投毒检测

D-SCAN 是一个用于检测 RAG（检索增强生成）系统中文档投毒攻击的分析框架。它在生成过程中收集 LLM 内部状态（token 概率和注意力权重），提取多维特征，并训练分类器以区分干净与投毒的检索文档。

1. 安装与要求

环境要求

• Python >= 3.9
• 支持 CUDA 的 GPU（建议显存 >= 24GB，用于 8B 参数模型推理）

模型准备

本项目默认使用 Llama-3.1-8B-Instruct。将模型下载到本地路径，并更新 collect_inner_state.py 中的 MODEL_ID 变量：

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

2. 快速开始

完整流程包含三个步骤：收集内部状态 → 计算特征 → 训练分类器。 问题及相关检索文档可在 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（每个生成 token 的概率）计算：

• ppl_mean_prob：所有生成 token 的平均概率
• ppl_std_prob：概率标准差
• ppl_min_prob：最小概率（极端不确定性）
• ppl_low_prob_ratio：低概率 token 的比例（<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% token 所占据的注意力份额
• 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：token 级注意力的跨样本余弦相似度
• sample_doc_consistency：文档级注意力的跨样本余弦相似度
• sample_doc_js_divergence：文档级注意力的跨样本 JS 散度

6. AttentionDynamicsMetrics（注意力动态）

• attn_doc_switch_mean/max：主导文档切换次数
• attn_entropy_change_mean/std：注意力熵的逐步骤变化

7. TokenLevelAttentionMetrics（Token 级注意力波动）

• tla_std_mean/std/max/median/p90/p99/high_ratio/cv：跨生成步骤的逐 token 注意力标准差统计
• tla_ent_mean/std/max/median/p90/p99/high_ratio/cv：跨生成步骤的逐 token 注意力熵统计

8. AnswerProbabilityMetrics（答案概率深度统计）

• aprob_p10/p25/p50/p75/p90/iqr：概率分位数
• aprob_high_ratio_05/08：高概率 token 的比例
• aprob_low_ratio_01/001：低概率 token 的比例
• 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：单步最大下降/跳升
• 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：概率序列的跨样本 Pearson 相关
• 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）

Notebook 工作流程：

1. 加载数据：读取 compute_feature.py 输出的 full_analysis_results.json
2. 特征选择：使用全部特征或按类别筛选（例如仅注意力特征）
3. 训练分类器：
   • 逻辑回归（L2 正则化）：线性分类器，适用于特征较少的情况
   • 随机森林（100 棵树，max_depth=5）：非线性分类器，能捕捉特征交互
4. 5 折交叉验证：报告 AUC、准确率、精确率、召回率、F1
5. 特征重要性分析：输出 Top-10 重要特征
6. 测试集评估：使用训练时的标准化器和分类器在独立测试集上评估迁移性能
7. 可视化：ROC 曲线、特征重要性柱状图、训练集与测试集对比

分类器的特征选择

Notebook 中的 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 执行期间还会打印每个类别分类器的性能：