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:

Campo	Tipo	Descrição
outer_ppl_probs	List[Tensor]	Probabilidade de geração de token para cada amostra
inner_ppl_matrix	List[List[Tensor]]	Pesos de atenção sobre a sequência de entrada em cada passo de geração (média por camada)
doc_ranges	Dict[str, List[int]]	Intervalo de posição de token para cada documento na sequência de entrada
generated_sequences	List[List[int]]	Sequências de IDs de tokens gerados para cada amostra

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

#	Categoria	Nome da Classe	# Características	Ideia Central
1	Estatísticas de Probabilidade de Geração	PerplexityMetrics	8	Documentos envenenados podem aumentar a incerteza do modelo durante a geração, refletida em mudanças na distribuição de probabilidade
2	Entropia de Atenção	AttentionEntropyMetrics	4	Alta entropia = atenção dispersa = possível informação conflitante; Baixa entropia = atenção focada
3	Concentração de Atenção	AttentionConcentrationMetrics	8	Mede se a atenção está concentrada em poucos tokens via proporção Top-K e coeficiente de Gini
4	Densidade de Atenção em Documentos	DocumentAttentionDensityMetrics	6	Soma de atenção dividida pelo comprimento do documento, eliminando viés de comprimento na alocação de atenção
5	Consistência Multi-Amostra	SampleConsistencyMetrics	4	Sob envenenamento, os padrões de atenção entre amostras podem ser inconsistentes (similaridade de cosseno, divergência JS)
6	Dinâmica de Atenção	AttentionDynamicsMetrics	4	Frequência de trocas de documento dominante e magnitude de mudança de entropia durante a geração
7	Flutuação de Atenção em Nível de Token	TokenLevelAttentionMetrics	16	Estabilidade da atenção em cada posição de token de entrada ao longo dos passos de geração (desvio padrão, entropia)
8	Estatísticas Profundas de Probabilidade de Resposta	AnswerProbabilityMetrics	18	Quantis de probabilidade, proporções de tokens de alta/baixa probabilidade, log de probabilidade, perplexidade, etc.
9	Dinâmica de Probabilidade	ProbabilityDynamicsMetrics	14	Inclinação de tendência, autocorrelação, volatilidade, proporção de picos na sequência de geração
10	Consistência de Probabilidade entre Amostras	CrossSampleProbabilityConsistencyMetrics	13	Similaridade de cosseno, correlação de Pearson, MSE, índice de divergência entre amostras

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)