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

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

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

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

• Python >= 3.9
• GPU с поддержкой CUDA (рекомендуется >= 24 ГБ видеопамяти для инференса модели с 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. Обучить классификаторы «Логистическая регрессия» / «Случайный лес» (5-кратная перекрёстная валидация)
3. Оценить переносимость на тестовом наборе

3. Детали конвейера

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

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

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

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 (Статистика вероятности генерации)

Вычисляется из 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: Стандартное отклонение, размах, максимум и минимум плотности внимания по документам
• doc_attn_dens_entropy: Энтропия распределения плотности внимания по документам
• doc_attn_dens_temporal_var_mean: Средняя временная дисперсия плотности внимания к документам

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

• sample_attn_consistency/std: Косинусное сходство между сэмплами для внимания на уровне токенов
• sample_doc_consistency: Косинусное сходство между сэмплами для внимания на уровне документов
• sample_doc_js_divergence: Расхождение JS между сэмплами для внимания на уровне документов

6. AttentionDynamicsMetrics (Динамика внимания)

• attn_doc_switch_mean/max: Количество переключений доминирующего документа
• attn_entropy_change_mean/std: Пошаговое изменение энтропии внимания

7. TokenLevelAttentionMetrics (Флуктуация внимания на уровне токенов)

• tla_std_mean/std/max/median/p90/p99/high_ratio/cv: Статистики стандартного отклонения внимания по токенам на протяжении шагов генерации
• tla_ent_mean/std/max/median/p90/p99/high_ratio/cv: Статистики энтропии внимания по токенам на протяжении шагов генерации

8. AnswerProbabilityMetrics (Глубокая статистика вероятности ответа)

• aprob_p10/p25/p50/p75/p90/iqr: Квантили вероятности
• aprob_high_ratio_05/08: Доля токенов с высокой вероятностью
• aprob_low_ratio_01/001: Доля токенов с низкой вероятностью
• aprob_log_mean/std/min: Статистики лог-вероятности
• aprob_ppl_mean/std/max: Перплексия на уровне последовательности
• aprob_geometric_mean: Геометрическое среднее вероятностей
• aprob_distribution_entropy: Информационная энтропия гистограммы вероятностей

9. ProbabilityDynamicsMetrics (Динамика вероятности)

• pdyn_diff_mean/abs_diff_mean/abs_diff_std/abs_diff_max: Статистики разностей вероятностей
• pdyn_max_drop/max_jump: Максимальное одношаговое падение/скачок
• pdyn_volatility_mean/std: Волатильность
• pdyn_trend_slope_mean/std: Наклон линейного тренда
• pdyn_autocorr_mean/std: Коэффициент автокорреляции
• pdyn_spike_ratio_01/03: Доля выбросов (точек резких изменений)

10. CrossSampleProbabilityConsistencyMetrics (Согласованность вероятностей между сэмплами)

• cspc_mean_prob_std/cv/range: Согласованность средних вероятностей между сэмплами
• cspc_ppl_std/cv/range: Согласованность перплексии между сэмплами
• cspc_min_prob_std/range: Согласованность минимальных вероятностей
• cspc_seq_cosine_mean/std: Косинусное сходство последовательностей вероятностей между сэмплами
• cspc_seq_pearson_mean: Корреляция Пирсона последовательностей вероятностей между сэмплами
• cspc_seq_mse_mean: Среднеквадратичная ошибка (MSE) последовательностей вероятностей между сэмплами
• cspc_divergence_index: Индекс расхождения между сэмплами

Аргументы командной строки 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

Выходные файлы

3.3 Обучение классификатора (fit_D-SCAN.ipynb)

Рабочий процесс в ноутбуке:

1. Загрузка данных: Чтение выходного файла full_analysis_results.json из compute_feature.py
2. Выбор признаков: Использование всех признаков или фильтрация по категории (например, только признаки внимания)
3. Обучение классификаторов:
   • Логистическая регрессия (L2-регуляризация): линейный классификатор, подходит для небольшого числа признаков
   • Случайный лес (100 деревьев, max_depth=5): нелинейный классификатор, учитывает взаимодействия признаков
4. 5-кратная перекрёстная валидация: Выдаёт AUC, Accuracy, Precision, Recall, F1
5. Анализ важности признаков: Выводит Top-10 важных признаков
6. Оценка на тестовом наборе: Использует масштабировщик и классификатор из обучения для оценки переносимости на независимом тестовом наборе
7. Визуализация: ROC-кривые, столбчатые диаграммы важности признаков, сравнение Train vs Test

Выбор признаков для классификаторов

Переменная use_features в ноутбуке обеспечивает гибкое управление подмножеством признаков, используемых классификаторами:

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

Производительность классификатора по каждой категории также выводится во время выполнения compute_feature.py: