Quand le contexte mord : Détection de l'empoisonnement RAG par effondrement de l'attention au niveau document

D-SCAN est un framework d'analyse permettant de détecter les attaques par 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 (>= 24 Go de VRAM recommandé pour l'inférence de modèles à 8B de paramètres)

Préparation du modèle

Ce projet utilise Llama-3.1-8B-Instruct par défaut. Téléchargez le modèle dans 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 comprend trois étapes : Collecter les états internes → Calculer les caractéristiques → Entraîner le classifieur. Les questions et les documents récupérés associés sont fournis sur 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 Régression logistique / Forêt aléatoire (validation croisée à 5 plis)
3. Évaluer la performance 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 (défaut : 10 échantillons, temperature=1.0) et collecte les états internes suivants :

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 enregistré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 sont présents parmi les documents récupérés pour une requête donnée.

Aperçu des métriques

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 de la probabilité
• 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 les échantillons
• ppl_coef_variation : Coefficient de variation (CV = std/mean)
• ppl_skewness : Asymétrie de la distribution de probabilité
• ppl_kurtosis : Aplatissement de la distribution de probabilité

2. AttentionEntropyMetrics (Entropie de l'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 de l'attention

3. AttentionConcentrationMetrics (Concentration de l'attention)

• attn_top5/10/20_ratio_mean/std : Part d'attention capturée par les tokens Top-K %
• attn_gini_mean/std : Coefficient de Gini de la distribution d'attention (mesure d'inégalité)

4. DocumentAttentionDensityMetrics (Densité d'attention par document)

• doc_attn_dens_std/range/max/min : Écart-type, étendue, max et min de la densité d'attention entre documents
• doc_attn_dens_entropy : Entropie de la distribution de densité d'attention des documents
• doc_attn_dens_temporal_var_mean : Variance temporelle moyenne de la densité d'attention des documents

5. SampleConsistencyMetrics (Cohérence multi-échantillons)

• sample_attn_consistency/std : Similarité cosinus inter-échantillons de l'attention au niveau token
• sample_doc_consistency : Similarité cosinus inter-échantillons de l'attention au niveau document
• sample_doc_js_divergence : Divergence JS inter-échantillons de l'attention au niveau document

6. AttentionDynamicsMetrics (Dynamique de l'attention)

• attn_doc_switch_mean/max : Nombre de basculements du document dominant
• attn_entropy_change_mean/std : Changement d'entropie de l'attention par étape

7. TokenLevelAttentionMetrics (Fluctuation de l'attention au niveau token)

• tla_std_mean/std/max/median/p90/p99/high_ratio/cv : Statistiques de l'écart-type de l'attention par token au fil des étapes de génération
• tla_ent_mean/std/max/median/p90/p99/high_ratio/cv : Statistiques de l'entropie de l'attention par token au fil des étapes de génération

8. AnswerProbabilityMetrics (Statistiques approfondies de probabilité de réponse)

• aprob_p10/p25/p50/p75/p90/iqr : Quantiles de probabilité
• aprob_high_ratio_05/08 : Ratio de tokens à haute probabilité
• aprob_low_ratio_01/001 : Ratio de tokens à faible probabilité
• aprob_log_mean/std/min : Statistiques de log-probabilité
• aprob_ppl_mean/std/max : Perplexité au niveau de la séquence
• aprob_geometric_mean : Moyenne géométrique des probabilités
• aprob_distribution_entropy : Entropie informationnelle de l'histogramme de probabilité

9. ProbabilityDynamicsMetrics (Dynamique de probabilité)

• pdyn_diff_mean/abs_diff_mean/abs_diff_std/abs_diff_max : Statistiques des différences de probabilité
• pdyn_max_drop/max_jump : Baisse/saut maximal en une seule étape
• pdyn_volatility_mean/std : Volatilité
• pdyn_trend_slope_mean/std : Pente de tendance linéaire
• pdyn_autocorr_mean/std : Coefficient d'autocorrélation
• pdyn_spike_ratio_01/03 : Ratio de pics (points de mutation)

10. CrossSampleProbabilityConsistencyMetrics (Cohérence de probabilité inter-échantillons)

• cspc_mean_prob_std/cv/range : Cohérence des probabilités moyennes entre échantillons
• cspc_ppl_std/cv/range : Cohérence de la perplexité entre échantillons
• cspc_min_prob_std/range : Cohérence des probabilités minimales
• cspc_seq_cosine_mean/std : Similarité cosinus inter-échantillons des séquences de probabilité
• cspc_seq_pearson_mean : Corrélation de Pearson inter-échantillons des séquences de probabilité
• cspc_seq_mse_mean : MSE inter-échantillons des séquences de probabilité
• cspc_divergence_index : Indice de divergence inter-échantillons

Arguments de la ligne de commande de 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

Fichiers de sortie

3.3 Entraînement du classifieur (fit_D-SCAN.ipynb)

Workflow du notebook :

1. Charger les données : Lire la sortie full_analysis_results.json de compute_feature.py
2. Sélection des caractéristiques : Utiliser toutes les caractéristiques ou filtrer par catégorie (par exemple, caractéristiques d'attention uniquement)
3. Entraîner les classifieurs :
   • Régression logistique (régularisation L2) : classifieur linéaire, adapté à un nombre réduit de caractéristiques
   • Forêt aléatoire (100 arbres, max_depth=5) : classifieur non linéaire, capture les interactions entre caractéristiques
4. Validation croisée à 5 plis : Rapporte AUC, exactitude, précision, rappel, F1
5. Analyse de l'importance des caractéristiques : Affiche les 10 caractéristiques les plus importantes
6. Évaluation sur l'ensemble de test : Utilise le scaler et le classifieur entraînés pour évaluer la performance de transfert sur un ensemble de test indépendant
7. Visualisation : Courbes ROC, diagrammes en barres de l'importance des caractéristiques, comparaison Entraînement vs Test

Sélection des caractéristiques pour les classifieurs

La variable use_features du notebook permet de contrôler de manière flexible le sous-ensemble de caractéristiques utilisé par les classifieurs :

# Use all features
use_features = [f for f in feature_names if f not in exclude_cols]

# Use only attention-related features
use_features = [f for f in use_features if 'attn' in f]

# Combine by metric category (example)
use_features = [f for f in feature_names if f.startswith(('ppl_', 'doc_', 'sample_'))]

La performance du classifieur par catégorie est également affichée lors de l'exécution de compute_feature.py :