When Context Bites: Detectando Envenenamento de RAG via Colapso de Atenção em Nível de Documento

D-SCAN é um framework de análise para detectar ataques de envenenamento de documentos em sistemas RAG (Retrieval-Augmented Generation). Ele coleta estados internos do LLM durante a geração (probabilidades de tokens e pesos de atenção), extrai características multidimensionais e treina classificadores para distinguir entre documentos recuperados limpos e envenenados.

1. Instalação e Requisitos

Pré-requisitos

• Python >= 3.9
• GPU compatível com CUDA (>= 24GB de VRAM recomendado para inferência de modelo com 8B parâmetros)

Preparação do Modelo

Este projeto usa Llama-3.1-8B-Instruct por padrão. Baixe o modelo para um caminho local e atualize a variável MODEL_ID em collect_inner_state.py:

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

2. Início Rápido

O fluxo de trabalho completo consiste em três etapas: Coletar Estados Internos → Calcular Características → Treinar Classificador. A pergunta e os documentos recuperados relacionados são fornecidos em https://huggingface.co/datasets/An998/D-SCAN.

Etapa 1: Coletar Estados Internos do Modelo

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

Etapa 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

Etapa 3: Treinar Classificador e Avaliar

Abra fit_D-SCAN.ipynb e execute as células sequencialmente para:

1. Carregar os dados de características gerados na Etapa 2
2. Treinar classificadores de Regressão Logística / Random Forest (5-Fold CV)
3. Avaliar o desempenho de transferência no conjunto de teste

3. Detalhes do Pipeline

3.1 Coleta de Dados (collect_inner_state.py)

Para cada entrada de pergunta-documento, o LLM realiza geração multi-amostra (padrão: 10 amostras, temperature=1.0) e coleta os seguintes estados internos:

Parâmetros de Configuração Principais

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'

As saídas são salvas como arquivos data_{batch_id}_reppl.pt e arquivos de estatísticas results_stats_*.jsonl.

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

Extrai 10 categorias com mais de 100 características dimensionais dos estados internos coletados. O classificador usa essas características para determinar se existem documentos envenenados entre os documentos recuperados para uma determinada consulta.

Visão Geral das Métricas

Descrições Detalhadas das Métricas

1. PerplexityMetrics (Estatísticas de Probabilidade de Geração)

Calculado a partir de outer_ppl_probs (probabilidade de cada token gerado):

• ppl_mean_prob: Probabilidade média em todos os tokens gerados
• ppl_std_prob: Desvio padrão da probabilidade
• ppl_min_prob: Probabilidade mínima (incerteza extrema)
• ppl_low_prob_ratio: Proporção de tokens de baixa probabilidade (<0.1)
• ppl_cross_sample_var: Variância das probabilidades médias entre amostras
• ppl_coef_variation: Coeficiente de variação (CV = std/mean)
• ppl_skewness: Assimetria da distribuição de probabilidade
• ppl_kurtosis: Curtose da distribuição de probabilidade

2. AttentionEntropyMetrics (Entropia de Atenção)

Calculado a partir de inner_ppl_matrix (pesos de atenção em cada passo):

• attn_entropy_mean/std/max: Média, desvio padrão e máximo da entropia da distribuição de atenção
• attn_entropy_cv: Coeficiente de variação da entropia de atenção

3. AttentionConcentrationMetrics (Concentração de Atenção)

• attn_top5/10/20_ratio_mean/std: Fração de atenção capturada pelos tokens Top-K%
• attn_gini_mean/std: Coeficiente de Gini da distribuição de atenção (medida de desigualdade)

4. DocumentAttentionDensityMetrics (Densidade de Atenção em Documentos)

• doc_attn_dens_std/range/max/min: Desvio padrão, amplitude, máximo e mínimo da densidade de atenção entre documentos
• doc_attn_dens_entropy: Entropia da distribuição de densidade de atenção em documentos
• doc_attn_dens_temporal_var_mean: Média da variância temporal da densidade de atenção em documentos

5. SampleConsistencyMetrics (Consistência Multi-Amostra)

• sample_attn_consistency/std: Similaridade de cosseno entre amostras da atenção em nível de token
• sample_doc_consistency: Similaridade de cosseno entre amostras da atenção em nível de documento
• sample_doc_js_divergence: Divergência JS entre amostras da atenção em nível de documento

6. AttentionDynamicsMetrics (Dinâmica de Atenção)

• attn_doc_switch_mean/max: Contagem de trocas de documento dominante
• attn_entropy_change_mean/std: Mudança passo a passo da entropia de atenção

