RAGSieve

RAGSieve: 신뢰할 수 있는 참조 없이 RAG에서 지식 중독 문서 탐지를 위한 코드입니다.

RAGSieve는 검색 증강 생성(RAG) 시스템의 두 지점에서 지식 중독 문서를 탐지합니다. RSQ는 온라인 필터로, 생성에 선택된 다섯 문서를 현재 쿼리에 대한 순위 6--20과 비교하여 점수를 매깁니다. RSG는 오프라인 스캐너로, 모든 코퍼스 문서를 자체 의미--어휘 이웃과 비교하여 점수를 매깁니다. 이 패키지에는 논문에서 사용된 RSQ 및 RSG 구현, 정확한 전체 코퍼스 검색, Top-5 필터링 및 리필, 문서 수준 평가, 그리고 간결하게 실행 가능한 예제가 포함되어 있습니다.

저장소 내용

data/datasets/       NQ, HotpotQA, MS MARCO 텍스트 지식 베이스
data/demo/           네 개의 대상 쿼리와 20개의 예제 중독 문서
src/ragsieve/     RSQ, RSG, 검색, 필터링, 메트릭
install.sh           환경 설치
data_preparation.sh  아홉 개의 dense 인덱스 로컬 구축
run_demo.sh          종단 간 RSQ 및 RSG 기능 검사

공개된 세 지식 베이스에는 128,044개의 NQ 문서, 9,961개의 HotpotQA 문서, 8,239개의 MS MARCO 패시지가 포함되어 있습니다. 여기에는 원시 텍스트, 샘플링된 1,000개의 쿼리, qrels, 대상 답변, 그리고 고정된 100개 쿼리 공격 하위 집합이 포함됩니다. Dense 벡터는 로컬에서 생성되며 저장소의 일부가 아닙니다.

데모에는 NQ에서 가져온 두 개의 PR-W와 두 개의 CEM-C 실현이 포함되어 있으며, 쿼리당 다섯 개의 중독 문서가 있습니다. 이러한 공격은 전체 공격 컬렉션이나 공격 생성 구현을 제공하지 않고도 두 탐지기를 통과하는 짧고 검사 가능한 경로를 제공합니다.

설치

Python 3.11 또는 3.12, uv, 그리고 CUDA 지원 GPU를 권장합니다.

bash install.sh

탐지기는 처음 사용할 때 논문에 명시된 Hugging Face 모델을 다운로드합니다.

종단 간 QA를 위해서는 공유된 OpenAI 호환 구성을 복사하고 세 값을 채우십시오:

cp .env.example .env

OPENAI_BASE_URL=https://api.openai.com/v1
OPENAI_MODEL=your-model-name
OPENAI_API_KEY=your-api-key

동일한 엔드포인트와 모델이 답변을 생성하고 의미론적 ASR 판정을 수행합니다.

빠른 평가

bash run_demo.sh

이 명령은 520개 문서 RSG 스냅샷에 대한 로컬 BGE-M3 인덱스를 구축하고, 준비된 전체 코퍼스 검색 컨텍스트에 대해 RSQ를 실행하며, 생성 Top-5를 필터링 및 리필하고, 두 탐지기를 모두 평가합니다. 결과는 outputs/demo/에 기록됩니다:

• rsq-metrics.json: RSQ AUROC 및 5% 클린 제거 시 중독 탐지;
• rsg-metrics.json: 쿼리 없는 코퍼스 스냅샷에 대한 RSG 메트릭;
• filtered-top5.jsonl: RSQ 후 생성에 유지된 다섯 문서.
• rsg-*-top5.jsonl 및 joint-*-top5.jsonl: RSG 단독 및 공동 필터링 후 각 데모 공격에 대해 유지된 컨텍스트. QA는 별도의 옵트인 명령입니다.

공개된 예제에서 RSQ는 PR-W와 CEM-C 모두에 대해 100% AUROC와 5% 클린 제거 시 100% 중독 탐지에 도달합니다. RSG는 동일한 클린 제거 예산에서 100% AUROC와 100% 중독 탐지에 도달합니다; 고정 규칙은 중독 문서의 85%와 클린 문서의 0%를 제거합니다. 이 값들은 선별된 데모에 대한 기능 검사이며, 논문의 집계 결과가 아닙니다.

