Skip to content

Repository files navigation

AB-GCG: Does Smarter Position Selection Improve GCG?

TL;DR: We replace GCG's all-coordinates competition with four different position pre-commitment strategies — two attention-based heuristics and two contextual bandits trained via REINFORCE. All four make GCG dramatically worse (32–50pp ASR drop). A mechanistic decomposition reveals 60–67% of the gap comes from optimization instability, not search quality.

📄 Read the full write-up → I Tried Four Smarter Ways to Select Positions in GCG. They All Made It Worse — and Here’s Why That’s Interesting


What This Project Does

GCG (Greedy Coordinate Gradient) is the most prominent white-box adversarial attack against safety-aligned LLMs. At each optimization step, GCG samples 512 candidate token replacements across all suffix positions and keeps the best one via argmin — an implicit all-coordinates competition where positions compete for the replacement slot.

This project asks: can we do better by explicitly selecting positions? We replace GCG's cross-position competition with position pre-commitment — a policy selects one position before candidate evaluation, and all 512 candidates are confined to that position. We test four pre-commitment strategies while keeping every other GCG component identical:

Method Strategy ASR
Vanilla GCG (baseline) All-coordinates competition 78%
AB-GCG Contextual bandit (gradient + attention + time) 46%
Gradient-only bandit Contextual bandit (gradient + time) 40%
Attention-inverse Softmax over inverse attention importance 38%
Attention-only Softmax over attention importance 28%

Key findings:

  • All pre-commitment strategies suffer severe ASR drops (32–50pp), regardless of the selection signal
  • 60–67% of the performance gap is attributable to optimization instability (loss regression), not search quality
  • Within the pre-commitment regime, mixed-signal selection (AB-GCG) outperforms all single-signal variants
  • To our knowledge, this is the first controlled ablation isolating cross-position competition as the critical factor in GCG's effectiveness

Setup

# Install dependencies
pip install nanogcg torch transformers tqdm numpy matplotlib

# Download AdvBench data
mkdir -p data
wget -O data/harmful_behaviors.csv \
  https://raw.githubusercontent.com/llm-attacks/llm-attacks/main/data/advbench/harmful_behaviors.csv

Quick Start

Single-prompt smoke test (~2 min)

python run_single_test.py --steps 50 --search_width 64

Run full experiments

# Vanilla GCG baseline (50 prompts × 500 steps)
python run_baseline.py --num_prompts 50 --num_steps 500

# AB-GCG (contextual bandit, 3D features)
python run_ab_gcg.py --num_prompts 50 --num_steps 500

# Ablation variants
python run_ablations.py --variant attention_only
python run_ablations.py --variant attention_inverse
python run_ablations.py --variant gradient_only

# Analyze results
python analyze_results.py --baseline results/baseline --ab_gcg results/ab_gcg

Multi-GPU parallel execution

# 4-GPU sharding (splits 50 prompts across GPUs)
bash scripts/run_baseline_parallel.sh
bash scripts/run_ab_gcg_parallel.sh
bash scripts/run_attn_only_parallel.sh
bash scripts/run_attn_inverse_parallel.sh
bash scripts/run_gradient_only_parallel.sh

# Merge shard results
python scripts/merge_summaries.py --method baseline

Project Structure

AB-GCG/
├── ab_gcg/                    # Core library
│   ├── attention.py           # Attention importance scores (adapted from Mask-GCG)
│   ├── bandit.py              # PositionBandit MLP + REINFORCE training
│   ├── gcg.py                 # ABGCG class (extends nanogcg's GCG)
│   ├── evaluate.py            # Keyword-based ASR evaluation
│   └── utils.py               # Context vectors, step logging, serialization
├── run_baseline.py            # Vanilla GCG baseline experiment
├── run_ab_gcg.py              # AB-GCG experiment
├── run_ablations.py           # Attention-only, attention-inverse, gradient-only
├── run_single_test.py         # Quick single-prompt smoke test
├── analyze_results.py         # Post-hoc analysis and plots
├── configs.py                 # Shared experiment configuration
├── data_utils.py              # AdvBench data loading
├── scripts/                   # Multi-GPU parallel execution scripts
└── data/
    └── harmful_behaviors.csv  # AdvBench dataset

Design Decisions

Controlled ablation via hook methods. ABGCG inherits from nanogcg's GCG class and overrides only the candidate generation pipeline. Position selection (_select_position) and policy update (_post_step_update) are implemented as hook methods that ablation variants override — the optimization loop itself is shared across all five conditions.

Online learning, no pretraining. The bandit MLP learns from scratch within each GCG run. Cross-prompt transfer is enabled by carrying bandit weights across prompts.

Shared forward pass. Attention importance scores are extracted from the same forward pass used for gradient computation (output_attentions=True), adding negligible overhead.

~120 lines of new code. The entire modification on top of nanogcg is deliberately minimal to keep the ablation clean.

Experimental Configuration

Parameter Value
Target model Qwen-2.5-3B-Instruct
Suffix length 20 tokens
Optimization steps 500 per prompt
Candidate batch size 512
Top-k 256
Evaluation prompts 50 (AdvBench, ≤30 words)
Success criterion Keyword-based (no refusal prefix + target prefix match)
Hardware 4 × NVIDIA A100 80GB
Random seed 42 (shared across all methods)

References

License

MIT

About

An solo research exploration about GCG and RL :) -- Cheney

Resources

Stars

41 stars

Watchers

1 watching

Forks

Releases

Packages

Contributors

Languages