Quand le contexte mord : Détecter l'empoisonnement RAG via l'effondrement de l'attention au niveau du document

D-SCAN est un framework d'analyse pour détecter les attaques d'empoisonnement de documents dans les systèmes RAG (Retrieval-Augmented Generation). Il collecte les états internes du LLM pendant la génération (probabilités de tokens et poids d'attention), extrait des caractéristiques multidimensionnelles, et entraîne des classifieurs pour distinguer les documents récupérés sains des documents empoisonnés.

1. Installation et prérequis

Prérequis

• Python >= 3.9
• GPU compatible CUDA (>= 24GB VRAM recommandé pour l'inférence d'un modèle à 8B paramètres)

Préparation du modèle

Ce projet utilise Llama-3.1-8B-Instruct par défaut. Téléchargez le modèle vers un chemin local et mettez à jour la variable MODEL_ID dans collect_inner_state.py :

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

---

2. Démarrage rapide

Le workflow complet se compose de trois étapes : Collecter les états internes → Calculer les caractéristiques → Entraîner le classifieur. La question et les documents récupérés associés sont fournis dans https://huggingface.co/datasets/An998/D-SCAN.

Étape 1 : Collecter les états internes du modèle

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

Étape 2 : Calculer les métriques de caractéristiques

# 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

Étape 3 : Entraîner le classifieur et évaluer

Ouvrez fit_D-SCAN.ipynb et exécutez les cellules séquentiellement pour :

1. Charger les données de caractéristiques générées à l'étape 2
2. Entraîner les classifieurs Logistic Regression / Random Forest (validation croisée 5-Fold)
3. Évaluer les performances de transfert sur l'ensemble de test

---

3. Détails du pipeline

3.1 Collecte de données (collect_inner_state.py)

Pour chaque entrée question-document, le LLM effectue une génération multi-échantillons (par défaut : 10 échantillons, temperature=1.0) et collecte les états internes suivants :

Champ	Type	Description
outer_ppl_probs	List[Tensor]	Probabilité de génération de token pour chaque échantillon
inner_ppl_matrix	List[List[Tensor]]	Poids d'attention sur la séquence d'entrée à chaque étape de génération (moyennés par couche)
doc_ranges	Dict[str, List[int]]	Plage de positions de tokens pour chaque document dans la séquence d'entrée
generated_sequences	List[List[int]]	Séquences d'IDs de tokens générés pour chaque échantillon

Paramètres de configuration clés

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'

Les sorties sont sauvegardées sous forme de fichiers data_{batch_id}_reppl.pt et de fichiers de statistiques results_stats_*.jsonl.

---

3.2 Calcul des caractéristiques (compute_feature.py)

Extrait 10 catégories avec plus de 100 dimensions de caractéristiques à partir des états internes collectés. Le classifieur utilise ces caractéristiques pour déterminer si des documents empoisonnés existent parmi les documents récupérés pour une requête donnée.

Aperçu des métriques

#	Catégorie	Nom de classe	Nb de caractéristiques	Idée principale
1	Statistiques de probabilité de génération	PerplexityMetrics	8	Les documents empoisonnés peuvent augmenter l'incertitude du modèle pendant la génération, reflétée par des changements dans la distribution de probabilité
2	Entropie d'attention	AttentionEntropyMetrics	4	Entropie élevée = attention dispersée = possible conflit d'informations ; Entropie faible = attention concentrée
3	Concentration d'attention	AttentionConcentrationMetrics	8	Mesure si l'attention est concentrée sur quelques tokens via le ratio Top-K et le coefficient de Gini
4	Densité d'attention par document	DocumentAttentionDensityMetrics	6	Somme d'attention divisée par la longueur du document, éliminant le biais de longueur dans l'allocation d'attention
5	Cohérence multi-échantillons	SampleConsistencyMetrics	4	Sous empoisonnement, les motifs d'attention entre échantillons peuvent être incohérents (similarité cosinus, divergence JS)
6	Dynamique d'attention	AttentionDynamicsMetrics	4	Fréquence des changements de document dominant et amplitude de variation d'entropie pendant la génération
7	Fluctuation d'attention au niveau des tokens	TokenLevelAttentionMetrics	16	Stabilité de l'attention à chaque position de token d'entrée à travers les étapes de génération (écart-type, entropie)
8	Statistiques approfondies de probabilité de réponse	AnswerProbabilityMetrics	18	Quantiles de probabilité, ratios de tokens à haute/faible probabilité, log probabilité, perplexité, etc.
9	Dynamique de probabilité	ProbabilityDynamicsMetrics	14	Pente de tendance, autocorrélation, volatilité, ratio de pics dans la séquence de génération
10	Cohérence de probabilité inter-échantillons	CrossSampleProbabilityConsistencyMetrics	13	Similarité cosinus, corrélation de Pearson, MSE, indice de divergence entre échantillons

Descriptions détaillées des métriques

1. PerplexityMetrics (Statistiques de probabilité de génération)

Calculées à partir de outer_ppl_probs (probabilité de chaque token généré) :

• ppl_mean_prob : Probabilité moyenne sur tous les tokens générés
• ppl_std_prob : Écart-type des probabilités
• ppl_min_prob : Probabilité minimale (incertitude extrême)
• ppl_low_prob_ratio : Ratio de tokens à faible probabilité (<0.1)
• ppl_cross_sample_var : Variance des probabilités moyennes entre échantillons
• ppl_coef_variation : Coefficient de variation (CV = std/mean)
• ppl_skewness : Asymétrie de la distribution de probabilité
• ppl_kurtosis : Kurtosis de la distribution de probabilité

2. AttentionEntropyMetrics (Entropie d'attention)

Calculées à partir de inner_ppl_matrix (poids d'attention à chaque étape) :

• attn_entropy_mean/std/max : Moyenne, écart-type et maximum de l'entropie de la distribution d'attention
• attn_entropy_cv : Coefficient de variation de l'entropie d'attention

3. AttentionConcentrationMetrics (Concentration d'attention)