albert-y1n/pismithPISmith: Reinforcement-Learning-basiertes Red Teaming für Prompt-Injection-Abwehrmaßnahmen
Dies ist die offizielle Implementierung von PISmith: Reinforcement Learning-basiertes Red Teaming für Prompt-Injection-Abwehrmechanismen
PISmith wurde mit Python 3.10 und CUDA Version 12.9 getestet.
1. Erstellen Sie eine Conda-Umgebung mit Python 3.10
conda create -n PISmith python=3.10 -y
conda activate PISmith
2. Abhängigkeiten installieren
pip install -r requirements.txt
3. (Optional) Den Meta-SecAlign-Modell-Checkpoint vorbereiten
Führen Sie für Experimente gegen die secalign-Abwehr das mitgelieferte Merge-Skript aus, um das Basismodell mit dem SecAlign-Adapter herunterzuladen und zusammenzuführen:
python merge_meta_secalign.py
Dies lädt meta-llama/Llama-3.1-8B-Instruct und facebook/Meta-SecAlign-8B von HuggingFace herunter, führt sie zusammen und speichert das Ergebnis unter checkpoints/Meta-SecAlign-8B-merged/.
PIArena unterstützt Training und Evaluierung gegen eine Reihe von Prompt-Injection-Abwehrmechanismen. Verwenden Sie das Argument defense, um die Zielabwehr auszuwählen.
Unterstützte Abwehrmechanismen:
secalign, none, promptguard, promptarmor, sandwich, instructional, datasentinel, piguard, datafilter
bash scripts/train_piarena.sh <defense> [train_gpus] [target_gpu] [target_port]
| Argument | Standard | Beschreibung |
|---|---|---|
defense | secalign | Abwehr, gegen die trainiert wird |
train_gpus | "1,2,3" | GPU-Indizes für das RL-Training |
target_gpu | 0 | GPU für den Ziel-vLLM-Server |
target_port | 8010 | Port für den Ziel-vLLM-Server |
Beispiele:
# 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 | Standard | Beschreibung |
|---|---|---|
checkpoint | — | Pfad zum trainierten Angreifer-Checkpoint |
defense | secalign | Abwehr, gegen die evaluiert wird |
target_port | 8000 | Port für den Ziel-vLLM-Server |
target_gpu | 0 | GPU für den Ziel-vLLM-Server |
attacker_gpu | 1 | GPU für den Angreifer-vLLM-Server |
attacker_port | 8001 | Port für den Angreifer-vLLM-Server |
num_samples | 10 | Pass@k: Anzahl der Stichproben pro Testfall |
Beispiele:
# 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
Unterstützt GPT-4o-mini, GPT-4o, GPT-5-nano und lokale vLLM-Ziele.
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
Evaluierung:
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 basiert auf AgentDojo. Installieren Sie es separat, bevor Sie AgentDyn-Experimente ausführen:
git clone https://github.com/SaFo-Lab/AgentDyn.git
cd AgentDyn
pip install -e . --no-deps
Das Training unterstützt AgentDyn-Suiten wie github, dailylife und shopping.
bash scripts/train_agentdyn.sh [target_type] [suites] [train_gpus]
# Example
bash scripts/train_agentdyn.sh gpt5-nano github "0,1,2,3"
Die Evaluierung berichtet sowohl pass@k als auch die durchschnittliche ASR auf Stichprobenebene und unterstützt datenparalleles Serving des Angreifer-vLLM:
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
Unterstützt ein lokales vLLM-Ziel, GPT-4o-mini oder den Multi-Ziel-Modus.
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
Evaluierung:
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 wird gegen 7 Baselines aus den Kategorien statisch, suchbasiert und RL-basiert evaluiert. Alle RL-basierten Methoden berichten ASR@10 / ASR@1; statische und suchbasierte Methoden berichten ASR@1.
| Methode | Kategorie | Avg. ASR@10 | Avg. ASR@1 |
|---|---|---|---|
| Direct | Statisch | — | 0.04 |
| Combined | Statisch | — | 0.07 |
| TAP | Suchbasiert | — | 0.11 |
| PAIR | Suchbasiert | — | 0.16 |
| Strategy | Suchbasiert | — | 0.21 |
| Vanilla GRPO | RL-basiert | 0.13 | 0.05 |
| RL-Hammer | RL-basiert | 0.70 | 0.48 |
| PISmith (unser) | RL-basiert | 1.00 | 0.87 |
PISmith ASR@1 gemittelt über 13 Benchmarks. Der Nutzen misst die Aufgaben-Genauigkeit ohne Angriff.
| Abwehr | Typ | Nutzen | PISmith ASR@1 |
|---|---|---|---|
| Keine Abwehr | — | 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 | Filter | 0.72 | 0.82 |
| PromptGuard | Filter | 0.66 | 0.89 |
| DataSentinel | Filter | 0.55 | 0.52 |
Es bleibt eine Herausforderung für moderne Abwehrmechanismen, gleichzeitig einen hohen Nutzen (≥0.70) und eine niedrige ASR (≤0.60) zu erreichen, was einen grundlegenden Nutzen–Robustheits-Konflikt offenbart.
| Zielmodell | Direct ASR@1 | 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 |
| Zielmodell | Beste statische ASR@1 | 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 |