DOCSEARCHbyAlgolia
Submit Docs

refusal_direction

repository·main·Indexed 19 days ago

https://github.com/andyrdt/refusal_direction

Code and results to reproduce the research paper 'Refusal in Language Models Is Mediated by a Single Direction'. This repository provides a pipeline to identify, select, and evaluate refusal directions in LLMs, including tools for processing harmful behavior datasets such as AdvBench, MaliciousInstruct, TDC2023, JailbreakBench, HarmBench, StrongReject, and Alpaca.

Tokens
1.3K
Snippets
5
Records
6
Agent score
15%

What's inside refusal_direction

  1. Reproduce main results using the pipeline

    main

    You can reproduce the main results from the paper by running the pipeline.run_pipeline module. You must provide the --model_path argument, which should be the HuggingFace model path (e.g., meta-llama/Meta-Llama-3-8B-Instruct).

    The pipeline executes five sequential steps and saves artifacts in pipeline/runs/{model_alias}/:

    1. Extract candidate refusal directions: Saved to generate_directions.
    2. Select the most effective refusal direction: Saved to select_direction. The final direction is saved as direction.pt.
    3. Generate completions over harmful prompts: Evaluates refusal metrics; saved to completions.
    4. Generate completions over harmless prompts: Evaluates refusal metrics; saved to completions.
    5. Evaluate CE loss metrics: Saved to loss_evals.
    python3 -m pipeline.run_pipeline --model_path {model_path}

  2. Setup the refusal_direction environment

    main

    To use this repository, clone it and run the provided setup script. The setup.sh script will perform the following:

    1. Prompt you for a HuggingFace token (required for gated models).
    2. Prompt you for a Together AI token (required for evaluating jailbreak safety scores via the Together AI API).
    3. Create a virtual environment.
    4. Install all required dependencies.

    Note: This repository contains text that is offensive, harmful, or otherwise inappropriate in nature.

    git clone https://github.com/andyrdt/refusal_direction.git
cd refusal_direction
source setup.sh

  3. Access the minimal demo Colab

    main

    A minimal demo of bypassing refusal is available as a Google Colab notebook. This is useful for a quick, interactive demonstration of the methodology described in the project's blog post.

    https://colab.research.google.com/drive/1a-aQvKC9avdZpdyBn4jgRQFObTPy1JZw

  4. Standardized dataset JSON format

    main

    All processed datasets in this project are stored as JSON files where each entry is a dictionary with the following keys:

    • instruction (string): The text prompt or behavior to be evaluated.
    • category (string or null): The semantic category of the instruction (e.g., from HarmBench or JailbreakBench). If the source dataset does not provide a category, this is set to null.

  5. Construct harmful and harmless dataset splits

    main

    After processing the raw datasets, you can use the following functions to create organized train, validation, and test splits.

    Harmful Dataset Splits

    construct_harmful_dataset_splits() creates three files in the splits/ directory:

    • harmful_train.json: Composed of samples from advbench.json, malicious_instruct.json, and tdc2023.json. It limits each source to a max_train_subset_size of 128 samples to prevent bias.
    • harmful_val.json: Uses harmbench_val.json.
    • harmful_test.json: Composed of jailbreakbench.json, harmbench_test.json, and strongreject.json.

    Note: The function automatically removes duplicate instructions across the train, val, and test sets to ensure data integrity.

    Harmless Dataset Splits

    construct_harmless_dataset_splits() creates three files in the splits/ directory using the alpaca.json dataset:

    • harmless_train.json (60%)
    • harmless_val.json (20%)
    • harmless_test.json (20%)

    It uses a fixed random seed of 42 for reproducibility.

    # Generate the final splits for training and evaluation
construct_harmful_dataset_splits()
construct_harmless_dataset_splits()

  6. Download and process various harmful behavior datasets

    main

    The project provides several helper functions to download raw datasets from GitHub and convert them into a standardized JSON format. Each processed JSON file follows a schema containing instruction and category keys.

    Supported datasets include:

    • AdvBench: Downloads harmful_behaviors.csv and extracts the goal column.
    • MaliciousInstruct: Downloads MaliciousInstruct.txt and treats each line as an instruction.
    • TDC2023: Downloads both dev and test behavior JSON files and merges them.
    • JailbreakBench: Downloads behaviors.csv and extracts Goal and Category.
    • HarmBench: Downloads harmbench_behaviors_text_{split}.csv (where split is val or test). It filters out instructions where FunctionalCategory contains 'copyright' or Tags contains 'context'.
    • StrongReject: Downloads strongreject_dataset.csv and extracts forbidden_prompt and category.
    • Alpaca: Loads the tatsu-lab/alpaca dataset from Hugging Face and filters for instructions that do not have an input string.
    # Example of downloading and processing datasets
download_advbench()
download_malicious_instruct()
download_tdc2023()

download_jailbreakbench()
download_harmbench(split='val')
download_harmbench(split='test')
download_strongreject()

download_alpaca()