정답 필드는 평가 명령에 의해서만 읽힙니다. RSQ나 RSG에 전달되지 않습니다.

완전한 로컬 인덱스 구축

bash data_preparation.sh

기본적으로 이는 논문에서 평가된 3개 데이터셋 x 3개 검색기(BGE-M3, E5-large-v2, all-MiniLM-L6-v2)를 구축합니다. 출력은 무시된 data/indices/ 디렉토리에 배치됩니다. 스크립트를 편집하지 않고 하위 집합을 선택할 수 있습니다:

DATASETS="nq" ENCODERS="bge-m3" DEVICE="cuda:0" bash data_preparation.sh

각 인덱스는 논문의 풀링, 접두사, 절단, 코사인 정규화 설정을 사용하여 완전한 공개 코퍼스에서 생성됩니다.

RSQ: 온라인 쿼리 시점 필터링

번들 공격이 포함된 클린 검색과 검색을 준비하십시오:

uv run ragsieve prepare-contexts \
  --dataset-dir data/datasets/nq \
  --index-dir data/indices/nq/bge-m3 \
  --subset data/demo/subset.json \
  --attacks data/demo/attacks.jsonl \
  --output outputs/demo/contexts.jsonl \
  --device cuda:0

prepare-contexts는 각 쿼리에 대해 전체 코퍼스에 대한 정확한 검색을 한 번 수행하고, 주입된 문서를 별도로 점수 매기며, 상위 100개를 원래 순서대로 저장합니다. RSQ는 후보 1--5를 순위 6--20과 비교하여 점수를 매깁니다; 나머지 순위는 리필에 사용할 수 있습니다. RSQ를 실행하고, 문서 수준 메트릭을 계산하며, 생성 컨텍스트를 리필하십시오:

uv run ragsieve detect \
  --input outputs/demo/contexts.jsonl \
  --output outputs/demo/rsq-predictions.jsonl \
  --device cuda:0

uv run ragsieve evaluate \
  --predictions outputs/demo/rsq-predictions.jsonl \
  --output outputs/demo/rsq-metrics.json

uv run ragsieve filter-contexts \
  --contexts outputs/demo/contexts.jsonl \
  --predictions outputs/demo/rsq-predictions.jsonl \
  --output outputs/demo/filtered-top5.jsonl

uv run ragsieve qa \
  --input outputs/demo/filtered-top5.jsonl \
  --output outputs/demo/qa-records.jsonl \
  --summary outputs/demo/qa-summary.json \
  --env-file .env

RSQ는 네 가지 증거 소스를 결합합니다: 답변 앵커 집중도, 문자 체계 무결성, 로컬 언어 모델 서프라이절, 그리고 쿼리 정렬 전환. CLI 기본값의 모든 설정은 논문 설정입니다.

RSG: 오프라인 코퍼스 검사

RSG는 문서 텍스트와 RAG 인덱스가 이미 유지 관리하는 벡터를 사용합니다:

uv run ragsieve detect-graph \
  --documents data/demo/corpus.jsonl \
  --embeddings data/indices/demo/bge-m3/embeddings.npy \
  --output outputs/demo/rsg-predictions.jsonl \
  --device cuda:0

uv run ragsieve evaluate-graph \
  --predictions outputs/demo/rsg-predictions.jsonl \
  --labels data/demo/labels.jsonl \
  --output outputs/demo/rsg-metrics.json

RSG는 정확한 코사인 top-16 그래프를 구축하고, 의미상 가깝고 어휘상 먼 이웃을 유지하며, 각 문서의 자체 이웃 바닥 대비 밀도 상승을 측정하고, 경험적 코퍼스 꼬리와 문자 체계 무결성을 결합합니다.

공동 배포

공동 배포는 직렬입니다: RSG가 먼저 문서를 제외한 다음, RSQ가 새로운 상위 다섯을 생존한 순위 6--20과 비교하여 점수를 매깁니다. 플래그된 후보는 제거되고 컨텍스트는 생존한 순위에서 한 번 리필됩니다; 대체 문서는 다시 점수 매기지 않습니다. bash run_demo.sh 후 PR-W 예제는 다음과 같습니다:

