D-SCAN 是一个分析框架,用于检测 RAG(检索增强生成)系统中的文档投毒攻击。它在生成过程中收集 LLM 内部状态(token 概率和注意力权重),提取多维特征,并训练分类器来区分干净和投毒的检索文档。
本项目默认使用 Llama-3.1-8B-Instruct。将模型下载到本地路径,并更新 collect_inner_state.py 中的 MODEL_ID 变量:
MODEL_ID = "/your/path/to/Llama-3.1-8B-Instruct"
完整工作流包含三个步骤:收集内部状态 → 计算特征 → 训练分类器。 问题和相关检索文档提供于 https://huggingface.co/datasets/An998/D-SCAN。
# Update configuration parameters in collect_inner_state.py, then run
python collect_inner_state.py
# 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
打开 fit_D-SCAN.ipynb 并依次执行单元格以:
collect_inner_state.py)对于每个问题-文档输入,LLM 执行多样本生成(默认:10 个样本,temperature=1.0)并收集以下内部状态:
| 字段 | 类型 | 描述 |
|---|---|---|
outer_ppl_probs | List[Tensor] | 每个样本的 token 生成概率 |
inner_ppl_matrix | List[List[Tensor]] | 每个生成步骤对输入序列的注意力权重(层平均) |
doc_ranges | Dict[str, List[int]] | 输入序列中每个文档的 token 位置范围 |
generated_sequences | List[List[int]] | 每个样本生成的 token ID 序列 |
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 统计文件。
compute_feature.py)从收集的内部状态中提取 10 个类别、100+ 维特征。分类器使用这些特征来判断给定查询的检索文档中是否存在投毒文档。
| # | 类别 | 类名 | 特征数 | 核心思想 |
|---|---|---|---|---|
| 1 | 生成概率统计 | PerplexityMetrics | 8 | 投毒文档可能在生成过程中增加模型不确定性,反映在概率分布变化中 |
| 2 | 注意力熵 | AttentionEntropyMetrics | 4 | 高熵 = 注意力分散 = 可能存在冲突信息;低熵 = 注意力集中 |
| 3 | 注意力集中度 | AttentionConcentrationMetrics | 8 | 通过 Top-K 比率和基尼系数衡量注意力是否集中在少数 token 上 |
| 4 | 文档注意力密度 | DocumentAttentionDensityMetrics | 6 | 注意力总和除以文档长度,消除注意力分配中的长度偏差 |
| 5 | 多样本一致性 | SampleConsistencyMetrics | 4 | 在投毒情况下,各样本间的注意力模式可能不一致(余弦相似度、JS 散度) |
| 6 | 注意力动态 | AttentionDynamicsMetrics | 4 | 生成过程中主导文档切换的频率和熵变化幅度 |
| 7 | Token 级注意力波动 | TokenLevelAttentionMetrics | 16 | 每个输入 token 位置的注意力在各生成步骤间的稳定性(标准差、熵) |
| 8 | 答案概率深度统计 | AnswerProbabilityMetrics | 18 | 概率分位数、高/低概率 token 比率、对数概率、困惑度等 |
| 9 | 概率动态 | ProbabilityDynamicsMetrics | 14 | 生成序列中的趋势斜率、自相关、波动性、尖峰比率 |
| 10 | 跨样本概率一致性 | CrossSampleProbabilityConsistencyMetrics | 13 | 各样本间的余弦相似度、皮尔逊相关、MSE、散度指数 |
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:概率序列的跨样本皮尔逊相关cspc_seq_mse_mean:概率序列的跨样本 MSEcspc_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
| 参数 | 默认值 | 描述 |
|---|---|---|
--attack_dir | - | 攻击数据目录(包含 data_*_reppl.pt 文件) |
--clean_dir | - | 干净数据目录 |
--output_dir | - | 输出目录 |
--max_attack_samples | 3000 | 最大攻击样本数 |
--max_clean_samples | 3000 | 最大干净样本数 |
--min_correct_count | 0 | 干净数据中的最小正确答案数(用于过滤) |
--min_attack_target_count | 0 | 攻击数据中的最小目标答案命中数 |
--num_use_samples | None | 用于指标计算的每个问题的样本数(默认:全部) |
--model_path | None | 模型路径(用于加载 tokenizer 解码生成序列) |
| 文件名 | 描述 |
|---|---|
single_metric_analysis.json | 每个指标的 AUC、p 值、Cohen's d 等 |
full_analysis_results.json | 完整特征矩阵 + 标签 |
document_detailed_metrics.json | 每文档详细指标(JSON 摘要) |
document_detailed_metrics_full.pkl | 完整文档级指标(含步骤级注意力) |