Adv프롬프터

AdvPrompter: LLM을 위한 빠른 적응형 적대적 프롬프트

이 저장소는 AdvPrompter (arxiv:2404.16873)의 공식 구현입니다.

별표를 남겨주세요. 이 저장소를 방문하고 우리 논문을 인용해 주세요. 우리 작업이 마음에 드신다면(및/또는 사용하신다면) 감사합니다!

0. 설치

• 옵션 1. Singularity를 ​​사용하여 advprompter.def 컨테이너 설치
• 옵션 2. python3.11 및 요구 사항.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에 지정되어 있으며 다양한 옵션이 이미 구현되어 있습니다.

코드베이스는 conf/base.yaml에서 해당 옵션을 설정하여 선택적으로 wandb를 지원합니다.

하드웨어 사양에 대한 참고 사항: 우리가 수행한 모든 실험에서는 두 개의 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

특정 데이터 세트의 TargetLLM에 대해 지정된 AdvPrompter의 성능을 테스트합니다. conf/base.yaml에서 TargetLLM 및 AdvPrompter를 지정해야 합니다. 또한 이전에 AdvPrompter가 미세 조정된 경우 peft_checkpoint에 대한 경로를 지정할 수도 있습니다.

// conf/prompter/llama2.yaml을 참조하세요.
lora_params:
웜스타트: 사실
lora_checkpoint: "path_to_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

TargetLLM에 대해 지정된 AdvPrompter를 교육합니다. 위에서 지정한 평가를 정기적으로 자동 수행하고, AdvPrompter의 중간 버전을 나중에 웜스타트를 위해 ./exp/.../checkpoints 아래의 실행 디렉터리에 저장합니다. 체크포인트는 모델 구성의 lora_checkpoint 매개변수를 사용하여 지정할 수 있습니다(1.1 평가에 설명된 대로). 또한 훈련은 AdvPrompterOpt로 생성된 대상 접미사를 각 시대마다 ./exp/.../suffix_opt_dataset 에 저장합니다. 이를 통해 train.yaml 의 사전 훈련에서 해당 경로를 지정하여 이러한 접미사 데이터 세트에 대한 사전 훈련이 가능합니다.

conf/train.yaml에서 고려해야 할 몇 가지 중요한 하이퍼파라미터: [epochs, lr, top_k, num_chunks, lambda_val]

예

참고: target_llm.llm_params.checkpoint를 로컬 경로로 바꿀 수 있습니다.

• 예 1: Vicuna-7B의 AdvPrompter:

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

• 예 2: Vicuna-13B의 AdvPrompter:

   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: Mistral-7B-chat의 AdvPrompter:

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

• 예 4: Llama2-7B-chat의 AdvPrompter:

   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**

(* = 제1저자 동일, ** = 조언 동일)

3. 라이센스

우리의 소스 코드는 CC-BY-NC 4.0 라이센스를 따릅니다.