uv run ragsieve quarantine-contexts \
  --contexts data/demo/contexts.jsonl \
  --graph-predictions outputs/demo/rsg-predictions.jsonl \
  --condition pr_w \
  --output outputs/demo/serial-pr_w-contexts.jsonl
uv run ragsieve detect \
  --input outputs/demo/serial-pr_w-contexts.jsonl \
  --output outputs/demo/serial-pr_w-rsq.jsonl --device cuda:0
uv run ragsieve filter-contexts \
  --contexts outputs/demo/serial-pr_w-contexts.jsonl \
  --predictions outputs/demo/serial-pr_w-rsq.jsonl \
  --output outputs/demo/joint-pr_w-top5.jsonl

uv run ragsieve qa \
  --input outputs/demo/joint-pr_w-top5.jsonl \
  --output outputs/demo/joint-pr_w-qa.jsonl \
  --summary outputs/demo/joint-pr_w-qa-summary.json \
  --env-file .env

RSG 단독의 경우, 원래 컨텍스트에 대해 --graph-predictions와 --condition만으로 filter-contexts를 실행하십시오. CEM-C의 경우 --condition ipi_cem_c와 별도의 출력 경로를 사용하십시오. --condition은 RSG 스캔을 평가 중인 검색 조건에 바인딩합니다; 중독되지 않은 QA는 클린 코퍼스의 별도 RSG 스캔을 사용합니다. 번들 RSG 스냅샷에는 두 데모 공격이 모두 포함되어 있습니다. 전체 코퍼스 및 클린 QA 단계는 아티팩트 가이드에 있습니다.

구성 요소 절제

저장된 증거는 모델을 다시 로드하지 않고도 게시된 절제를 실행하기에 충분합니다:

uv run ragsieve ablate \
  --predictions outputs/demo/rsq-predictions.jsonl \
  --variant without-answer-anchor \
  --output outputs/demo/rsq-without-answer-anchor.jsonl
uv run ragsieve evaluate \
  --predictions outputs/demo/rsq-without-answer-anchor.jsonl \
  --output outputs/demo/rsq-without-answer-anchor-metrics.json

uv run ragsieve ablate \
  --predictions outputs/demo/rsg-predictions.jsonl \
  --variant corpus-local-only \
  --output outputs/demo/rsg-corpus-local-only.jsonl
uv run ragsieve evaluate-graph \
  --predictions outputs/demo/rsg-corpus-local-only.jsonl \
  --labels data/demo/labels.jsonl \
  --output outputs/demo/rsg-corpus-local-only-metrics.json

RSQ는 또한 without-script-integrity, without-surprisal, without-query-alignment을 지원합니다; RSG는 또한 script-integrity-only를 지원합니다. 코퍼스-로컬 절제는 전체 5% 경보 예산을 그래프 분기에 할당합니다. 절제된 예측은 동일한 필터링 및 QA 명령에 전달될 수 있습니다.

탐지 비용

uv run ragsieve benchmark --mode rsq \
  --input data/demo/contexts.jsonl \
  --output outputs/demo/rsq-cost.json --device cuda:0
uv run ragsieve benchmark --mode rsg \
  --input data/demo/corpus.jsonl \
  --embeddings data/indices/demo/bge-m3/embeddings.npy \
  --output outputs/demo/rsg-cost.json --device cuda:0

측정은 모델/데이터 로딩, 검색, QA를 제외합니다. 모델과 임베딩은 상주 상태로 유지됩니다; CUDA 작업은 각 측정 주변에서 동기화됩니다. 이 명령은 평균, 중앙값, P95 실행 시간, 처리량, 최대 GPU 할당을 보고합니다. 논문의 코퍼스 스캔 워크로드에는 전체 NQ 인덱스를 사용하십시오; 작은 데모는 기능 검사입니다. 워크로드 및 환경 세부 사항은 아티팩트 가이드에 있습니다.

논문과의 매핑

이 릴리스는 RAGSieve 항목과 절제를 재현합니다; 기준선 항목은 논문에서 인용된 구현에서 얻습니다. 아래 매핑은 현재 프리프린트를 따릅니다.