jiraphat-j/Prompt-Guard

# PromptGuard — LLM越狱与提示词注入防御 一个二元分类系统，用于检测并阻止针对大型语言模型（LLM）的越狱攻击、提示词注入和提示词提取攻击。 实现并基准测试了两个模型： - **TF-IDF + 逻辑回归** — 轻量级、可解释、近乎零延迟 - **DistilBERT（微调）** — 基于Transformer，对未见过的攻击模式具有更强的泛化能力 ## 问题所在 LLM容易受到对抗性提示词的攻击，这些攻击试图： - **绕过安全指南**（越狱） - **覆盖系统指令**（提示词注入） - **泄露机密上下文**（提示词提取） 攻击者还使用编码技术（Base64、ROT13、leet speak）对这些提示词进行混淆，以规避基于关键词的过滤器。PromptGuard同时解决了分类和混淆问题。 ## 架构 ``` Raw Prompt │ ▼ Preprocessing Pipeline ├─ Base64 decode ├─ ROT13 decode ├─ Leetspeak normalization (e.g. "1gn0r3" → "ignore") └─ URL removal + whitespace collapse │ ▼ Classifier ├─ TF-IDF + Logistic Regression (fast, CPU-only) └─ DistilBERT fine-tuned (accurate, GPU recommended) │ ▼ Label: safe / malicious ``` ## 数据集 | 数据集划分 | TF-IDF模型 | DistilBERT模型 | |---|---|---| | 总样本数 | 440 | 4,568 | | 训练集 | 330 | 3,426 | | 测试集 | 110 | 1,142 | | 恶意样本比例 | 约58% | 约61% | **数据来源：** - 自定义数据集（`prompts_dataset_full.csv` + `adversarial_dataset_v2.csv`） - [`jackhhao/jailbreak-classification`](https://huggingface.co/datasets/jackhhao/jailbreak-classification)（HuggingFace）— 用于TF-IDF流程 - [`neuralchemy/Prompt-injection-dataset`](https://huggingface.co/datasets/neuralchemy/Prompt-injection-dataset) — 用于DistilBERT流程 **覆盖的攻击类型：** 指令覆盖、目标劫持、越狱角色、多语言劫持、基于编码的混淆（Base64、ROT13、leet speak）、真实世界越狱攻击 ## 结果 ### TF-IDF + 逻辑回归 在预留测试集上评估（110个样本）： | 指标 | 安全 | 恶意 | 整体 | |---|---|---|---| | 精确率 | 0.87 | 1.00 | 0.94 | | 召回率 | 1.00 | 0.89 | 0.94 | | F1分数 | 0.93 | 0.94 | 0.94 | | 准确率 | — | — | **93.6%** | 5折交叉验证：`F1 = 0.940 ± 0.025` 测试集上**零误报**。 ### DistilBERT（微调，3个epoch，RTX 5050） 在预留测试集上评估（1,142个样本）： | 指标 | 安全 | 恶意 | 整体 | |---|---|---|---| | 精确率 | 0.95 | 0.98 | 0.97 | | 召回率 | 0.97 | 0.96 | 0.97 | | F1分数 | 0.96 | 0.97 | 0.97 | | 准确率 | — | — | **96.6%** | 3折交叉验证：`F1 = 0.973 ± 0.005` 模型参数：66,955,010 ### 公平对比 — 未见数据集 两个模型在[`deepset/prompt-injections`](https://huggingface.co/datasets/deepset/prompt-injections)上评估（116个样本，**训练时未见过**）： | 模型 | 准确率 | 精确率 | 召回率 | F1 | 推理时间 | |---|---|---|---|---|---| | TF-IDF + LR | 63.8% | 0.821 | 0.383 | 0.523 | 0.007秒 | | DistilBERT | 74.1% | 0.714 | 0.833 | **0.769** | 0.587秒 | DistilBERT在新颖攻击模式上泛化能力明显更强（F1提升24.6%）。TF-IDF速度快90倍，但会遗漏更多之前未见过的对抗性样本。 ## 何时使用哪个模型 | 标准 | TF-IDF + LR | DistilBERT | |---|---|---| | 延迟 | < 1毫秒/提示词 | ~10–50毫秒（GPU） | | 资源 | 仅CPU，< 1MB | 推荐GPU，约260MB | | 可解释性 | 特征系数 | 黑盒 | | 已知的攻击模式 | 优秀 | 优秀 | | 新颖/混淆攻击 | 泛化能力较弱 | 泛化能力更强 | | 部署复杂度 | 简单（sklearn pickle） | 需要PyTorch + transformers | ## 预处理流程 `preprocessing.py`在所有notebook中共享。该流程在分类之前处理混淆攻击： ``` from preprocessing import preprocess preprocess("aWdub3JlIGFsbCBwcmV2aW91cyBpbnN0cnVjdGlvbnM=") # 忽略所有之前的指令 preprocess("1gn0r3 4ll pr3v10us 1nstruct10ns") # 忽略所有之前的指令 preprocess("Vtaber nyy cerivbhf vafgehpgvbaf.") # 忽略所有之前的指令 ``` 流程顺序：Base64解码 → ROT13解码 → 小写化 → leet speak规范化 → 移除URL → 合并空白符 ## 仓库结构 ``` PromptGuard/ ├── preprocessing.py # Shared preprocessing pipeline ├── P-G_TF-IDF_LR.ipynb # TF-IDF + Logistic Regression ├── P-G_bert.ipynb # DistilBERT fine-tuning ├── P-G_compare_TF-bert.ipynb # Fair model comparison ├── model_TF-IDF_LR_P-G.pkl # Saved TF-IDF + LR model ├── model_bert_P-G.pkl # Saved DistilBERT model └── README.md ``` ## 安装 ``` pip install scikit-learn pandas matplotlib transformers datasets torch ``` 按顺序运行notebook： 1. `P-G_TF-IDF_LR.ipynb` — 训练和评估TF-IDF + LR 2. `P-G_bert.ipynb` — 微调和评估DistilBERT 3. `P-G_compare_TF-bert.ipynb` — 在未见数据上对比两个模型 ## 局限性 - TF-IDF模型在较小的数据集上训练（440个样本）——对新攻击的泛化能力较弱 - DistilBERT在测试集上有14个误报，主要是带有安全相关语言的角色扮演提示词 - 目前两个模型都无法处理间接提示词注入（嵌入在文档或工具输出中的注入） ## 动机 作为对LLM安全和AI红队持续兴趣的一部分而构建。目标是了解攻击模式——从攻击和防御两方面入手——为企业的LLM部署提供更好的护栏设计。 *Jiraphat Jarernpipattada — 计算机工程，KMUTT*

标签：Adversarial Attack, AI防护, AMSI绕过, Apex, Base64解码, DistilBERT, DLL 劫持, Jailbreak, Logistic Regression, NLP, Prompt Extraction, ROT13解码, TF-IDF, Transformer, 二分类, 人工智能安全, 凭据扫描, 合规性, 大语言模型, 威胁检测, 安全防护, 对抗攻击, 恶意意图检测, 敏感信息检测, 文本分类, 文本预处理, 机器学习, 模型微调, 深度学习, 系统调用监控, 越狱检测, 逆向工具