Cuando el contexto muerde: Detección de envenenamiento de RAG mediante el colapso de atención a nivel de documento

D-SCAN es un marco de análisis para detectar ataques de envenenamiento de documentos en sistemas RAG (Retrieval-Augmented Generation). Recopila los estados internos del LLM durante la generación (probabilidades de tokens y pesos de atención), extrae características multidimensionales y entrena clasificadores para distinguir entre documentos recuperados limpios y envenenados.

1. Instalación y requisitos

Requisitos previos

• Python >= 3.9
• GPU compatible con CUDA (se recomiendan >= 24 GB de VRAM para la inferencia de modelos de 8B parámetros)

Preparación del modelo

Este proyecto usa Llama-3.1-8B-Instruct por defecto. Descarga el modelo a una ruta local y actualiza la variable MODEL_ID en collect_inner_state.py:

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

2. Inicio rápido

El flujo de trabajo completo consta de tres pasos: Recopilar estados internos → Calcular características → Entrenar el clasificador. Las preguntas y los documentos recuperados correspondientes se proporcionan en https://huggingface.co/datasets/An998/D-SCAN.

Paso 1: Recopilar los estados internos del modelo

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

Paso 2: Calcular las métricas de características

# 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

Paso 3: Entrenar el clasificador y evaluar

Abre fit_D-SCAN.ipynb y ejecuta las celdas en orden para:

1. Cargar los datos de características generados en el Paso 2
2. Entrenar clasificadores de Regresión Logística / Bosque Aleatorio (validación cruzada de 5 pliegues)
3. Evaluar el rendimiento de transferencia en el conjunto de prueba

3. Detalles del pipeline

3.1 Recopilación de datos (collect_inner_state.py)

Para cada entrada de pregunta-documento, el LLM realiza una generación de múltiples muestras (por defecto: 10 muestras, temperature=1.0) y recopila los siguientes estados internos:

Parámetros de configuración clave

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'

Las salidas se guardan como archivos data_{batch_id}_reppl.pt y archivos de estadísticas results_stats_*.jsonl.

3.2 Cálculo de características (compute_feature.py)

Extrae 10 categorías con más de 100 características dimensionales de los estados internos recopilados. El clasificador usa estas características para determinar si existen documentos envenenados entre los documentos recuperados para una consulta determinada.

Resumen de métricas

Descripciones detalladas de las métricas

1. PerplexityMetrics (Estadísticas de probabilidad de generación)

Calculadas a partir de outer_ppl_probs (probabilidad de cada token generado):

• ppl_mean_prob: Probabilidad media de todos los tokens generados
• ppl_std_prob: Desviación estándar de la probabilidad
• ppl_min_prob: Probabilidad mínima (incertidumbre extrema)
• ppl_low_prob_ratio: Proporción de tokens de baja probabilidad (<0.1)
• ppl_cross_sample_var: Varianza de las probabilidades medias entre muestras
• ppl_coef_variation: Coeficiente de variación (CV = std/mean)
• ppl_skewness: Asimetría de la distribución de probabilidad
• ppl_kurtosis: Curtosis de la distribución de probabilidad

2. AttentionEntropyMetrics (Entropía de atención)

Calculadas a partir de inner_ppl_matrix (pesos de atención en cada paso):

• attn_entropy_mean/std/max: Media, desviación estándar y máximo de la entropía de la distribución de atención
• attn_entropy_cv: Coeficiente de variación de la entropía de atención

3. AttentionConcentrationMetrics (Concentración de atención)

• attn_top5/10/20_ratio_mean/std: Proporción de atención capturada por los tokens del Top-K%
• attn_gini_mean/std: Coeficiente de Gini de la distribución de atención (medida de desigualdad)

4. DocumentAttentionDensityMetrics (Densidad de atención por documento)

• doc_attn_dens_std/range/max/min: Desviación estándar, rango, máximo y mínimo de la densidad de atención entre documentos
• doc_attn_dens_entropy: Entropía de la distribución de densidad de atención por documento
• doc_attn_dens_temporal_var_mean: Varianza temporal media de la densidad de atención por documento

5. SampleConsistencyMetrics (Consistencia entre muestras)

