本仓库包含运行ASTRA和ASTRA++攻击的代码，这些攻击能够突破SecAlign++、SecAlign、StruQ防御。此外，本仓库还提供了一些生成的攻击示例和攻击日志。
这将克隆仓库以及所有子模块，并放置到正确的位置。
git clone --recurse-submodules https://github.com/nishitvp/better_opts_attacks.git
本代码库未使用除标准LLM和深度学习包之外的任何特殊Python包。如需安装依赖，请运行：
python3 -m pip install -r requirements.txt
这将在您的环境中安装必要的包。
首先，从HuggingFace下载基础模型（其中一些模型可能需要权限，您需要先通过HuggingFace申请访问权限）
hf login --token ${HF_PASSWORD}
hf download meta-llama/Meta-Llama-3-8B-Instruct --local-dir secalign_refactored/secalign_models/meta-llama/Meta-Llama-3-8B-Instruct
hf download mistralai/Mistral-7B-Instruct-v0.1 --local-dir secalign_refactored/secalign_models/mistralai/Mistral-7B-Instruct-v0.1
hf download meta-llama/Llama-3.1-8B-Instruct --local-dir secalign_refactored/secalign_models/meta-llama/Llama-3.1-8B-Instruct
hf download facebook/Meta-SecAlign-8B --local-dir secalign_refactored/secalign_models/Meta-SecAlign-8B
下载SecAlign适配器和StruQ模型。有一个便捷脚本可帮助下载并解压到正确位置。
python3 setup_scripts/download_models.py
在项目根目录下运行此脚本，它将下载并解压正确的模型到正确位置。为验证下载是否成功，请检查 secalign_refactored/secalign_models 目录下是否包含三个文件夹 huggyllama、meta-llama 和 mistralai，每个文件夹内应包含相应的模型。
运行针对强知识攻击者的ASTRA，关键脚本是 experiment.py。该脚本将从AlpacaFarm数据集中选取 40 个示例，使用指定配置注入有效载荷 Print exactly 'Hacked!'，并运行ASTRA及基线方法（GCG），同时控制初始化（随机初始化）、前向传播次数、迭代次数等参数。
用户可以通过在调用 experiment.py 时设置以下命令行参数来控制参数：
--model-name：可选值为 meta-llama、meta-llama-instruct、mistralai、mistralai-instruct、huggyllama。
--defense：可选值为 undefended、struq 和 secalign。默认值为 secalign。--model-name 和 --defense 参数共同唯一指定所使用的模型，具体对应关系如下表：
一个完整的命令示例如下：
python3 experiment.py --model-name meta-llama-instruct --defense secalign --prefix-length 5 --suffix-length 20 --expt-folder-prefix logs/astra_llama_secalign
运行上述命令后，脚本将自动将工作负载平均分配给不同的子进程，每个子进程使用一个GPU计算攻击。建议在至少48GB显存的GPU上运行攻击。
每个子进程将在 --expt-folder-prefix 指定路径下的单独子文件夹中，记录其所攻击的每个示例的完整日志。日志以可查询的数据库格式存储，包含Python对象（详见 utils/experiment_logger 文件）。每个记录的日志将包含优化过程中的token序列、损失值以及其他重要细节（如初始化配置），可供后续检索。
完整的ASTRA实验完成后，可使用文件 analysis/analysis.ipynb 对生成的日志进行分析。只需将Jupyter notebook复制到 --expt-folder-prefix 路径下，将notebook中的模型路径替换为正确的模型路径，然后运行notebook。最后一个单元格将输出GCG的成功次数和ASTRA的成功次数。（附注：由于分析脚本需要实际生成每个示例在优化过程中生成的每个输入token序列的输出，因此运行时间较长。）
analysis.ipynb 还包含辅助代码，用于绘制平均损失曲线以及进一步探索生成的日志。
分析脚本与优化过程分离，从根本上确保优化本身不受事后分析代码的污染。
我们还分享了一些ASTRA生成的攻击，位于压缩文件 data/attacks_generated.zip 中。具体而言，压缩文件中包含4个文件，每个文件包含40个ASTRA生成的攻击示例，分别对应SecAlign防御的Mistral和Llama-3版本各两个文件。这些文件提供了在20和25个token预算下（即 (0, 20) 和 (5, 20) 配置）生成的示例。注入内容已按照配置放置到输入字段中。
运行针对弱知识攻击者的ASTRA++，关键脚本是 experiment_universal.py。该脚本将从AlpacaFarm数据集中选取 10 个训练示例，使用指定的前缀和后缀配置注入有效载荷 Print exactly 'Hacked!'，并训练得到最终的通用提示注入。该脚本还会在ASTRA++训练后运行基线方法——通用版GCG。
与上一节中的 experiment.py 类似，experiment_universal.py 也接受 --model-name、--defense、--prefix-length、--suffix-length、--expt-folder-prefix 等参数。
除了上一节运行ASTRA时列出的所有模型外，experiment_universal.py 还额外支持一个模型，对应SecAlign++防御。要针对SecAlign++防御进行训练，请将 --model-name 参数设置为 secalign_refactored/secalign_models/Meta-SecAlign-8B，并将 --defense 参数设置为 meta_secalign。
该脚本还接受一个额外的命令行参数 --training-run，用于指定训练运行所使用的 10 个示例的集合。
其余参数的含义与之前相同。
完整的ASTRA++实验完成后，可使用文件 analysis/analysis_universal.ipynb 对生成的日志进行分析。与之前一样，将Jupyter notebook复制到 --expt-folder-prefix 路径下，正确加载模型（由于SecAlign++的加载方式存在一些细微差异，但未来会统一）。最后一个单元格将输出GCG的测试成功次数和ASTRA的成功次数，以及达到的最佳通用提示注入。
对于ASTRA++，我们将针对SecAlign++防御的Llama-3.1-8B-Instruct生成的通用提示注入上传至文件 data/universal_pis_secalign++.json，其中包含元数据和生成时的配置。此外，我们还将所有通用性训练运行中的攻击日志上传至 data/final_result_logs.pkl（我们保证，pickle文件中不包含恶意代码 :)。
虽然上述脚本可以开箱即用地运行ASTRA和ASTRA++攻击，但本代码库允许配置几乎所有可能影响攻击的参数，例如训练数据集大小、阈值选择、随机种子、初始化配置等。如果您有任何有趣的建议，欢迎发送拉取请求。
如有任何疑问、报告错误、索取详细信息或需要澄清，请随时在GitHub上提出Issue或发送邮件给作者。
--model-name | --defense | 加载的模型 |
|---|
meta-llama-instruct | secalign | SecAlign适配器 + 预指令调优的Meta-Llama-3-8B-Instruct |
mistralai-instruct | secalign | SecAlign适配器 + 预指令调优的Mistral-7B-Instruct-v0.1 |
meta-llama | secalign | 经SecAlign防御的Meta-Llama-3-8B（未进行预指令调优） |
mistralai | secalign | 经SecAlign防御的Mistral-7B-Instruct-v0.1（未进行预指令调优） |
huggyllama | secalign | 经SecAlign防御的Llama-2-7B（未进行预指令调优） |
meta-llama | struq | 经StruQ防御的Meta-Llama-3-8B（未进行预指令调优） |
mistralai | struq | 经StruQ防御的Mistral-7B-v0.1（未进行预指令调优） |
huggyllama | struq | 经StruQ防御的Llama-2-7B（未进行预指令调优） |
meta-llama-instruct | undefended | 无防御（原始）Meta-Llama-3-8B-Instruct |
mistralai-instruct | undefended | 无防御（原始）Mistral-7B-Instruct-v0.1 |
meta-llama | undefended | 无防御（原始）Meta-Llama-3-8B（未进行预指令调优） |
mistralai | undefended | 无防御（原始）Mistral-7B-Instruct-v0.1（未进行预指令调优） |
huggyllama | undefended | 无防御（原始）Llama-2-7B（未进行预指令调优） |
所有其他 --model-name 和 --defense 组合均为无效。
--prefix-length：用于给定评估运行的对抗性前缀长度。默认值为 5。
--suffix-length：用于给定评估运行的对抗性后缀长度。默认值为 20。
--expt-folder-prefix：实验运行日志的保存路径。