AdvPrompter：针对大语言模型的快速自适应对抗性提示

本仓库是 AdvPrompter (arxiv:2404.16873) 的官方实现。

如果您喜欢（和/或使用）我们的工作，请为本仓库点 🌟星🌟 并引用我们的论文 📜，谢谢！

0. 安装

选项 1. 使用 Singularity 安装 advprompter.def 容器
选项 2. 手动安装 python3.11 和 requirements.txt：

  conda create -n advprompter python=3.11.4
  conda activate advprompter
  pip install -r requirements.txt

1. 运行 AdvPrompter

我们使用 hydra 作为配置管理工具。主要配置文件：./conf/{train,eval,eval_suffix_dataset,base}.yaml AdvPrompter 和 TargetLLM 在 conf/base.yaml 中指定，已实现多种选项。

代码库可选择性地支持 wandb，通过在 conf/base.yaml 中设置相应选项。

硬件规格说明：我们进行的所有实验都使用了两个 NVIDIA A100 GPU，一个用于 AdvPrompter，一个用于 TargetLLM。您可以在 conf/target_llm/base_target_llm.yaml 和 conf/prompter/base_prompter.yaml 中管理设备。

1.1 评估

运行

python3 main.py --config-name=eval

以测试指定 AdvPrompter 在给定数据集上对抗 TargetLLM 的性能。您需要在 conf/base.yaml 中指定 TargetLLM 和 AdvPrompter。此外，如果 AdvPrompter 之前经过微调，您可能需要指定 peft_checkpoint 的路径：

// 参见 conf/prompter/llama2.yaml
lora_params:
warmstart: true
lora_checkpoint: "peft_checkpoint_路径"

评估过程中生成的后缀将保存到运行目录下的 ./exp/.../suffix_dataset 中的新数据集，以便后续使用。这样的数据集也可用于评估基线或手工制作的后缀对 TargetLLM 的效果，可以通过运行以下命令进行评估：

python3 main.py --config-name=eval_suffix_dataset

在填充 eval_suffix_dataset.yaml 中的 suffix_dataset_pth_dct 后。

1.2. 训练

运行

python3 main.py --config-name=train

以训练指定的 AdvPrompter 对抗 TargetLLM。它会自动定期执行上述评估，并将 AdvPrompter 的中间版本保存到运行目录下的 ./exp/.../checkpoints 中，以便后续热启动。可以通过模型配置中的 lora_checkpoint 参数指定检查点（如 1.1 评估中所示）。训练过程还会将每个 epoch 使用 AdvPrompterOpt 生成的目标后缀保存到 ./exp/.../suffix_opt_dataset。这允许通过在 train.yaml 的 pretrain 下指定相应路径来对这样的后缀数据集进行预训练。

conf/train.yaml 中需要考虑的一些重要超参数：[epochs, lr, top_k, num_chunks, lambda_val]

示例

注意：您可能需要将 target_llm.llm_params.checkpoint 替换为本地路径。

示例 1： AdvPrompter 在 Vicuna-7B 上：

 python3 main.py --config-name=train target_llm=vicuna_chat target_llm.llm_params.model_name=vicuna-7b-v1.5

示例 2： AdvPrompter 在 Vicuna-13B 上：

 python3 main.py --config-name=train target_llm=vicuna_chat target_llm.llm_params.model_name=vicuna-13b-v1.5 target_llm.llm_params.checkpoint=lmsys/vicuna-13b-v1.5 train.q_params.num_chunks=2

示例 3： AdvPrompter 在 Mistral-7B-chat 上：

 python3 main.py --config-name=train target_llm=mistral_chat

示例 4： AdvPrompter 在 Llama2-7B-chat 上：

 python3 main.py --config-name=train target_llm=llama2_chat train.q_params.lambda_val=150

2. 贡献者

Anselm Paulus*, Arman Zharmagambetov*, Chuan Guo, Brandon Amos**, Yuandong Tian**

(* = 共同第一作者，** = 共同指导)

3. 许可证

我们的源代码采用 CC-BY-NC 4.0 许可证。