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
输出以 JSONL 格式保存在 result/inference/<model>/... 下,包含 {prompt, response} 对。
每个回答由 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 | 随机数种子 |
--output_dir | ./data/batch_prompt | 输出根目录 |