컨텍스트가 물어뜯을 때: 문서 수준 어텐션 붕괴를 통한 RAG 포이즈닝 탐지

D-SCAN은 RAG(검색 증강 생성) 시스템에서 문서 포이즈닝 공격을 탐지하기 위한 분석 프레임워크입니다. 생성 중 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)을 수행하고 다음 내부 상태를 수집합니다:

필드	타입	설명
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	생성 단계에 걸친 각 입력 토큰 위치에서의 어텐션 안정성 (표준편차, 엔트로피)
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: 문서 간 어텐션 밀도의 표준편차, 범위, 최댓값, 최솟값
• 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