role-confusion/prompt-injection-as-role-confusion
GitHub: role-confusion/prompt-injection-as-role-confusion
该研究项目从「角色混淆」视角解释提示词注入攻击的机理，提供线性探针训练、角色空间投影和CoT Forgery攻击评估等工具来量化分析LLM的安全漏洞。
Stars: 72 | Forks: 11
Prompt Injection as Role Confusion
## 概述
LLM 将世界视为单一的文本流，并将其划分为诸如 `
` 或 `` 的*角色*。我们将 **prompt injection** 追溯为**角色混淆**：模型是通过文本的“语气”而非其被标记的角色来感知文本来源的。隐藏在网页中的命令之所以能劫持 agent，仅仅是因为它听起来像 `` 文本，尽管它带有 `` 标签。
本仓库提供：
- **Role probes**：训练并应用线性探针，以衡量模型内部如何感知“谁在说话”
- **CoT Forgery**：在聊天和 agent 设置下，复现这种欺骗模型推理的零样本攻击
- **Role-space 分析**：将 prompt injection 投影到角色空间中，以可视化和量化角色混淆
- **论文复现**：重新生成论文中的所有图表
[快速入门演示](#quickstart) 是最快的上手方式；[完整的复现笔记本](#full-reproduction) 涵盖了论文中的每一个实验。
## ⚡ 快速入门
我们建议从这里开始，而不是克隆整个仓库。
**Role probes 演示** - 训练探针并在真实示例上测试角色混淆：
- 下载并在同一目录下运行 `demo/role-probe-demo.ipynb` 和 `demo/simple_test_helpers.py`
**CoT Forgery 演示** - 在几个示例上运行攻击：
- 下载并运行 `demo/cot-forgery-demo.ipynb`
以下所有说明仅在逐个实验完整复现论文时才需要。
## 🔨 要求
### 硬件
完整实验假定拥有 CUDA GPU。论文实验最初是在 H200 上运行的。较小的本地实验可以在其他 CUDA GPU 上运行，但更大模型的激活导出工作流非常占用内存，可能需要类似的高 VRAM。
激活导出和模型输出也可能需要大量的磁盘空间。对于快速的初步尝试，请在运行完整的复现工作流之前先从 `demo/` 笔记本开始。
### 软件
- Python 3.12+
- CUDA 12.8
- R（可选），用于某些分析和绘图笔记本
安装依赖项：
```
bash setup_python.sh # Change paths in file to local venv
bash setup_r.sh
```
在仓库根目录下创建一个 `.env` 文件。根据您运行的笔记本，您可能需要：
```
OPENROUTER_API_KEY=...
OPENAI_API_KEY=...
HF_TOKEN=...
```
基于 API 的实验可能会产生服务商费用。
## 🔁 完整复现
以下笔记本按实验系列组织。在每个系列中，除非该步骤被标记为可选，否则请按顺序运行笔记本。
1. [Role Space 分析](#1-role-space-analysis)
2. [CoT Forgery 评估](#2-cot-forgery-evaluations)
3. [角色混淆：CoT Forgery](#3-role-confusion-cot-forgery)
4. [角色混淆：标准 Prompt Injection](#4-role-confusion-standard-prompt-injections)
## 1. Role Space 分析
此工作流分析模型的内部角色感知。笔记本和输出是特定于模型的；请在相应的笔记本中设置模型选择。
支持的模型包括：
- `gpt-oss-20b`
- `gpt-oss-120b`
- `Nemotron-3-Nano`
- `Qwen3-30B-A3B`
- `Jamba-Reasoning-3B`
| 步骤 | 笔记本 | 需求 | 主要输出 |
| --- | --- | --- | --- |
| 生成对话数据 | `role-analysis/01-get-conversations-data.ipynb` | OpenRouter 或本地模型访问权限 | `convs/{model_name}.csv` |
| 训练和评估 role probes | `role-analysis/02-train-role-probes.ipynb` | 对话数据 | `outputs/probes/{model_name}.pkl`, `outputs/probe-training/*.csv`, `outputs/probe-projections/*.csv` |
| 绘制对话投影 | `role-analysis/03-analyze-probes.ipynb` | 训练好的探针和投影 | `role-analysis/plots/*` |
| 绘制园艺对话投影 | `role-analysis/04-tomato-probe-results.ipynb` | 训练好的探针和投影 | `role-analysis/plots/*` |
第一个笔记本从 `toxicchat` 和 `oasst` 构建特定于模型的对话数据，并通过 OpenRouter 或本地后备方案重新生成模型响应。
第二个笔记本在由标签引发的角色几何结构上训练 role probes，并将普通对话和受控示例投影到角色空间中。
## 2. CoT Forgery 评估
此工作流在聊天和 agent 设置下运行 CoT Forgery。聊天笔记本评估越狱式的 prompt；agent 笔记本评估在 ReAct 式工具循环中的 prompt injection。
### 2.1 聊天越狱
聊天笔记本生成 CoT Forgery prompt，运行本地和闭源权重模型的生成，对攻击成功与否进行分类，并绘制最终的攻击成功率。
| 步骤 | 笔记本 | 需求 | 主要输出 |
| --- | --- | --- | --- |
| 生成 CoT Forgery prompt | `cot-forgery-chat-evals/01-generate-forgeries.ipynb` | `OPENROUTER_API_KEY` | `base-harmful-policies.csv` |
| 运行本地模型生成 | `cot-forgery-chat-evals/02-export-jailbreak-generations.ipynb` | 已生成的 prompt，本地模型访问权限 | `base-harmful-responses-classified.csv` |
| 运行闭源模型生成 | `cot-forgery-chat-evals/03-run-openrouter-generations.ipynb` | 已生成的 prompt，`OPENROUTER_API_KEY` | `openrouter-generations/harmful-responses-classified.csv` |
| 绘制结果 | `cot-forgery-chat-evals/04-plot-jailbreak-stats.ipynb` | 已分类的生成输出 | `cot-forgery-chat-evals/plots/*` |
第一个笔记本通过 OpenRouter 调用 LLM，为 StrongREJECT 中的有害问题生成 CoT Forgery prompt 和比较基线。生成笔记本在本地 `gpt-oss-*` 模型和闭源权重模型上评估这些 prompt。分类步骤通过 OpenRouter 使用 LLM 分类器。
### 2.2 Agent Prompt Injection
agent 笔记本评估在 ReAct 式工具使用循环中的 CoT Forgery 式 prompt injection。
| 步骤 | 笔记本 | 需求 | 主要输出 |
| --- | --- | --- | --- |
| 运行本地 agent 评估 | `cot-forgery-agent-evals/01-run-injections-gpt-oss.ipynb` | 本地 `gpt-oss-*` 模型访问权限 | `local-agent-outputs-{model_name}-classified.csv` |
| 运行托管 agent 评估 | `cot-forgery-agent-evals/02-run-injections-openai.ipynb` | 托管模型 API 访问权限 | `api-agents-output-classified.csv` |
| 绘制结果 | `cot-forgery-agent-evals/03-plot-injections.ipynb` | 已分类的 agent 输出 | `cot-forgery-agent-evals/plots/*` |
输出包括完整的 agent 循环记录和最终的成功分类。
## 3. 角色混淆：CoT Forgery
此工作流使用在 role-space 工作流中训练的探针来分析来自聊天和 agent 实验的 CoT Forgery 激活，展示样式化/去样式化如何影响角色混淆，以及角色混淆如何预测攻击的成功。
| 步骤 | 笔记本 | 需求 | 主要输出 |
| --- | --- | --- | --- |
| 导出聊天攻击激活 | `cot-forgery-role-confusion/02-export-user-injection-activations.ipynb` | `cot-forgery-chat-evals/02-export-jailbreak-generations.ipynb` | `activations-redteam/{model_name}` |
| 导出 agent 攻击激活 | `cot-forgery-role-confusion/01-export-agent-activations.ipynb` | `cot-forgery-agent-evals/01-run-injections-gpt-oss.ipynb` | `activations-agent/{model_name}` |
| 将攻击投影到角色空间 | `cot-forgery-role-confusion/03-project-role-probes.ipynb` | 训练好的探针和导出的激活 | `cot-forgery-role-confusion/exports/*` |
| 绘制聊天和 agent 角色分析 | `cot-forgery-role-confusion/04-plot-injection-probe-results.ipynb`, `cot-forgery-role-confusion/05-plot-agent-probe-results.ipynb` | 角色投影导出 | `cot-forgery-role-confusion/plots/*` |
如果您只关心聊天攻击，可以跳过 agent 激活导出和 agent 绘图笔记本。投影笔记本也包含一个 agent 分析部分，当 agent 激活不可用时可以跳过该部分。
## 4. 角色混淆：标准 Prompt Injection
此工作流分析角色感知如何预测标准 agent prompt injection 的成功，而不是 CoT Forgery。这对应于论文中标准的 prompt injection 角色混淆分析，我们在其中外生改变 Userness，并展示角色混淆能够预测典型 agent prompt injection 中攻击的成功。
| 步骤 | 笔记本 | 需求 | 主要输出 |
| --- | --- | --- | --- |
| 运行 prompt injection 变体并投影 Userness | `agent-injections/01-run-user-injections-gpt-oss.ipynb` | 训练好的 role probes | `outputs/agent-outputs-classified-{model_name}.csv` |
| 绘制结果 | `agent-injections/02-analyze-injections.ipynb` | 带有 Userness 的已分类输出 | `outputs/plots/*` |
第一个笔记本创建 prompt injection 变体，运行 `gpt-oss-*` 的 ReAct 循环，对 agent 的有害级别进行分类，并提取每个注入查询的平均 Userness。
## 引用
```
@inproceedings{ye2026promptinjectionroleconfusion,
title = {Prompt Injection as Role Confusion},
author = {Ye, Charles and Cui, Jasmine and Hadfield-Menell, Dylan},
booktitle = {International Conference on Machine Learning (ICML)},
year = {2026},
url = {https://arxiv.org/abs/2603.12277}
}
```
## 许可证
有关许可证信息，请参阅 `LICENSE`。标签：AI安全, Apex, Chat Copilot, DLL 劫持, IaC 扫描, NoSQL, Petitpotam, Vectored Exception Handling, 凭据扫描, 可解释性, 大语言模型, 机器学习, 逆向工具