When Context Bites: Erkennung von RAG-Vergiftung durch Aufmerksamkeitskollaps auf Dokumentebene

D-SCAN ist ein Analyseframework zur Erkennung von Dokumentvergiftungsangriffen in RAG-Systemen (Retrieval-Augmented Generation). Es sammelt interne LLM-Zustände während der Generierung (Token-Wahrscheinlichkeiten und Aufmerksamkeitsgewichte), extrahiert mehrdimensionale Merkmale und trainiert Klassifikatoren, um zwischen sauberen und vergifteten abgerufenen Dokumenten zu unterscheiden.

1. Installation & Voraussetzungen

Voraussetzungen

• Python >= 3.9
• CUDA-kompatible GPU (>= 24GB VRAM empfohlen für Inferenz mit 8B-Parameter-Modell)

Modellvorbereitung

Dieses Projekt verwendet standardmäßig Llama-3.1-8B-Instruct. Laden Sie das Modell in einen lokalen Pfad herunter und aktualisieren Sie die Variable MODEL_ID in collect_inner_state.py:

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

---

2. Schnellstart

Der vollständige Workflow besteht aus drei Schritten: Interne Zustände sammeln → Merkmale berechnen → Klassifikator trainieren. Die Frage und die zugehörigen abgerufenen Dokumente sind unter https://huggingface.co/datasets/An998/D-SCAN verfügbar.

Schritt 1: Interne Modellzustände sammeln

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

Schritt 2: Merkmalsmetriken berechnen

# 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

Schritt 3: Klassifikator trainieren und evaluieren

Öffnen Sie fit_D-SCAN.ipynb und führen Sie die Zellen nacheinander aus, um:

1. Die in Schritt 2 generierten Merkmalsdaten zu laden
2. Logistic Regression / Random Forest Klassifikatoren zu trainieren (5-Fold CV)
3. Die Transferleistung auf dem Testset zu evaluieren

---

3. Pipeline-Details

3.1 Datensammlung (collect_inner_state.py)

Für jede Frage-Dokument-Eingabe führt das LLM eine Multi-Sample-Generierung durch (Standard: 10 Samples, temperature=1.0) und sammelt die folgenden internen Zustände:

Feld	Typ	Beschreibung
outer_ppl_probs	List[Tensor]	Token-Generierungswahrscheinlichkeit für jedes Sample
inner_ppl_matrix	List[List[Tensor]]	Aufmerksamkeitsgewichte über die Eingabesequenz bei jedem Generierungsschritt (layer-gemittelt)
doc_ranges	Dict[str, List[int]]	Token-Positionsbereich für jedes Dokument in der Eingabesequenz
generated_sequences	List[List[int]]	Generierte Token-ID-Sequenzen für jedes Sample

Wichtige Konfigurationsparameter

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'

Ausgaben werden als data_{batch_id}_reppl.pt-Dateien und results_stats_*.jsonl-Statistikdateien gespeichert.

---

3.2 Merkmalsberechnung (compute_feature.py)

Extrahiert 10 Kategorien mit über 100-dimensionalen Merkmalen aus den gesammelten internen Zuständen. Der Klassifikator verwendet diese Merkmale, um zu bestimmen, ob vergiftete Dokumente unter den abgerufenen Dokumenten für eine gegebene Anfrage vorhanden sind.

Metrik-Übersicht

#	Kategorie	Klassenname	# Merkmale	Kernidee
1	Generierungswahrscheinlichkeits-Statistiken	PerplexityMetrics	8	Vergiftete Dokumente können die Modellunsicherheit während der Generierung erhöhen, was sich in Änderungen der Wahrscheinlichkeitsverteilung widerspiegelt
2	Aufmerksamkeitsentropie	AttentionEntropyMetrics	4	Hohe Entropie = verteilte Aufmerksamkeit = möglicherweise widersprüchliche Informationen; Niedrige Entropie = fokussierte Aufmerksamkeit
3	Aufmerksamkeitskonzentration	AttentionConcentrationMetrics	8	Misst, ob die Aufmerksamkeit über Top-K-Verhältnis und Gini-Koeffizient auf wenige Token konzentriert ist
4	Dokument-Aufmerksamkeitsdichte	DocumentAttentionDensityMetrics	6	Aufmerksamkeitssumme geteilt durch Dokumentlänge, wodurch Längenverzerrung bei der Aufmerksamkeitsverteilung eliminiert wird
5	Multi-Sample-Konsistenz	SampleConsistencyMetrics	4	Bei Vergiftung können Aufmerksamkeitsmuster über Samples hinweg inkonsistent sein (Kosinus-Ähnlichkeit, JS-Divergenz)
6	Aufmerksamkeitsdynamik	AttentionDynamicsMetrics	4	Häufigkeit von Wechseln des dominanten Dokuments und Ausmaß der Entropieänderung während der Generierung
7	Token-Level-Aufmerksamkeitsfluktuation	TokenLevelAttentionMetrics	16	Stabilität der Aufmerksamkeit an jeder Eingabe-Token-Position über Generierungsschritte hinweg (Std, Entropie)
8	Antwortwahrscheinlichkeits-Tiefenstatistiken	AnswerProbabilityMetrics	18	Wahrscheinlichkeitsquantile, Verhältnisse von Token mit hoher/niedriger Wahrscheinlichkeit, Log-Wahrscheinlichkeit, Perplexität usw.
9	Wahrscheinlichkeitsdynamik	ProbabilityDynamicsMetrics	14	Trendsteigung, Autokorrelation, Volatilität, Spike-Verhältnis in der Generierungssequenz
10	Cross-Sample-Wahrscheinlichkeitskonsistenz	CrossSampleProbabilityConsistencyMetrics	13	Kosinus-Ähnlichkeit, Pearson-Korrelation, MSE, Divergenzindex über Samples hinweg

Detaillierte Metrik-Beschreibungen

1. PerplexityMetrics (Generierungswahrscheinlichkeits-Statistiken)

Berechnet aus outer_ppl_probs (Wahrscheinlichkeit jedes generierten Tokens):

• ppl_mean_prob: Mittlere Wahrscheinlichkeit über alle generierten Token
• ppl_std_prob: Standardabweichung der Wahrscheinlichkeit
• ppl_min_prob: Minimale Wahrscheinlichkeit (extreme Unsicherheit)
• ppl_low_prob_ratio: Verhältnis von Token mit niedriger Wahrscheinlichkeit (<0.1)
• ppl_cross_sample_var: Varianz der mittleren Wahrscheinlichkeiten über Samples hinweg
• ppl_coef_variation: Variationskoeffizient (CV = std/mean)
• ppl_skewness: Schiefe der Wahrscheinlichkeitsverteilung
• ppl_kurtosis: Kurtosis der Wahrscheinlichkeitsverteilung

2. AttentionEntropyMetrics (Aufmerksamkeitsentropie)

Berechnet aus inner_ppl_matrix (Aufmerksamkeitsgewichte bei jedem Schritt):

• attn_entropy_mean/std/max: Mittelwert, Standardabweichung und Maximum der Entropie der Aufmerksamkeitsverteilung
• attn_entropy_cv: Variationskoeffizient der Aufmerksamkeitsentropie

3. AttentionConcentrationMetrics (Aufmerksamkeitskonzentration)

• attn_top5/10/20_ratio_mean/std: Anteil der Aufmerksamkeit, der von den Top-K%-Token erfasst wird
• attn_gini_mean/std: Gini-Koeffizient der Aufmerksamkeitsverteilung (Ungleichheitsmaß)