When Context Bites: Detecting RAG Poisoning via Document-Level Attention Collapse

D-SCAN es un marco de análisis para detectar ataques de envenenamiento de documentos en sistemas RAG (Retrieval-Augmented Generation). Recopila 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 recomienda >= 24GB de VRAM para la inferencia de modelos de 8B parámetros)

Preparación del Modelo

Este proyecto utiliza Llama-3.1-8B-Instruct por defecto. Descargue el modelo a una ruta local y actualice 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 Clasificador. La pregunta y los documentos recuperados relacionados se proporcionan en https://huggingface.co/datasets/An998/D-SCAN.

Paso 1: Recopilar Estados Internos del Modelo

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

Paso 2: Calcular 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

Abra fit_D-SCAN.ipynb y ejecute las celdas secuencialmente para:

1. Cargar los datos de características generados en el Paso 2
2. Entrenar clasificadores de Regresión Logística / Random Forest (5-Fold CV)
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 multi-muestra (por defecto: 10 muestras, temperature=1.0) y recopila los siguientes estados internos:

Campo	Tipo	Descripción
outer_ppl_probs	List[Tensor]	Probabilidad de generación de tokens para cada muestra
inner_ppl_matrix	List[List[Tensor]]	Pesos de atención sobre la secuencia de entrada en cada paso de generación (promediados por capa)
doc_ranges	Dict[str, List[int]]	Rango de posiciones de tokens para cada documento en la secuencia de entrada
generated_sequences	List[List[int]]	Secuencias de IDs de tokens generados para cada muestra

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 utiliza estas características para determinar si existen documentos envenenados entre los documentos recuperados para una consulta dada.

Resumen de Métricas

#	Categoría	Nombre de Clase	# Características	Idea Central
1	Estadísticas de Probabilidad de Generación	PerplexityMetrics	8	Los documentos envenenados pueden aumentar la incertidumbre del modelo durante la generación, reflejado en cambios en la distribución de probabilidad
2	Entropía de Atención	AttentionEntropyMetrics	4	Alta entropía = atención dispersa = posible información conflictiva; Baja entropía = atención enfocada
3	Concentración de Atención	AttentionConcentrationMetrics	8	Mide si la atención se concentra en unos pocos tokens mediante la proporción Top-K y el coeficiente de Gini
4	Densidad de Atención por Documento	DocumentAttentionDensityMetrics	6	Suma de atención dividida por la longitud del documento, eliminando el sesgo de longitud en la asignación de atención
5	Consistencia Multi-Muestra	SampleConsistencyMetrics	4	Bajo envenenamiento, los patrones de atención entre muestras pueden ser inconsistentes (similitud coseno, divergencia JS)
6	Dinámica de Atención	AttentionDynamicsMetrics	4	Frecuencia de cambios de documento dominante y magnitud del cambio de entropía durante la generación
7	Fluctuación de Atención a Nivel de Token	TokenLevelAttentionMetrics	16	Estabilidad de la atención en cada posición de token de entrada a través de los pasos de generación (std, entropía)
8	Estadísticas Profundas de Probabilidad de Respuesta	AnswerProbabilityMetrics	18	Cuantiles de probabilidad, proporciones de tokens de alta/baja probabilidad, log probabilidad, perplejidad, etc.
9	Dinámica de Probabilidad	ProbabilityDynamicsMetrics	14	Pendiente de tendencia, autocorrelación, volatilidad, proporción de picos en la secuencia de generación
10	Consistencia de Probabilidad entre Muestras	CrossSampleProbabilityConsistencyMetrics	13	Similitud coseno, correlación de Pearson, MSE, índice de divergencia entre muestras

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 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)