Когда контекст кусается: обнаружение отравления RAG через коллапс внимания на уровне документов

D-SCAN — это аналитический фреймворк для обнаружения атак отравления документов в системах RAG (Retrieval-Augmented Generation). Он собирает внутренние состояния LLM во время генерации (вероятности токенов и веса внимания), извлекает многомерные признаки и обучает классификаторы для различения чистых и отравленных извлечённых документов.

1. Установка и требования

Предварительные требования

• Python >= 3.9
• GPU с поддержкой CUDA (рекомендуется >= 24GB VRAM для инференса модели с 8B параметрами)

Подготовка модели

Этот проект по умолчанию использует Llama-3.1-8B-Instruct. Скачайте модель в локальный путь и обновите переменную MODEL_ID в collect_inner_state.py:

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

---

2. Быстрый старт

Полный рабочий процесс состоит из трёх шагов: Сбор внутренних состояний → Вычисление признаков → Обучение классификатора. Вопрос и связанные извлечённые документы предоставлены в https://huggingface.co/datasets/An998/D-SCAN.

Шаг 1: Сбор внутренних состояний модели

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

Шаг 2: Вычисление метрик признаков

# 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

Шаг 3: Обучение классификатора и оценка

Откройте fit_D-SCAN.ipynb и выполняйте ячейки последовательно, чтобы:

1. Загрузить данные признаков, сгенерированные на Шаге 2
2. Обучить классификаторы Logistic Regression / Random Forest (5-Fold CV)
3. Оценить производительность переноса на тестовом наборе

---

3. Детали пайплайна

3.1 Сбор данных (collect_inner_state.py)

Для каждого входного вопроса-документа LLM выполняет многосэмпловую генерацию (по умолчанию: 10 сэмплов, temperature=1.0) и собирает следующие внутренние состояния:

Поле	Тип	Описание
outer_ppl_probs	List[Tensor]	Вероятность генерации токена для каждого сэмпла
inner_ppl_matrix	List[List[Tensor]]	Веса внимания по входной последовательности на каждом шаге генерации (усреднённые по слоям)
doc_ranges	Dict[str, List[int]]	Диапазон позиций токенов для каждого документа во входной последовательности
generated_sequences	List[List[int]]	Сгенерированные последовательности ID токенов для каждого сэмпла

Ключевые параметры конфигурации

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'

Выходные данные сохраняются в виде файлов data_{batch_id}_reppl.pt и файлов статистики results_stats_*.jsonl.

---

3.2 Вычисление признаков (compute_feature.py)

Извлекает 10 категорий с признаками размерностью 100+ из собранных внутренних состояний. Классификатор использует эти признаки для определения наличия отравленных документов среди извлечённых документов для данного запроса.

Обзор метрик

#	Категория	Имя класса	Кол-во признаков	Основная идея
1	Статистика вероятности генерации	PerplexityMetrics	8	Отравленные документы могут увеличить неопределённость модели во время генерации, что отражается в изменениях распределения вероятностей
2	Энтропия внимания	AttentionEntropyMetrics	4	Высокая энтропия = рассеянное внимание = возможная конфликтующая информация; Низкая энтропия = сфокусированное внимание
3	Концентрация внимания	AttentionConcentrationMetrics	8	Измеряет, сконцентрировано ли внимание на нескольких токенах, через Top-K долю и коэффициент Джини
4	Плотность внимания к документу	DocumentAttentionDensityMetrics	6	Сумма внимания, делённая на длину документа, устраняет смещение по длине в распределении внимания
5	Согласованность по нескольким сэмплам	SampleConsistencyMetrics	4	При отравлении паттерны внимания по сэмплам могут быть несогласованными (косинусное сходство, JS-дивергенция)
6	Динамика внимания	AttentionDynamicsMetrics	4	Частота переключений доминирующего документа и величина изменения энтропии во время генерации
7	Флуктуация внимания на уровне токенов	TokenLevelAttentionMetrics	16	Стабильность внимания на каждой позиции входного токена по шагам генерации (std, энтропия)
8	Глубокая статистика вероятности ответа	AnswerProbabilityMetrics	18	Квантили вероятностей, доли токенов с высокой/низкой вероятностью, логарифмическая вероятность, перплексия и т.д.
9	Динамика вероятностей	ProbabilityDynamicsMetrics	14	Наклон тренда, автокорреляция, волатильность, доля всплесков в последовательности генерации
10	Согласованность вероятностей между сэмплами	CrossSampleProbabilityConsistencyMetrics	13	Косинусное сходство, корреляция Пирсона, MSE, индекс дивергенции между сэмплами

Подробные описания метрик

1. PerplexityMetrics (Статистика вероятности генерации)

Вычисляется из outer_ppl_probs (вероятность каждого сгенерированного токена):

• ppl_mean_prob: Средняя вероятность по всем сгенерированным токенам
• ppl_std_prob: Стандартное отклонение вероятности
• ppl_min_prob: Минимальная вероятность (экстремальная неопределённость)
• ppl_low_prob_ratio: Доля токенов с низкой вероятностью (<0.1)
• ppl_cross_sample_var: Дисперсия средних вероятностей между сэмплами
• ppl_coef_variation: Коэффициент вариации (CV = std/mean)
• ppl_skewness: Асимметрия распределения вероятностей
• ppl_kurtosis: Эксцесс распределения вероятностей

2. AttentionEntropyMetrics (Энтропия внимания)

Вычисляется из inner_ppl_matrix (веса внимания на каждом шаге):

• attn_entropy_mean/std/max: Среднее, стандартное отклонение и максимум энтропии распределения внимания
• attn_entropy_cv: Коэффициент вариации энтропии внимания

3. AttentionConcentrationMetrics (Концентрация внимания)

• attn_top5/10/20_ratio_mean/std: Доля внимания, приходящаяся на Top-K% токенов
• attn_gini_mean/std: Коэффициент Джини распределения внимания (мера неравенства)

4. DocumentAttentionDensityMetrics (Плотность внимания к документу)

• doc_attn_dens_std/range/max/min: Std, размах, максимум и минимум плотности внимания по документам
• doc_attn_dens_entropy: Энтропия распределения плотности внимания к документам
• doc_attn_dens_temporal_var_mean: Средняя временная дисперсия плотности внимания к документам

5. SampleConsistencyMetrics (Согласованность по нескольким сэмплам)