When Context Bites: Rilevare il RAG Poisoning tramite il Collasso dell'Attenzione a Livello di Documento

D-SCAN è un framework di analisi per rilevare attacchi di avvelenamento dei documenti nei sistemi RAG (Retrieval-Augmented Generation). Raccoglie gli stati interni dell'LLM durante la generazione (probabilità dei token e pesi di attenzione), estrae feature multidimensionali e addestra classificatori per distinguere tra documenti recuperati puliti e avvelenati.

1. Installazione e Requisiti

Prerequisiti

• Python >= 3.9
• GPU compatibile con CUDA (>= 24GB VRAM consigliati per l'inferenza di un modello con 8B parametri)

Preparazione del Modello

Questo progetto utilizza Llama-3.1-8B-Instruct per impostazione predefinita. Scarica il modello in un percorso locale e aggiorna la variabile MODEL_ID in collect_inner_state.py:

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

---

2. Avvio Rapido

Il flusso di lavoro completo consiste in tre passaggi: Raccogliere gli Stati Interni → Calcolare le Feature → Addestrare il Classificatore. La domanda e i documenti recuperati correlati sono forniti in https://huggingface.co/datasets/An998/D-SCAN.

Passaggio 1: Raccogliere gli Stati Interni del Modello

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

Passaggio 2: Calcolare le Metriche delle Feature

# 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

Passaggio 3: Addestrare il Classificatore e Valutare

Apri fit_D-SCAN.ipynb ed esegui le celle in sequenza per:

1. Caricare i dati delle feature generati nel Passaggio 2
2. Addestrare classificatori Logistic Regression / Random Forest (5-Fold CV)
3. Valutare le prestazioni di trasferimento sul set di test

---

3. Dettagli della Pipeline

3.1 Raccolta Dati (collect_inner_state.py)

Per ogni input domanda-documento, l'LLM esegue una generazione multi-campione (predefinito: 10 campioni, temperature=1.0) e raccoglie i seguenti stati interni:

Campo	Tipo	Descrizione
outer_ppl_probs	List[Tensor]	Probabilità di generazione dei token per ciascun campione
inner_ppl_matrix	List[List[Tensor]]	Pesi di attenzione sulla sequenza di input a ogni passo di generazione (media sui layer)
doc_ranges	Dict[str, List[int]]	Intervallo di posizioni dei token per ciascun documento nella sequenza di input
generated_sequences	List[List[int]]	Sequenze di ID dei token generati per ciascun campione

Parametri di Configurazione Chiave

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'

Gli output vengono salvati come file data_{batch_id}_reppl.pt e file di statistiche results_stats_*.jsonl.

---

3.2 Calcolo delle Feature (compute_feature.py)

Estrae 10 categorie con oltre 100 feature dimensionali dagli stati interni raccolti. Il classificatore utilizza queste feature per determinare se esistono documenti avvelenati tra i documenti recuperati per una data query.

Panoramica delle Metriche

#	Categoria	Nome della Classe	# Feature	Idea Centrale
1	Statistiche di Probabilità di Generazione	PerplexityMetrics	8	I documenti avvelenati possono aumentare l'incertezza del modello durante la generazione, riflessa nei cambiamenti della distribuzione di probabilità
2	Entropia dell'Attenzione	AttentionEntropyMetrics	4	Entropia alta = attenzione dispersa = possibile informazione conflittuale; Entropia bassa = attenzione focalizzata
3	Concentrazione dell'Attenzione	AttentionConcentrationMetrics	8	Misura se l'attenzione è concentrata su pochi token tramite rapporto Top-K e coefficiente di Gini
4	Densità di Attenzione sui Documenti	DocumentAttentionDensityMetrics	6	Somma dell'attenzione divisa per la lunghezza del documento, eliminando il bias di lunghezza nell'allocazione dell'attenzione
5	Coerenza Multi-Campione	SampleConsistencyMetrics	4	Sotto avvelenamento, i pattern di attenzione tra i campioni possono essere incoerenti (similarità coseno, divergenza JS)
6	Dinamica dell'Attenzione	AttentionDynamicsMetrics	4	Frequenza dei cambi di documento dominante e ampiezza della variazione di entropia durante la generazione
7	Fluttuazione dell'Attenzione a Livello di Token	TokenLevelAttentionMetrics	16	Stabilità dell'attenzione a ogni posizione di token di input attraverso i passi di generazione (std, entropia)
8	Statistiche Approfondite sulla Probabilità della Risposta	AnswerProbabilityMetrics	18	Quantili di probabilità, rapporti di token ad alta/bassa probabilità, log probabilità, perplessità, ecc.
9	Dinamica della Probabilità	ProbabilityDynamicsMetrics	14	Pendenza del trend, autocorrelazione, volatilità, rapporto di picchi nella sequenza di generazione
10	Coerenza di Probabilità tra Campioni	CrossSampleProbabilityConsistencyMetrics	13	Similarità coseno, correlazione di Pearson, MSE, indice di divergenza tra i campioni

Descrizioni Dettagliate delle Metriche

1. PerplexityMetrics (Statistiche di Probabilità di Generazione)

Calcolate da outer_ppl_probs (probabilità di ciascun token generato):

• ppl_mean_prob: Probabilità media su tutti i token generati
• ppl_std_prob: Deviazione standard della probabilità
• ppl_min_prob: Probabilità minima (incertezza estrema)
• ppl_low_prob_ratio: Rapporto di token a bassa probabilità (<0.1)
• ppl_cross_sample_var: Varianza delle probabilità medie tra i campioni
• ppl_coef_variation: Coefficiente di variazione (CV = std/mean)
• ppl_skewness: Asimmetria della distribuzione di probabilità
• ppl_kurtosis: Curtosi della distribuzione di probabilità

2. AttentionEntropyMetrics (Entropia dell'Attenzione)

Calcolate da inner_ppl_matrix (pesi di attenzione a ogni passo):

• attn_entropy_mean/std/max: Media, deviazione standard e massimo dell'entropia della distribuzione di attenzione
• attn_entropy_cv: Coefficiente di variazione dell'entropia dell'attenzione

3. AttentionConcentrationMetrics (Concentrazione dell'Attenzione)

• attn_top5/10/20_ratio_mean/std: Quota di attenzione catturata dai token Top-K%
• attn_gini_mean/std: Coefficiente di Gini della distribuzione di attenzione (misura di disuguaglianza)

4. DocumentAttentionDensityMetrics (Densità di Attenzione sui Documenti)