• sample_attn_consistency/std: Similitud coseno entre muestras de la atención a nivel de token
• sample_doc_consistency: Similitud coseno entre muestras de la atención a nivel de documento
• sample_doc_js_divergence: Divergencia JS entre muestras de la atención a nivel de documento

6. AttentionDynamicsMetrics (Dinámica de atención)

• attn_doc_switch_mean/max: Número de cambios de documento dominante
• attn_entropy_change_mean/std: Cambio de entropía de atención paso a paso

7. TokenLevelAttentionMetrics (Fluctuación de atención a nivel de token)

• tla_std_mean/std/max/median/p90/p99/high_ratio/cv: Estadísticos de la desviación estándar de atención por token a lo largo de los pasos de generación
• tla_ent_mean/std/max/median/p90/p99/high_ratio/cv: Estadísticos de la entropía de atención por token a lo largo de los pasos de generación

8. AnswerProbabilityMetrics (Estadísticas profundas de probabilidad de respuesta)

• aprob_p10/p25/p50/p75/p90/iqr: Cuantiles de probabilidad
• aprob_high_ratio_05/08: Proporción de tokens de alta probabilidad
• aprob_low_ratio_01/001: Proporción de tokens de baja probabilidad
• aprob_log_mean/std/min: Estadísticos de log-probabilidad
• aprob_ppl_mean/std/max: Perplejidad a nivel de secuencia
• aprob_geometric_mean: Media geométrica de las probabilidades
• aprob_distribution_entropy: Entropía de la información del histograma de probabilidad

9. ProbabilityDynamicsMetrics (Dinámica de probabilidad)

• pdyn_diff_mean/abs_diff_mean/abs_diff_std/abs_diff_max: Estadísticos de diferencias de probabilidad
• pdyn_max_drop/max_jump: Caída/salto máximo en un solo paso
• pdyn_volatility_mean/std: Volatilidad
• pdyn_trend_slope_mean/std: Pendiente de tendencia lineal
• pdyn_autocorr_mean/std: Coeficiente de autocorrelación
• pdyn_spike_ratio_01/03: Proporción de picos (puntos de mutación)

10. CrossSampleProbabilityConsistencyMetrics (Consistencia de probabilidad entre muestras)

• cspc_mean_prob_std/cv/range: Consistencia de las probabilidades medias entre muestras
• cspc_ppl_std/cv/range: Consistencia de la perplejidad entre muestras
• cspc_min_prob_std/range: Consistencia de las probabilidades mínimas
• cspc_seq_cosine_mean/std: Similitud coseno entre muestras de las secuencias de probabilidad
• cspc_seq_pearson_mean: Correlación de Pearson entre muestras de las secuencias de probabilidad
• cspc_seq_mse_mean: MSE entre muestras de las secuencias de probabilidad
• cspc_divergence_index: Índice de divergencia entre muestras

Argumentos de línea de comandos 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

Archivos de salida

3.3 Entrenamiento del clasificador (fit_D-SCAN.ipynb)

Flujo de trabajo del notebook:

1. Cargar datos: Leer la salida full_analysis_results.json de compute_feature.py
2. Selección de características: Usar todas las características o filtrar por categoría (p. ej., solo características de atención)
3. Entrenar clasificadores:
   • Regresión Logística (regularización L2): clasificador lineal, adecuado para menos características
   • Bosque Aleatorio (100 árboles, max_depth=5): clasificador no lineal, captura interacciones entre características
4. Validación cruzada de 5 pliegues: Reporta AUC, Accuracy, Precision, Recall, F1
5. Análisis de importancia de características: Muestra las 10 características más importantes
6. Evaluación en conjunto de prueba: Usa el escalador y el clasificador entrenados para evaluar el rendimiento de transferencia en un conjunto de prueba independiente
7. Visualización: Curvas ROC, gráficos de barras de importancia de características, comparación Entrenamiento vs Prueba

Selección de características para los clasificadores

La variable use_features del notebook proporciona un control flexible sobre el subconjunto de características utilizado por los clasificadores:

# 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_'))]

El rendimiento del clasificador por categoría también se imprime durante la ejecución de compute_feature.py: