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

Korea Advanced Institute of Science and Technology (KAIST)

Notícias

• 2026.09 Nosso artigo foi aceito na AACL-IJCNLP 2026 (Main)! 🎉
• 2026.08 O artigo está disponível no arXiv!
• 2026.08 O código foi disponibilizado!

Visão Geral

O batch prompting é uma estratégia prática de inferência que agrupa múltiplas consultas em uma única chamada. Mostramos que seu sucesso em termos de utilidade não se estende à segurança: uma pergunta prejudicial que é recusada de forma confiável isoladamente pode eliciar uma resposta prejudicial quando inserida em um lote de perguntas benignas.

Este repositório fornece o código oficial para:

1. Gerar dados de ataque com batch-prompt
2. Executar inferência em modelos de código aberto (Phi / Llama / Qwen) e APIs comerciais (GPT / Gemini / Claude)
3. Avaliar o sucesso do ataque com um ensemble de LlamaGuard-4, WildGuard e Beaver-Dam

Preparação

1. Clone este repositório.

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

2. Configuração do ambiente

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

É necessária GPU para inferência de modelos locais e para os três modelos juízes. Testado com Python 3.10 + PyTorch 2.4 + transformers==5.2.

Nota: src/evaluate/llamaguard4/ inclui uma cópia vendorizada do transformers para o LlamaGuard-4 (Llama4ForConditionalGeneration). Ela é adicionada ao sys.path apenas enquanto o LlamaGuard está carregando.

3. Registre modelos / chaves de API

Edite configs/models.yaml com os caminhos locais dos checkpoints do HuggingFace (geradores + juízes), e defina as chaves de API em configs/api.yaml (ou via variáveis de ambiente OPENAI_API_KEY / GEMINI_API_KEY / ANTHROPIC_API_KEY / HF_TOKEN).

Dados e modelos necessários

• Preenchimento benigno: data/context_dataset/gsm8k.jsonl (cada linha tem uma chave "question")
• Datasets prejudiciais (baixados automaticamente): JailbreakBench/JBB-Behaviors, walledai/StrongREJECT
• Geradores: Phi-4 / Llama-3.1-8B-Instruct / Qwen3-8B
• Juízes: Llama-Guard-4-12B, PKU-Alignment/beaver-dam-7b, allenai/wildguard

Gerando Dados de Batch-Prompt

Padrão: --harmful_dataset JBB, tamanho do lote n = 1..10, 313 amostras. Os espaços benignos são preenchidos a partir do GSM8K (single_math).

python -m src.utils.generate_batch_question_prompt

As saídas são gravadas em data/batch_prompt/single_math/{n}/pos_{k}.jsonl.

Executando Inferência

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

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

Varredura em lote 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 comerciais (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

As saídas são salvas em result/inference/<model>/... como JSONL com pares {prompt, response}.

Avaliando o Sucesso do Ataque

Cada resposta é pontuada por LlamaGuard-4, WildGuard e Beaver-Dam, e então combinada por votação majoritária:

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

Você também pode executar a avaliação no mesmo loop da inferência:

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

Início 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ção

Considere citar nosso artigo se nosso trabalho ajudar em sua pesquisa.

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

Contato

Para quaisquer dúvidas sobre nosso artigo ou código, envie um e-mail para [email protected].