Accelerating Greedy Coordinate Gradient and General Prompt Optimization via Probe Sampling

TL;DR

This work tackles the computational bottleneck of Greedy Coordinate Gradient (GCG) for adversarial prompt optimization by introducing Probe sampling, which adaptively filters candidate prompts using a smaller draft model based on a probe agreement with the target model. The method significantly reduces forward computations, achieving up to 5.6× speedups (with simulated annealing) and improving attack success rates on AdvBench, while also accelerating other discrete prompt optimization methods like AutoPrompt, APE, and AutoDAN. Through extensive experiments across multiple models and tasks, the authors show robust transferability and practical gains, along with detailed analyses of memory and time allocation, parameter sensitivity (e.g., filtered-set size and probe-set size), and robustness of the agreement metric. The approach offers a general mechanism for conditional computation in prompt optimization and suggests avenues for extending adaptive filtering to multi-modality and larger-scale draft models, albeit with limitations on large test sets and proprietary models.

Abstract

Safety of Large Language Models (LLMs) has become a critical issue given their rapid progresses. Greedy Coordinate Gradient (GCG) is shown to be effective in constructing adversarial prompts to break the aligned LLMs, but optimization of GCG is time-consuming. To reduce the time cost of GCG and enable more comprehensive studies of LLM safety, in this work, we study a new algorithm called $\texttt{Probe sampling}$. At the core of the algorithm is a mechanism that dynamically determines how similar a smaller draft model's predictions are to the target model's predictions for prompt candidates. When the target model is similar to the draft model, we rely heavily on the draft model to filter out a large number of potential prompt candidates. Probe sampling achieves up to $5.6$ times speedup using Llama2-7b-chat and leads to equal or improved attack success rate (ASR) on the AdvBench. Furthermore, probe sampling is also able to accelerate other prompt optimization techniques and adversarial methods, leading to acceleration of $1.8\times$ for AutoPrompt, $2.4\times$ for APE and $2.4\times$ for AutoDAN.

Figures (5)

• Figure 1: A brief illustration of the Greedy Coordinate Gradient (GCG) algorithm zou2023universal.
• Figure 2: Probe sampling mainly consists of three steps. (i) A batch of candidates ($\{a, b,\cdots, h\}$) is sampled. We determine the probe agreement score between the draft model and the target model on a probe set ($\{b,d,h\}$). The probe agreement score is used to compute the filtered set size. (ii) We obtain a filtered set ($\{e,f\}$) based on the losses on the draft model (iii) We test the losses of candidates in the filtered set using the target model.
• Figure 3: Memory usage on a single A100 with 80GB memory with (a) Llama2-7b-chat and (b) Vicuna-7b-v1.3 on $1$ instance. The memory consumption of probe sampling with or without simulated annealing is similar to that of the original setting. The computation with the target model still takes most of the memory.
• Figure 4: Wall time of GCG, probe sampling with and without simulated annealing. For the target model computation, the first part is done on the probe set and the second part is done on the filtered set. Draft model computation and computation of the target model on the probe set are suited to be done in parallel as they take similar time.
• Figure 5: Converge progress with different sizes of filtered set.