DOCSEARCHbyAlgolia
Submit Docs

AdvPrompter: Fast Adaptive Adversarial Prompting for LLMs

repository·main·Indexed 41 days ago

https://github.com/facebookresearch/advprompter

An implementation of Fast Adaptive Adversarial Prompting designed to generate adversarial suffixes to test the robustness of Large Language Models (LLMs). It features a Hydra-based configuration system for managing training and evaluation workflows, support for PEFT checkpoints, and integration with Weights & Biases for metric tracking.

Tokens
2.2K
Snippets
5
Records
11
Agent score
2%

What's inside AdvPrompter

  1. Configure AdvPrompter using Hydra

    main

    AdvPrompter uses Hydra for configuration management. All settings are managed via YAML files located in the ./conf/ directory.

    Key configuration files:

    • ./conf/base.yaml: Specifies the AdvPrompter and the TargetLLM.
    • ./conf/train.yaml: Contains training hyperparameters like epochs, lr, top_k, num_chunks, and lambda_val.
    • ./conf/eval.yaml: Configuration for evaluation tasks.
    • ./conf/eval_suffix_dataset.yaml: Configuration for evaluating specific suffix datasets.

    Hardware Management: Experiments typically require two NVIDIA A100 GPUs (one for AdvPrompter and one for TargetLLM). You can manage device assignments in:

    • conf/target_llm/base_target_llm.yaml
    • conf/prompter/base_prompter.yaml

    Weights & Biases (wandb): Support for wandb is optional and can be enabled by setting the corresponding options in conf/base.yaml.

  2. Evaluate AdvPrompter against a TargetLLM

    main

    To test the performance of a specified AdvPrompter against a TargetLLM on a given dataset, run the eval config.

    Requirements:

    • You must specify the TargetLLM and AdvPrompter in conf/base.yaml.
    • If the AdvPrompter was previously finetuned, you must provide the path to the PEFT checkpoint using the following parameters in the model config (e.g., conf/prompter/llama2.yaml):
      • lora_params.warmstart: true
      • lora_params.lora_checkpoint: "path_to_peft_checkpoint"

    Outputs:

    • Suffixes generated during evaluation are saved to ./exp/.../suffix_dataset for later use.
    python3 main.py --config-name=eval

  3. Train AdvPrompter

    main

    To train an AdvPrompter against a TargetLLM, run the train config. The process automatically performs evaluations at regular intervals and saves intermediate AdvPrompter versions to ./exp/.../checkpoints for later warmstarting.

    Key Features:

    • Warmstarting: Use the lora_checkpoint parameter in model configs to load a previous checkpoint.
    • Pretraining: Training saves target suffixes to ./exp/.../suffix_opt_dataset. You can use these for pretraining by specifying the path under pretrain in train.yaml.
    • Hyperparameters: Important tuning parameters are located in conf/train.yaml: epochs, lr, top_k, num_chunks, and lambda_val.
    python3 main.py --config-name=train

  4. Install AdvPrompter

    main

    You can install AdvPrompter using one of two methods:

    1. Singularity: Use the provided advprompter.def container.
    2. Manual Python Installation: Create a Conda environment with Python 3.11.4 and install the dependencies from requirements.txt.
    conda create -n advprompter python=3.11.4
conda activate advprompter
pip install -r requirements.txt

  5. Evaluate a specific suffix dataset

    main

    You can evaluate baselines or hand-crafted suffixes against a TargetLLM by using the eval_suffix_dataset configuration. This requires populating the suffix_dataset_pth_dct key in eval_suffix_dataset.yaml with the path to your dataset.

    python3 main.py --config-name=eval_suffix_dataset

  6. The Workspace class manages the AdvPrompter lifecycle

    main

    The Workspace class is the central orchestrator for AdvPrompter. It initializes the models, manages logging (including WandB), and handles the high-level execution of training and evaluation loops.

    Key responsibilities:

    • Model Initialization: Loads the prompter (the adversarial model) and the target_llm (the model being attacked) using the LLM class.
    • Logging: Initializes wandb tables for tracking training and evaluation metrics if cfg.wandb_params.enable_wandb is true.
    • Workflow Execution: Provides high-level methods like train(), pretrain(), and eval() that implement the core logic of the project.
    • Data Management: Handles reading test/affirmative prefixes and saving suffix datasets to CSV files.

  7. Understand the AdvPrompter evaluation workflow

    main

    The evaluation workflow in Workspace.eval() is designed to assess the prompter's ability to generate adversarial suffixes across different settings.

    1. Suffix Dataset Generation: The generate_suffix_datasets() method iterates through datasets defined in cfg.eval.data.dataset_pth_dct. For each dataset, it generates suffixes using the prompter at various token lengths (defined in cfg.eval.prompter.max_new_tokens_list) and saves them as CSV files.
    2. Suffix Dataset Evaluation: The eval_suffix_datasets() method then loads these generated CSV files and evaluates them using eval_suffix_dataset(). This involves:
      • Loading the instruct, target, suffix, and full_instruct sequences.
      • Running evaluate_prompt to get teacher-forced and autoregressive results.
      • Checking for jailbreaks using check_jailbroken against test_prefixes.
      • Calculating metrics like jailbreak success rate, perplexity (PPL), and hit rate @ N.
      • Logging results to WandB and printing summaries to the console.

  8. Understand the AdvPrompter training workflow

    main

    The training workflow in Workspace.train() follows these steps:

    1. Initial Evaluation (Optional): If cfg.train.do_initial_eval is set, an evaluation is run before training starts.
    2. Pretraining (Optional): If cfg.pretrain.enable is true, the pretrain() method is called to warmstart the prompter.
    3. Main Training Loop: Iterates through epochs. In each epoch:
      • Generates an initial suffix using the prompter's autoregressive capabilities.
      • Uses advPrompterOpt to optimize that suffix against the target LLM.
      • Evaluates the optimized suffix using evaluate_prompt.
      • Adds successful or high-priority samples to a ReplayBuffer using add_to_replay_buffer.
      • Performs finetune_prompter(): Samples from the ReplayBuffer and updates the prompter's weights via finetune_prompter_step.
    4. Model Saving: Saves the prompter model to cfg.train.model_save_dir at specified intervals or at the end of training.
    5. Final Evaluation: Runs a final evaluation after training completes.

  9. Training Examples

    main

    The following examples demonstrate how to run training with different target models and hyperparameter overrides. Note that target_llm.llm_params.checkpoint may need to be replaced with a local path if not using a remote repository.

    Example 1: AdvPrompter on Vicuna-7B

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

    Example 2: AdvPrompter on Vicuna-13B with chunking override

    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

    Example 3: AdvPrompter on Mistral-7B-chat

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

    Example 4: AdvPrompter on Llama2-7B-chat with lambda_val override

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

  10. Reference: Workspace training configuration keys

    main

    When configuring the train mode, the following keys are used within the Hydra config:

    KeyDescription
    cfg.train.epochsTotal number of training epochs
    cfg.train.dataset_pthPath to the training dataset
    cfg.train.batch_sizeBatch size for training
    cfg.train.model_save_dirDirectory where trained prompter models are saved
    cfg.train.eval_everyFrequency (in epochs) to run evaluation during training
    cfg.train.replay_buffer.sizeMaximum capacity of the prioritized replay buffer
    cfg.train.replay_buffer.priority_alphaAlpha parameter for the PrioritizedSampler
    cfg.train.replay_buffer.priority_factor.loss_deltaMultiplier for the loss difference in priority calculation
    cfg.train.replay_buffer.priority_factor.jailbreakingWeight added to priority if a jailbreak is detected
    cfg.train.prompter_optim_paramsDictionary of parameters passed to the Adam optimizer for the prompter

  11. Run AdvPrompter via CLI

    main

    AdvPrompter is executed using a Hydra-based CLI entrypoint. The behavior of the run is determined by the mode configuration key. You can trigger different workflows by setting the mode in your Hydra configuration file.

    Supported modes:

    • train: Executes the training workflow, which may include optional pretraining.
    • eval: Executes the evaluation workflow, generating suffix datasets and evaluating them.
    • eval_suffix_dataset: Evaluates existing suffix datasets specified in the cfg.eval.suffix_dataset_pth_dct configuration.

    To run the script, ensure you have your configuration files prepared in the conf directory.

    # Example of how the CLI is invoked (assuming hydra is configured)
python main.py mode=train
python main.py mode=eval
python main.py mode=eval_suffix_dataset