PoisonCraft

이 저장소는 POISONCRAFT: Practical Poisoning of Retrieval-Augmented Generation for Large Language Models. 의 공식 구현을 제공합니다.

개요

POISONCRAFT는 악의적인 행위자가 RAG(Retrieval-Augmented Generation) 파이프라인에서 사용하는 코퍼스에 "오염된(poisoned)" 콘텐츠를 심어 대규모 언어 모델이 악성 콘텐츠를 할루시네이션하거나 참조하도록 유도할 수 있는 방법을 시연하는 것을 목표로 합니다. 이 코드베이스는 다음을 위한 스크립트를 제공합니다:

• 표준 데이터셋(예: Natural Questions, MS MARCO, HotpotQA)을 준비하고 전처리합니다.

• 쿼리 유사 텍스트에 적대적 접미사(즉, "poisons")를 주입하여 검색 결과를 저하시키거나 조작합니다.

• 다양한 리트리버(예: Contriever, SimCSE, BGE)에서 오염 공격의 효과를 평가하고 전이 가능성을 측정합니다.

빠른 사용법

환경 설정

다음은 환경을 구성하기 위한 개괄적인 가이드입니다. 정확한 요구 사항은 requirements.txt에 나와 있습니다. 충돌을 피하기 위해 새 가상 환경을 만드는 것을 권장합니다:

# conda 사용 예시
conda create -n PoisonCraft python=3.10
conda activate PoisonCraft

# 의존성 설치
pip install torch==2.4.1 torchvision==0.19.1 torchaudio==2.4.1 --index-url https://download.pytorch.org/whl/cu118
pip install -r requirements.txt

API 키 및 모델 구성 설정

도메인 분류, 쿼리 분류와 같은 이 프로젝트의 일부 부분은 API 호출(예: OpenAI GPT 모델)에 의존합니다. API 키와 모델 설정을 다음과 같이 구성해야 합니다:

1. src/models_configs/ 디렉토리로 이동합니다.
2. 적절한 구성 파일(예: gpt3.5_config.py)을 찾습니다.
3. api_keys 필드를 API 키로 업데이트합니다:
   root@kitploit:~

   "api_key_info": {
       "api_keys": ["your_openai_api_key"],
       "api_key_use": 0
   }
   

참고: 이 프로젝트는 여러 LLM 공급자(예: OpenAI, Llama, PaLM2)를 지원합니다. 모델에 맞는 적절한 구성 파일을 선택해야 합니다.

코드 실행 방법

파이프라인의 다양한 단계를 위해 scripts 폴더 아래에 다양한 bash 스크립트를 제공합니다. 다음은 단계별 참조입니다:

데이터셋 준비

1. 데이터셋 다운로드 및 준비

scripts/run_prepare_dataset.sh를 실행하여 데이터셋을 다운로드하고 압축을 풀고 학습/테스트 분할을 생성합니다:

bash scripts/run_prepare_dataset.sh

이 스크립트는 내부적으로 다음을 호출합니다:

• experiments/prepare_datasets.py
• experiments/classify_queries_by_domain.py (필요한 경우 도메인 분류용)

2. 데이터셋 전처리

각 도메인에 대한 top-k ground truth 유사도 점수를 계산하려면 다음을 실행합니다:

bash scripts/run_process_data.sh

3. [선택 사항] BEIR 기반 Ground Truth 평가

리트리버 모델을 평가하려는 경우에만 이 단계를 실행하세요. 상당한 시간이 걸릴 수 있습니다. 다음 단계를 따르세요:

1. BEIR 평가 스크립트를 실행하여 검색 결과를 계산합니다:
   root@kitploit:~

   python src/beir_eval/eval_beir.py --model_code bge-small --dataset nq --top_k 100
   

2. 검색된 결과를 사용하여 ground truth 점수를 생성합니다:
   root@kitploit:~

   bash scripts/run_process_data.sh
   

오염 공격 훈련

적대적 접미사(즉, "poisons")를 생성하려면 다음을 사용할 수 있습니다:

bash scripts/run_multi_poisoncraft.sh

• 이 스크립트는 run_poisoncraft.sh의 여러 실행을 병렬로 시작하며, 각 실행은 서로 다른 적대적 접미사 길이와 도메인 인덱스를 탐색합니다.
• 핵심 로직은 experiments/poisoncraft_exp.py를 호출합니다.

다음과 같은 매개변수를 사용자 정의할 수 있습니다:

• DOMAIN_LIST (오염시키려는 쿼리의 도메인 인덱스)
• ADV_LENGTH_LIST (테스트할 적대적 접미사 길이)

검색 공격 평가 실행

적대적 접미사를 생성한 후, 몇 개의 쿼리가 "오염된" 참조를 가져오는지 평가하려면 다음을 실행합니다:

bash scripts/run_retrieval_attack.sh

이 스크립트는:

• 사전 계산된 적대적 접미사를 로드합니다.
• experiments/retrieval_attack.py를 사용하여 공격 테스트를 실행합니다.
• 오염된 텍스트를 검색하게 된 쿼리 수에 대한 결과를 기록합니다.

타겟 생성 공격 평가 실행

또는 LLM 수준 평가(즉, top-k 검색 결과를 포함한 최종 텍스트 생성)를 위해 다음을 사용할 수 있습니다:

bash scripts/run_target_attack.sh

결과 및 로깅

• 로그는 logs/ 디렉토리에 저장됩니다.
• 결과(예: 오염된 쿼리 및 공격 성공률)는 results/ 디렉토리에 저장됩니다.

감사(Acknowledgement)

• Open-Prompt-Injection: 기본 모델 코드는 이 저장소를 기반으로 합니다.
• PoisonedRAG: 일부 구현은 이 프로젝트에서 각색되었습니다.
• Beir Benchmark: 평가를 위한 벤치마크 프레임워크로 사용됩니다.
• Contriever: RAG(Retrieval-Augmented Generation) 기능을 위해 활용됩니다.
• llms-attacks: 적대적 훈련을 위한 특정 구현은 이 저장소에서 파생되었습니다.