AB-GCG

KO | EN

gitlite — search

#javascript #hacktoberfest #react #typescript #llm #go #golang #rust #deep-learning #nodejs #windows #kubernetes #awesome-list #pytorch

(★ 41)

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

AB-GCG Latest Version Download

Download Latest Version (.zip)

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# 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](https://medium.com/@cheneyshyu/i-tried-four-smarter-ways-to-select-positions-in-gcg-f0ed2fb64023)

---

## What This Project Does

[GCG](https://arxiv.org/abs/2307.15043) (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

```bash
# 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)
```bash
python run_single_test.py --steps 50 --search_width 64
```

### Run full experiments
```bash
# 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
```bash
# 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

- Zou et al., 2023. [Universal and Transferable Adversarial Attacks on Aligned Language Models](https://arxiv.org/abs/2307.15043)
- Mu et al., 2025. [Mask-GCG: Improving GCG Attacks through Mask Optimization](https://arxiv.org/abs/2509.06350)
- Li et al., 2025. [MAGIC: Exploiting the Index Gradients for Optimization-Based Jailbreaking on LLMs](https://arxiv.org/abs/2412.08615)
- Jia et al., 2024. [Improved Techniques for Optimization-Based Jailbreaking on LLMs](https://arxiv.org/abs/2405.21018)
- Williams, 1992. Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning
- [nanogcg](https://github.com/GraySwanAI/nanoGCG) by Gray Swan AI

## License

MIT

AB-GCG

AB-GCG Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

AB-GCG

AB-GCG Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

Download README