
이 저장소는 논문 "[Safety in Batches? Understanding and Mitigating Safety Failures in Batch Prompting]"의 공식 구현을 포함하고 있습니다.
Kihyun Kim, Hee-Seon Kim, Wonjun Lee, Changick Kim
한국과학기술원 (KAIST)
2026.09 본 논문이 AACL-IJCNLP 2026 (Main) 에 채택되었습니다! 🎉2026.08 논문이 arXiv에 공개되었습니다!2026.08 코드가 공개되었습니다!배치 프롬프팅은 여러 질의를 하나의 호출에 묶는 실용적인 추론 전략입니다. 우리는 이것이 유용성 측면에서 성공적이라는 점이 안전성 측면으로 확장되지 않는다는 것을 보입니다. 즉, 단독으로는 안정적으로 거부되는 유해한 질문이 무해한 질문들의 배치에 포함되면 유해한 응답을 유도할 수 있습니다.
이 저장소는 다음을 위한 공식 코드를 제공합니다:
git clone https://github.com/96kihyun/batch_jailbreak.git
cd batch_jailbreak
conda create -n batch python=3.10 -y
conda activate batch
pip install -r requirements.txt
로컬 모델 추론과 세 가지 판정 모델을 위해 GPU가 필요합니다. Python 3.10 + PyTorch 2.4 + transformers==5.2에서 테스트되었습니다.
참고:
src/evaluate/llamaguard4/에는 LlamaGuard-4용transformers의 벤더링된 복사본(Llama4ForConditionalGeneration)이 포함되어 있습니다. 이는 LlamaGuard가 로드되는 동안에만sys.path에 추가됩니다.
configs/models.yaml을 로컬 HuggingFace 체크포인트 경로(생성기 + 판정기)로 편집하고, configs/api.yaml에 API 키를 설정합니다(또는 환경 변수 OPENAI_API_KEY / GEMINI_API_KEY / ANTHROPIC_API_KEY / HF_TOKEN 사용).
필수 데이터 및 모델
data/context_dataset/gsm8k.jsonl (각 줄에 "question" 키 포함)JailbreakBench/JBB-Behaviors, walledai/StrongREJECTPKU-Alignment/beaver-dam-7b, allenai/wildguard기본값: --harmful_dataset JBB, 배치 크기 n = 1..10, 313개 샘플. 무해한 슬롯은 GSM8K(single_math)에서 채워집니다.
python -m src.utils.generate_batch_question_prompt
출력은 data/batch_prompt/single_math/{n}/pos_{k}.jsonl 아래에 기록됩니다.
python scripts/inference/run_inference_batch_prompt.py \
--model_name qwen \
--case single_math --n 9 --pos 2 \
--harmful_dataset JBB
(model, case, batch_size, pos)에 대한 배치 스윕:
bash scripts/inference/run_inference_batch_prompt_eval.sh \
--dataset jbb --gpu_id 0 \
--models phi qwen --cases single_math \
--batch_sizes 12 --pos_values 1 2
python scripts/inference/run_inference_api.py \
--api gpt \
--input_file data/batch_prompt_jbb/single_math/12/pos_2.jsonl \
--model_name gpt --output_subdir batch_prompt_jbb/single_math/12/pos_2 \
--harmful_dataset JBB --extract_pos 2
출력은 result/inference/<model>/... 아래에 {prompt, response} 쌍을 담은 JSONL로 저장됩니다.
각 응답은 LlamaGuard-4, WildGuard, Beaver-Dam에 의해 점수가 매겨진 후 다수결로 결합됩니다:
python scripts/eval/run_ensemble.py \
--input_file result/inference/qwen/batch_prompt_jbb/single_math/12/qwen_single_math_n12_pos_2_response.jsonl \
--output_dir result/evaluation/batch_prompt_jbb/qwen/single_math/12/pos_2
추론과 동일한 루프에서 평가를 실행할 수도 있습니다:
bash scripts/inference/run_inference_batch_prompt_eval.sh \
--dataset jbb --gpu_id 0 \
--models qwen --cases single_math \
--batch_sizes 12 --pos_values 2 \
--conda_env batch
# 1. data
python -m src.utils.generate_batch_question_prompt \
--harmful_dataset JBB --min_n 12 --max_n 12 --num_samples 100
# 2. inference
python scripts/inference/run_inference_batch_prompt.py \
--model_name llama --case single_math --n 12 --pos 2 --harmful_dataset JBB
# 3. evaluation
python scripts/eval/run_ensemble.py \
--input_file result/inference/llama/batch_prompt/llama_single_math_n12_pos_2_response.jsonl
본 연구가 귀하의 연구에 도움이 된다면 논문을 인용해 주시기 바랍니다.
@article{kim2026safety,
title={Safety in Batches? Understanding and Mitigating Safety Failures in Batch Prompting},
author={Kim, Kihyun and Kim, HeeSeon and Lee, Wonjun and Kim, Changick},
journal={arXiv preprint arXiv:2608.02681},
year={2026}
}
논문이나 코드에 관한 질문이 있으시면 [email protected]로 이메일을 보내주시기 바랍니다.
| 플래그 | 기본값 | 설명 |
|---|
--harmful_dataset | JBB | JBB 또는 StrongREJECT |
--min_n / --max_n | 1 / 10 | 배치 크기 N의 포함 범위 |
--num_samples | 313 | 유해 프롬프트 수 |
--seed | 42 | RNG 시드 |
--output_dir | ./data/batch_prompt | 출력 루트 |