abedegno/prompt-injection-lora

# prompt-injection-lora 一个基于 **Qwen3-0.6B** 的 LoRA 微调模型，用于 prompt 注入检测。使用完全公开的训练数据，并在标准的 900 条提示 `aiguard-lab` 语料库上进行了评估。 - **Hugging Face 上的训练适配器：** - **完整文章（四篇博客系列）：** ## 核心结果 在 [`aiguard-lab`](https://github.com/abedegno/aidr-aiguard-lab) 的 900 条提示 PI（提示注入）子集上： | 模型 | 大小 | F1 | |---|---:|---:| | AIDR (闭源参考，PA4002) | N/A | **0.979** | | **本方案** (Qwen3-0.6B + LoRA + class-tokens) | 600M | **0.9053** | | Sentinel v1 (ModernBERT-large) | 395M | 0.879 | | Vijil DOME (ModernBERT-base) | ~150M | 0.860 | | ProtectAI DeBERTa V2 | 184M | 0.714 | | Meta Prompt-Guard-2 (86M / 22M) | 86M / 22M | ~0.69 | 通过 LoRA + class-tokens 微调，该方案在此语料库上的表现超越了所有已发表的开源分类器，与闭源的当前最佳水平仅相差不到 0.074。它包含约 10M 可训练参数（约占 Qwen3-0.6B 的 1.7%），适配器仅几兆大小，并且无需自定义分类器头即可通过任何 chat-completions 端点提供服务。 每条提示的评估 CSV 位于 [`results/finetune-eval.csv`](results/finetune-eval.csv)。 ## 目录结构 ``` prompt-injection-lora/ ├── scripts/ │ ├── prep_dataset_dome.py # starter mix: WildGuardMix + safe-guard (~95k rows) │ ├── prep_dataset.py # final mix: + SPML / alespalla / verazuo (~111k rows) │ ├── audit_dataset.py # schema / dedup / length sanity check │ ├── eval_baseline.py # eval HF text-classification models (DOME, ProtectAI, etc.) │ ├── eval_finetune_local.py # local eval via transformers + peft │ └── train_modal.py # Modal H100 training entrypoint (Qwen3.5-0.8B base-swap test) ├── results/ │ └── finetune-eval.csv # per-prompt scores from the headline eval └── data/ # gitignored - regenerate via the prep scripts ``` ## 复现 训练数据（约 310 MB）已被 gitignore。可从公开数据源重新生成： ``` # 0. 安装 dev env uv sync # 1. 将 aiguard-lab corpus 克隆为 eval corpus 的同级目录 cd .. && git clone https://github.com/abedegno/aidr-aiguard-lab.git aiguard-lab && cd - # 2. 构建最终的 training corpus（约 11.1 万行；从 Hugging Face + GitHub 拉取） uv run scripts/prep_dataset.py # 3. Audit uv run scripts/audit_dataset.py # 4. 在 Fireworks 上进行 Fine-tune（博客详细介绍了此步骤） firectl supervised-fine-tuning-job create \ --base-model accounts/fireworks/models/qwen3-0p6b \ --dataset pi-final-mix \ --output-model pi-injection-qwen3-0p6b \ --epochs 3 --lora-rank 8 --learning-rate 1e-4 --batch-size-samples 32 \ --learning-rate-warmup-steps 100 # 5. Pull adapter 并在本地评估 uv run scripts/eval_finetune_local.py \ --base-model Qwen/Qwen3-0.6B \ --adapter-dir \ --out results/finetune-eval.csv # 6. （可选）在相同 corpus 上重新测量 open-source baselines uv run scripts/eval_baseline.py --model vijil/vijil_dome_prompt_injection_detection \ --tokenizer answerdotai/ModernBERT-base --injection-label 1 --threshold 0.9 ``` ## 方法概述（一段话） Decoder LLM 没有二分类输出头；其输出是包含约 15 万个 token 的完整词汇表上的分布。这里的技巧是：使用 SFT（监督微调）训练模型，使其在接收“你是一个安全分类器”的提示后，下一个 token 的概率分布会直接坍缩到模型在预训练中已经学过的两个特定 token（`unsafe` 和 `safe`）上。置信度来源于对这两个 token 的 logits 进行的二分类 softmax。结合 LoRA（约 10M 可训练参数，约占基础模型的 1.7%），该方案体积小巧、成本低廉，并且无需自定义分类器头就能通过任何 chat-completions 端点提供服务。完整的分步说明在该博客系列的第 1 部分：。 ## 许可证 Apache 2.0。详见 [`LICENSE`](LICENSE)。

标签：AI安全, AI模型评估, AMSI绕过, Chat Copilot, DLL 劫持, Hugging Face, LoRA微调, NLP, PEFT, Qwen3, SML/LLM防护, Transformer, 凭据扫描, 大语言模型, 威胁检测, 对抗性机器学习, 开源模型, 恶意提示词检测, 文本分类, 深度学习, 网络安全, 逆向工具, 隐私保护