컨텍스트가 물 때: 문서 수준 어텐션 붕괴를 통한 RAG 오염 탐지

D-SCAN은 RAG(Retrieval-Augmented Generation) 시스템에서 문서 오염 공격(document poisoning attack)을 탐지하기 위한 분석 프레임워크입니다. 생성 과정에서 LLM 내부 상태(토큰 확률 및 어텐션 가중치)를 수집하고, 다차원 특징을 추출하여 정상 문서와 오염된 검색 문서를 구분하는 분류기를 학습합니다.

1. 설치 및 요구 사항

사전 요구 사항

• Python >= 3.9
• CUDA 호환 GPU(8B 파라미터 모델 추론에는 24GB 이상의 VRAM 권장)

모델 준비

이 프로젝트는 기본적으로 Llama-3.1-8B-Instruct를 사용합니다. 모델을 로컬 경로에 다운로드하고 collect_inner_state.py의 MODEL_ID 변수를 업데이트하세요:

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-Fold CV)
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. 데이터 로드: compute_feature.py의 출력인 full_analysis_results.json 읽기
2. 특징 선택: 전체 특징 사용 또는 범주별 필터링(예: 어텐션 전용 특징)
3. 분류기 학습:
   • 로지스틱 회귀(L2 정규화): 선형 분류기, 특징 수가 적을 때 적합
   • 랜덤 포레스트(트리 100개, max_depth=5): 비선형 분류기, 특징 간 상호작용 포착
4. 5-Fold 교차 검증: AUC, Accuracy, Precision, Recall, F1 보고
5. 특징 중요도 분석: 상위 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 실행 중 범주별 분류기 성능도 출력됩니다: