Kihyun Kim, Hee-Seon Kim, Wonjun Lee, Changick Kim

Korea Advanced Institute of Science and Technology (KAIST)

Noticias

• 2026.09 ¡Nuestro artículo ha sido aceptado en AACL-IJCNLP 2026 (Main)! 🎉
• 2026.08 ¡El artículo está disponible en arXiv!
• 2026.08 ¡El código ha sido publicado!

Descripción general

El batch prompting es una estrategia de inferencia práctica que agrupa múltiples consultas en una sola llamada. Mostramos que su éxito en cuanto a utilidad no se extiende a la seguridad: una pregunta dañina que es rechazada de forma fiable de manera aislada puede provocar una respuesta dañina cuando se inserta en un lote de preguntas benignas.

Este repositorio proporciona el código oficial para:

1. Generar datos de ataque mediante batch-prompt
2. Ejecutar inferencia en modelos de código abierto (Phi / Llama / Qwen) y APIs comerciales (GPT / Gemini / Claude)
3. Evaluar el éxito del ataque con un conjunto de LlamaGuard-4, WildGuard y Beaver-Dam

Preparación

1. Clona este repositorio.

git clone https://github.com/96kihyun/batch_jailbreak.git
cd batch_jailbreak

2. Configuración del entorno

conda create -n batch python=3.10 -y
conda activate batch
pip install -r requirements.txt

Se requiere GPU para la inferencia de modelos locales y para los tres modelos jueces. Probado con Python 3.10 + PyTorch 2.4 + transformers==5.2.

Nota: src/evaluate/llamaguard4/ incluye una copia vendorizada de transformers para LlamaGuard-4 (Llama4ForConditionalGeneration). Se añade a sys.path solo mientras LlamaGuard se está cargando.

3. Registra modelos / claves de API

Edita configs/models.yaml con las rutas locales de los checkpoints de HuggingFace (generadores + jueces), y establece las claves de API en configs/api.yaml (o mediante las variables de entorno OPENAI_API_KEY / GEMINI_API_KEY / ANTHROPIC_API_KEY / HF_TOKEN).

Datos y modelos requeridos

• Relleno benigno: data/context_dataset/gsm8k.jsonl (cada línea tiene una clave "question")
• Conjuntos de datos dañinos (descarga automática): JailbreakBench/JBB-Behaviors, walledai/StrongREJECT
• Generadores: Phi-4 / Llama-3.1-8B-Instruct / Qwen3-8B
• Jueces: Llama-Guard-4-12B, PKU-Alignment/beaver-dam-7b, allenai/wildguard

Generación de datos de Batch-Prompt

Predeterminado: --harmful_dataset JBB, tamaño de lote n = 1..10, 313 muestras. Las ranuras benignas se rellenan desde GSM8K (single_math).

python -m src.utils.generate_batch_question_prompt

Las salidas se escriben en data/batch_prompt/single_math/{n}/pos_{k}.jsonl.

Ejecución de la inferencia

Modelos de código abierto (Phi / Llama / Qwen)

python scripts/inference/run_inference_batch_prompt.py \
  --model_name qwen \
  --case single_math --n 9 --pos 2 \
  --harmful_dataset JBB

Barrido por lotes sobre (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

APIs comerciales (GPT / Gemini / Claude)

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

Las salidas se guardan en result/inference/<model>/... como JSONL con pares {prompt, response}.

Evaluación del éxito del ataque

Cada respuesta es puntuada por LlamaGuard-4, WildGuard y Beaver-Dam, y luego se combinan mediante voto por mayoría:

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

También puedes ejecutar la evaluación en el mismo bucle que la inferencia:

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

Inicio rápido (Smoke Run)

# 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

Cita

Considera citar nuestro artículo si nuestro trabajo ayuda a tu investigación.

@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}
}

Contacto

Para cualquier pregunta sobre nuestro artículo o código, envía un correo a [email protected].