TAP

TAP — Tree of Attacks with Pruning (reproduction)

中文文档 →

A faithful, self-contained reproduction of the jailbreak method from Mehrotra et al., “Tree of Attacks: Jailbreaking Black-Box LLMs Automatically” (NeurIPS 2024).

The system prompts are reproduced verbatim from the paper (Tables 5/6/7). This configuration uses open-source models for the attacker and target, and Qwen3Guard as the evaluator.

Research use only. This code is for safety/red-teaming research and reproducing a published baseline. Do not use it to produce or distribute harmful content.

Algorithm (paper Algorithm 1)

Each iteration (up to depth rounds):

1. Branch — the attacker LLM generates branching_factor variations per leaf.
2. Prune-1 — the evaluator removes off-topic variations (skipped for the Qwen3Guard evaluator, matching the guard-style evaluator in the paper’s Appendix C.4).
3. Attack & Assess — query the target with each prompt; the evaluator scores the responses. A jailbroken response ends the search.
4. Prune-2 — keep the top-width leaves by score for the next iteration.

Paper defaults: branching_factor=4, width=10, depth=10.

Evaluator (Qwen3Guard)

With the default --evaluator-type guard, Qwen3Guard labels the target’s response as one of Safe / Unsafe / Controversial:

• Safe → not jailbroken (score 1)
• Unsafe / Controversial → jailbroken (score 10)

Any branch judged jailbroken returns success immediately.

You can switch to the paper’s graded LLM judge with --evaluator-type llm --evaluator-model gpt-4 (needs an OpenAI key, or a local model name). That mode enables 1–10 scoring and Off-Topic (Phase-1) pruning.

Install

pip install -r requirements.txt

Models load via modelscope (falling back to transformers).

Data format

Datasets live in ./Dataset/ (Adv.csv / Harm.csv / JBB.csv). Goals are read from the first column starting at the second row (header skipped). A second column, if present, is used as the target’s starting string (Sure, here is ...); otherwise one is synthesized. Input and output paths are both CLI arguments.

Run

Non-thinking (default)

python tap_main.py \
    --input  ./Dataset/Adv.csv \
    --output ./results/adv_nothink.json \
    --GPU 1,2,4 \
    --attack-model-path /path/to/models/Qwen/Qwen3-8B \
    --target-model-path /path/to/models/Qwen/Qwen3-8B \
    --guard-path        /path/to/models/Qwen/Qwen3Guard-Gen-8B

Thinking mode (give the target 2 cards)

python tap_main.py \
    --input  ./Dataset/Adv.csv \
    --output ./results/adv_think.json \
    --attack-gpu 1 --target-gpu 2,5 --guard-gpu 4 \
    --target-model-path /path/to/models/Qwen/Qwen3-8B \
    --guard-path        /path/to/models/Qwen/Qwen3Guard-Gen-4B \
    --enable-thinking

GPU convention (one model per card, no contention)

• --GPU 1,2,4: attacker → card 1, target → card 2, evaluator → card 4 (one model per card; enough for non-thinking).
• Give a single role multiple cards with --attack-gpu / --target-gpu / --guard-gpu. Thinking mode (32768 tokens, large KV cache) → give the target 2 cards (--target-gpu 2,5).
• Large branches (e.g. 40 attacker prompts at once) are split into chunks of --max-batch (default 8). Lower it to 4/2/1 on OOM.

Useful options

| Flag | Default | Meaning | |——|———|———| | --resume | off | Resume from an existing --output: skip completed goals, continue | | --limit | 0 | >0: only process the first N goals (quick test) | | --branching-factor | 4 | branching factor b | | --width | 10 | max leaves kept per round w | | --depth | 10 | max iterations d | | --enable-thinking | off | enable the target’s thinking mode (Qwen3) | | --target-max-tokens | 32768 | target generation length (leaves room for thinking) | | --max-batch | 8 | max sequences per forward pass (lower on OOM) | | --evaluator-type | guard | guard=Qwen3Guard; llm=graded judge |

Run python tap_main.py --help for the full list.

Output

A JSON file written incrementally (atomic write; safe to inspect/resume mid-run):

• config: models and hyper-parameters of the run.
• summary: attack_success_rate (ASR), mean_queries, mean_iterations, mean_time_sec.
• results: per goal — success, final_score, final_prompt, final_response, final_thinking, num_queries, iterations, elapsed_sec.

final_prompt is the adversarial prompt TAP converged on (the jailbreaking one on success, or the highest-scoring one otherwise).

Resume

Interrupted? Re-run the same command with --resume (keep --input/--output unchanged). Completed goals are skipped and the summary is recomputed over all results.

Files

• tap_main.py — entry point: args, data IO, TAP tree search, result writing.
• conversers.py — the three roles: AttackLM / TargetLM / GuardEvaluator (+LLMEvaluator).
• language_models.py — HuggingFaceChat (local, one model per card) / OpenAIModel (optional gpt-*).
• system_prompts.py — paper Tables 5/6/7, verbatim.
• common.py — JSON / score parsing, conversation helpers.
• tap_animation.html — standalone animated explainer of the attack flow.

Citation

@inproceedings{mehrotra2024tree,
  title={Tree of Attacks: Jailbreaking Black-Box LLMs Automatically},
  author={Mehrotra, Anay and Zampetakis, Manolis and Kassianik, Paul and Nelson, Blaine and Anderson, Hyrum and Singer, Yaron and Karbasi, Amin},
  booktitle={NeurIPS},
  year={2024}
}