When Context Bites: Detecting RAG Poisoning via Document-Level Attention Collapse

D-SCAN is an analysis framework for detecting document poisoning attacks in RAG (Retrieval-Augmented Generation) systems. It collects LLM internal states during generation (token probabilities and attention weights), extracts multi-dimensional features, and trains classifiers to distinguish between clean and poisoned retrieved documents.

1. Installation & Requirements

Prerequisites

• Python >= 3.9
• CUDA-compatible GPU (>= 24GB VRAM recommended for 8B-parameter model inference)

Model Preparation

This project uses Llama-3.1-8B-Instruct by default. Download the model to a local path and update the MODEL_ID variable in collect_inner_state.py:

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

---

2. Quick Start

The full workflow consists of three steps: Collect Internal States → Compute Features → Train Classifier. The question and related retrieved documents are provided in https://huggingface.co/datasets/An998/D-SCAN.

Step 1: Collect Model Internal States

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

Step 2: Compute Feature Metrics

# 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

Step 3: Train Classifier and Evaluate

Open fit_D-SCAN.ipynb and execute cells sequentially to:

1. Load the feature data generated in Step 2
2. Train Logistic Regression / Random Forest classifiers (5-Fold CV)
3. Evaluate transfer performance on the test set

---

3. Pipeline Details

3.1 Data Collection (collect_inner_state.py)

For each question-document input, the LLM performs multi-sample generation (default: 10 samples, temperature=1.0) and collects the following internal states:

Field	Type	Description
outer_ppl_probs	List[Tensor]	Token generation probability for each sample
inner_ppl_matrix	List[List[Tensor]]	Attention weights over input sequence at each generation step (layer-averaged)
doc_ranges	Dict[str, List[int]]	Token position range for each document in the input sequence
generated_sequences	List[List[int]]	Generated token ID sequences for each sample

Key Configuration Parameters

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'

Outputs are saved as data_{batch_id}_reppl.pt files and results_stats_*.jsonl statistics files.

---

3.2 Feature Computation (compute_feature.py)

Extracts 10 categories with 100+ dimensional features from the collected internal states. The classifier uses these features to determine whether poisoned documents exist among the retrieved documents for a given query.

Metric Overview

#	Category	Class Name	# Features	Core Idea
1	Generation Probability Stats	PerplexityMetrics	8	Poisoned docs may increase model uncertainty during generation, reflected in probability distribution changes
2	Attention Entropy	AttentionEntropyMetrics	4	High entropy = dispersed attention = possible conflicting information; Low entropy = focused attention
3	Attention Concentration	AttentionConcentrationMetrics	8	Measures whether attention is concentrated on a few tokens via Top-K ratio and Gini coefficient
4	Document Attention Density	DocumentAttentionDensityMetrics	6	Attention sum divided by document length, eliminating length bias in attention allocation
5	Multi-Sample Consistency	SampleConsistencyMetrics	4	Under poisoning, attention patterns across samples may be inconsistent (cosine similarity, JS divergence)
6	Attention Dynamics	AttentionDynamicsMetrics	4	Frequency of dominant-document switches and entropy change magnitude during generation
7	Token-Level Attention Fluctuation	TokenLevelAttentionMetrics	16	Stability of attention at each input token position across generation steps (std, entropy)
8	Answer Probability Deep Stats	AnswerProbabilityMetrics	18	Probability quantiles, high/low probability token ratios, log probability, perplexity, etc.
9	Probability Dynamics	ProbabilityDynamicsMetrics	14	Trend slope, autocorrelation, volatility, spike ratio in the generation sequence
10	Cross-Sample Probability Consistency	CrossSampleProbabilityConsistencyMetrics	13	Cosine similarity, Pearson correlation, MSE, divergence index across samples

Detailed Metric Descriptions

1. PerplexityMetrics (Generation Probability Statistics)

Computed from outer_ppl_probs (probability of each generated token):

• ppl_mean_prob: Mean probability across all generated tokens
• ppl_std_prob: Probability standard deviation
• ppl_min_prob: Minimum probability (extreme uncertainty)
• ppl_low_prob_ratio: Ratio of low-probability tokens (<0.1)
• ppl_cross_sample_var: Variance of mean probabilities across samples
• ppl_coef_variation: Coefficient of variation (CV = std/mean)
• ppl_skewness: Probability distribution skewness
• ppl_kurtosis: Probability distribution kurtosis

2. AttentionEntropyMetrics (Attention Entropy)

Computed from inner_ppl_matrix (attention weights at each step):

• attn_entropy_mean/std/max: Mean, standard deviation, and maximum of attention distribution entropy
• attn_entropy_cv: Coefficient of variation of attention entropy

3. AttentionConcentrationMetrics (Attention Concentration)

• attn_top5/10/20_ratio_mean/std: Attention share captured by the Top-K% tokens
• attn_gini_mean/std: Gini coefficient of the attention distribution (inequality measure)

4. DocumentAttentionDensityMetrics (Document Attention Density)

• doc_attn_dens_std/range/max/min: Std, range, max, and min of attention density across documents
• doc_attn_dens_entropy: Entropy of document attention density distribution
• doc_attn_dens_temporal_var_mean: Mean temporal variance of document attention density

5. SampleConsistencyMetrics (Multi-Sample Consistency)

• sample_attn_consistency/std: Cross-sample cosine similarity of token-level attention
• sample_doc_consistency: Cross-sample cosine similarity of document-level attention
• sample_doc_js_divergence: Cross-sample JS divergence of document-level attention

6. AttentionDynamicsMetrics (Attention Dynamics)