albert-y1n/pismithPISmith: Red Teaming basé sur l'apprentissage par renforcement pour les défenses contre l'injection de prompts
Ceci est une implémentation officielle de PISmith : Red Teaming basé sur l'apprentissage par renforcement pour les défenses contre l'injection de prompts
PISmith a été testé avec Python 3.10 et CUDA Version : 12.9
1. Créer un environnement conda Python 3.10
conda create -n PISmith python=3.10 -y
conda activate PISmith
2. Installer les dépendances
pip install -r requirements.txt
3. (Facultatif) Préparer le checkpoint du modèle Meta-SecAlign
Pour les expériences ciblant la défense secalign, exécutez le script de fusion fourni pour télécharger et fusionner le modèle de base avec l'adaptateur SecAlign :
python merge_meta_secalign.py
Ce script télécharge meta-llama/Llama-3.1-8B-Instruct et facebook/Meta-SecAlign-8B depuis HuggingFace, les fusionne et enregistre le résultat dans checkpoints/Meta-SecAlign-8B-merged/.
PIArena prend en charge l'entraînement et l'évaluation face à une gamme de défenses contre l'injection de prompts. Utilisez l'argument defense pour sélectionner la défense cible.
Défenses prises en charge :
secalign, none, promptguard, promptarmor, sandwich, instructional, datasentinel, piguard, datafilter
bash scripts/train_piarena.sh <defense> [train_gpus] [target_gpu] [target_port]
| Argument | Valeur par défaut | Description |
|---|---|---|
defense | secalign | Défense à utiliser pour l'entraînement |
train_gpus | "1,2,3" | Index GPU pour l'entraînement RL |
target_gpu | 0 | GPU du serveur vLLM cible |
target_port | 8010 | Port du serveur vLLM cible |
Exemples :
# 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]
| Argument | Valeur par défaut | Description |
|---|---|---|
checkpoint | — | Chemin du checkpoint de l'attaquant entraîné |
defense | secalign | Défense à évaluer |
target_port | 8000 | Port du serveur vLLM cible |
target_gpu | 0 | GPU du serveur vLLM cible |
attacker_gpu | 1 | GPU du serveur vLLM de l'attaquant |
attacker_port | 8001 | Port du serveur vLLM de l'attaquant |
num_samples | 10 | Pass@k : nombre d'échantillons par cas de test |
Exemples :
# 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
Prend en charge les cibles GPT-4o-mini, GPT-4o, GPT-5-nano et les serveurs vLLM locaux.
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
Évaluation :
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
AgentDyn est construit sur AgentDojo. Installez-le séparément avant d'exécuter des expériences AgentDyn :
git clone https://github.com/SaFo-Lab/AgentDyn.git
cd AgentDyn
pip install -e . --no-deps
L'entraînement prend en charge les suites AgentDyn telles que github, dailylife et shopping.
bash scripts/train_agentdyn.sh [target_type] [suites] [train_gpus]
# Example
bash scripts/train_agentdyn.sh gpt5-nano github "0,1,2,3"
L'évaluation fournit à la fois le pass@k et l'ASR moyen au niveau des échantillons, et prend en charge le service vLLM de l'attaquant en parallélisme de données :
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
Prend en charge une cible vLLM locale, GPT-4o-mini ou le mode multi-cibles.
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
Évaluation :
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
PISmith est évalué face à 7 méthodes de référence couvrant les catégories d'attaques statiques, basées sur la recherche et basées sur le RL. Toutes les méthodes basées sur le RL rapportent l'ASR@10 / ASR@1 ; les méthodes statiques et basées sur la recherche rapportent l'ASR@1.
| Méthode | Catégorie | Moy. ASR@10 | Moy. ASR@1 |
|---|---|---|---|
| Direct | Statique | — | 0.04 |
| Combined | Statique | — | 0.07 |
| TAP | Basée sur la recherche | — | 0.11 |
| PAIR | Basée sur la recherche | — | 0.16 |
| Strategy | Basée sur la recherche | — | 0.21 |
| Vanilla GRPO | Basée sur le RL | 0.13 | 0.05 |
| RL-Hammer | Basée sur le RL | 0.70 | 0.48 |
| PISmith (Ours) | Basée sur le RL | 1.00 | 0.87 |
ASR@1 de PISmith moyenné sur 13 benchmarks. L'utilité mesure la précision de la tâche sans attaque.
| Défense | Type | Utilité | PISmith ASR@1 |
|---|---|---|---|
| Sans défense | — | 0.74 | 0.92 |
| Sandwich | Prévention | 0.74 | 0.91 |
| Instructional | Prévention | 0.73 | 0.92 |
| PromptArmor | Prévention | 0.74 | 0.92 |
| DataFilter | Prévention | 0.63 | 0.49 |
| PIGuard | Filtre | 0.72 | 0.82 |
| PromptGuard | Filtre | 0.66 | 0.89 |
| DataSentinel | Filtre | 0.55 | 0.52 |
Il reste difficile pour les défenses de pointe d'atteindre simultanément une utilité élevée (≥0.70) et un faible ASR (≤0.60), ce qui révèle un compromis fondamental entre utilité et robustesse.
| Modèle cible | ASR@1 direct | PISmith ASR@10/1 |
|---|---|---|
| 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 |
| Modèle cible | Meilleure ASR@1 statique | PISmith ASR@10/1 |
|---|---|---|
| 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 |