albert-y1n/pismithPISmith: Red Teaming baseado em Aprendizado por Reforço para Defesas contra Injeção de Prompts
Esta é uma implementação oficial de PISmith: Red Teaming baseado em Aprendizado por Reforço para Defesas contra Injeção de Prompt
O PISmith foi testado usando Python 3.10 e CUDA Version: 12.9
1. Crie um ambiente conda Python 3.10
conda create -n PISmith python=3.10 -y
conda activate PISmith
2. Instale as dependências
pip install -r requirements.txt
3. (Opcional) Prepare o checkpoint do modelo Meta-SecAlign
Para experimentos que visam a defesa secalign, execute o script de mesclagem fornecido para baixar e mesclar o modelo base com o adaptador SecAlign:
python merge_meta_secalign.py
Isso baixa meta-llama/Llama-3.1-8B-Instruct e facebook/Meta-SecAlign-8B do HuggingFace, mescla-os e salva o resultado em checkpoints/Meta-SecAlign-8B-merged/.
O PIArena oferece suporte a treinamento e avaliação contra uma variedade de defesas contra injeção de prompt. Use o argumento defense para selecionar a defesa alvo.
Defesas suportadas:
secalign, none, promptguard, promptarmor, sandwich, instructional, datasentinel, piguard, datafilter
bash scripts/train_piarena.sh <defense> [train_gpus] [target_gpu] [target_port]
| Argumento | Padrão | Descrição |
|---|---|---|
defense | secalign | Defesa contra a qual treinar |
train_gpus | "1,2,3" | Índices de GPU para treinamento RL |
target_gpu | 0 | GPU para o servidor vLLM alvo |
target_port | 8010 | Porta para o servidor vLLM alvo |
Exemplos:
# Train against SecAlign defense
bash scripts/train_piarena.sh secalign
# Train against no defense (plain LLM)
bash scripts/train_piarena.sh none
bash scripts/eval_piarena.sh <checkpoint> <defense> [target_port] [target_gpu] [attacker_gpu] [attacker_port] [num_samples]
| Argumento | Padrão | Descrição |
|---|---|---|
checkpoint | — | Caminho para o checkpoint do atacante treinado |
defense | secalign | Defesa contra a qual avaliar |
target_port | 8000 | Porta para o servidor vLLM alvo |
target_gpu | 0 | GPU para o servidor vLLM alvo |
attacker_gpu | 1 | GPU para o servidor vLLM do atacante |
attacker_port | 8001 | Porta para o servidor vLLM do atacante |
num_samples | 10 | pass@k: número de amostras por caso de teste |
Exemplos:
# Evaluate against SecAlign (default settings)
bash scripts/eval_piarena.sh checkpoints/piarena/checkpoint-500 secalign
# Evaluate against no piguard, pass@10
bash scripts/eval_piarena.sh checkpoints/piarena_none/checkpoint-500 piguard
Suporta alvos GPT-4o-mini, GPT-4o, GPT-5-nano e vLLM local.
bash scripts/train_agentdojo.sh [target_type] [suites] [train_gpus]
# Default: GPT-4o-mini target on the firsr 7 injected task of workspace suite
bash scripts/train_agentdojo.sh
# Train on all suites (workspace, banking, travel, slack)
bash scripts/train_agentdojo.sh gpt4o-mini all
Avaliação:
bash scripts/eval_agentdojo.sh <checkpoint> [target_type] [eval_suites] [num_samples] [target_defense]
# Example
bash scripts/eval_agentdojo.sh checkpoints/agentdojo/checkpoint-500 gpt4o-mini
O AgentDyn é construído sobre o AgentDojo. Instale-o separadamente antes de executar experimentos com AgentDyn:
git clone https://github.com/SaFo-Lab/AgentDyn.git
cd AgentDyn
pip install -e . --no-deps
O treinamento suporta suítes do AgentDyn, como github, dailylife e shopping.
bash scripts/train_agentdyn.sh [target_type] [suites] [train_gpus]
# Example
bash scripts/train_agentdyn.sh gpt5-nano github "0,1,2,3"
A avaliação reporta tanto pass@k quanto ASR média em nível de amostra, e suporta servidor vLLM do atacante em paralelismo de dados:
ATTACKER_GPUS=0,1,2,3 ATTACKER_DP_SIZE=4 \
bash scripts/eval_agentdyn.sh checkpoints/agentdyn/checkpoint-500 gpt5-nano "github,dailylife,shopping" 5
Suporta um alvo vLLM local, GPT-4o-mini ou modo multi-alvo.
bash scripts/train_injecagent.sh [target_type] [train_gpus] [target_gpu] [target_port]
# Default: local vLLM target (Meta-SecAlign-8B)
bash scripts/train_injecagent.sh
# GPT-4o-mini API target
bash scripts/train_injecagent.sh gpt4o-mini
Avaliação:
bash scripts/eval_injecagent.sh <checkpoint> [target_type] [target_gpu] [target_port] [eval_gpu] [num_samples]
# Example
bash scripts/eval_injecagent.sh checkpoints/injecagent/checkpoint-500
O PISmith é avaliado contra 7 baselines que abrangem categorias de ataque estáticas, baseadas em busca e baseadas em RL. Todos os métodos baseados em RL reportam ASR@10 / ASR@1; métodos estáticos e baseados em busca reportam ASR@1.
| Método | Categoria | ASR@10 Médio | ASR@1 Médio |
|---|---|---|---|
| Direct | Estática | — | 0.04 |
| Combined | Estática | — | 0.07 |
| TAP | Baseada em busca | — | 0.11 |
| PAIR | Baseada em busca | — | 0.16 |
| Strategy | Baseada em busca | — | 0.21 |
| Vanilla GRPO | Baseada em RL | 0.13 | 0.05 |
| RL-Hammer | Baseada em RL | 0.70 | 0.48 |
| PISmith (Nosso) | Baseada em RL | 1.00 | 0.87 |
ASR@1 do PISmith calculado com média sobre 13 benchmarks. Utilidade mede a precisão da tarefa sem ataque.
| Defesa | Tipo | Utilidade | ASR@1 do PISmith |
|---|---|---|---|
| Sem Defesa | — | 0.74 | 0.92 |
| Sandwich | Prevenção | 0.74 | 0.91 |
| Instructional | Prevenção | 0.73 | 0.92 |
| PromptArmor | Prevenção | 0.74 | 0.92 |
| DataFilter | Prevenção | 0.63 | 0.49 |
| PIGuard | Filtro | 0.72 | 0.82 |
| PromptGuard | Filtro | 0.66 | 0.89 |
| DataSentinel | Filtro | 0.55 | 0.52 |
Permanece desafiador para defesas de última geração alcançar simultaneamente alta utilidade (≥0.70) e baixo ASR (≤0.60), revelando um trade-off fundamental entre utilidade e robustez.
| Modelo Alvo | ASR@1 Direto | ASR@10/1 do PISmith |
|---|---|---|
| Meta-SecAlign-8B | 0.00 | 1.00 / 0.99 |
| GPT-4o-mini | 0.02 | 1.00 / 0.99 |
| GPT-4.1-nano | 0.01 | 1.00 / 1.00 |
| GPT-5-nano | 0.00 | 1.00 / 0.95 |
| Modelo Alvo | Melhor ASR@1 Estático | ASR@10/1 do PISmith |
|---|---|---|
| GPT-4o-mini | 0.23 | 0.78 / 0.62 |
| GPT-4.1-nano | 0.20 | 0.81 / 0.64 |
| GPT-5-nano | 0.01 | 0.38 / 0.24 |