7. TokenLevelAttentionMetrics (Flutuação de Atenção em Nível de Token)

• tla_std_mean/std/max/median/p90/p99/high_ratio/cv: Estatísticas do desvio padrão de atenção por token ao longo dos passos de geração
• tla_ent_mean/std/max/median/p90/p99/high_ratio/cv: Estatísticas da entropia de atenção por token ao longo dos passos de geração

8. AnswerProbabilityMetrics (Estatísticas Profundas de Probabilidade de Resposta)

• aprob_p10/p25/p50/p75/p90/iqr: Quantis de probabilidade
• aprob_high_ratio_05/08: Proporção de tokens de alta probabilidade
• aprob_low_ratio_01/001: Proporção de tokens de baixa probabilidade
• aprob_log_mean/std/min: Estatísticas de log de probabilidade
• aprob_ppl_mean/std/max: Perplexidade em nível de sequência
• aprob_geometric_mean: Média geométrica das probabilidades
• aprob_distribution_entropy: Entropia de informação do histograma de probabilidade

9. ProbabilityDynamicsMetrics (Dinâmica de Probabilidade)

• pdyn_diff_mean/abs_diff_mean/abs_diff_std/abs_diff_max: Estatísticas de diferença de probabilidade
• pdyn_max_drop/max_jump: Queda/salto máximo em um único passo
• pdyn_volatility_mean/std: Volatilidade
• pdyn_trend_slope_mean/std: Inclinação de tendência linear
• pdyn_autocorr_mean/std: Coeficiente de autocorrelação
• pdyn_spike_ratio_01/03: Proporção de picos (pontos de mutação)

10. CrossSampleProbabilityConsistencyMetrics (Consistência de Probabilidade entre Amostras)

• cspc_mean_prob_std/cv/range: Consistência das probabilidades médias entre amostras
• cspc_ppl_std/cv/range: Consistência da perplexidade entre amostras
• cspc_min_prob_std/range: Consistência das probabilidades mínimas
• cspc_seq_cosine_mean/std: Similaridade de cosseno entre amostras das sequências de probabilidade
• cspc_seq_pearson_mean: Correlação de Pearson entre amostras das sequências de probabilidade
• cspc_seq_mse_mean: MSE entre amostras das sequências de probabilidade
• cspc_divergence_index: Índice de divergência entre amostras

Argumentos de Linha de Comando do 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

Arquivos de Saída

3.3 Treinamento do Classificador (fit_D-SCAN.ipynb)

Fluxo de trabalho do notebook:

1. Carregar Dados: Ler a saída full_analysis_results.json de compute_feature.py
2. Seleção de Características: Usar todas as características ou filtrar por categoria (ex.: apenas características de atenção)
3. Treinar Classificadores:
   • Regressão Logística (regularização L2): Classificador linear, adequado para menos características
   • Random Forest (100 árvores, max_depth=5): Classificador não linear, captura interações entre características
4. Validação Cruzada 5-Fold: Reporta AUC, Acurácia, Precisão, Recall, F1
5. Análise de Importância das Características: Exibe as 10 características mais importantes
6. Avaliação no Conjunto de Teste: Usa o scaler e classificador de treinamento para avaliar o desempenho de transferência em um conjunto de teste independente
7. Visualização: Curvas ROC, gráficos de barras de importância das características, comparação Treino vs Teste

Seleção de Características para Classificadores

A variável use_features no notebook oferece controle flexível sobre o subconjunto de características usado pelos classificadores:

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

O desempenho do classificador por categoria também é impresso durante a execução de compute_feature.py:

Citação

Se você achar nosso trabalho útil, considere citar nosso artigo:

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

@inproceedings{dscan,
author = {Ren, Yingtao and Zhao, Ziyi and Fu, Yiwei and Luo, Xiao and Chang, Yu-Cheng and Lin, Chin-Teng},
title = {When Context Bites: Detecting RAG Poisoning via Document-Level Attention Collapse},
year = {2026},
isbn = {9798400725999},
publisher = {Association for Computing Machinery},
address = {New York, NY, USA},
url = {https://doi.org/10.1145/3805712.3809904},
doi = {10.1145/3805712.3809904},
booktitle = {Proceedings of the 49th International ACM SIGIR Conference on Research and Development in Information Retrieval},
pages = {4080–4086},
numpages = {7},
keywords = {retrieval-augmented generation, information security, question answering system},
location = {Australia},
series = {SIGIR '26}
}