from __future__ import annotations
import logging
import math
import random
from typing import Callable, List, Literal, Optional, Tuple, Union
import torch
from jaxtyping import Float, Int
from torch import Tensor
from tropt.common import (
DEFAULT_INIT_TRIGGER,
Targets,
TextTemplates,
)
from tropt.loss import BaseLoss
from tropt.model import (
BaseModel,
GradientTokenAccessMixin,
LogitsTokenAccessMixin,
LossTextAccessMixin,
LossTokenAccessMixin,
TokenAccessMixin,
)
from tropt.optimizer import BaseOptimizer, OptimizerResult
from tropt.optimizer.utils.buffer import TriggerBuffer
from tropt.optimizer.utils.retokenization import retokenize_filtering
from tropt.optimizer.utils.running_best import RunningBest
from tropt.optimizer.utils.token_constraints import TokenConstraints
from tropt.optimizer.utils.token_initializers import get_printable_random_trigger
from tropt.tracker import BaseTracker
logger = logging.getLogger(__name__)
[docs]
class GCGPlusOptimizer(BaseOptimizer):
"""Flexible GCG implementation, supporting tricks from GCG, QCG, GASLITE, and UAT.
Two-stage design:
1. Candidate selection (on proxy model) — gradient-based, random, or focused.
2. Candidate evaluation (on target model) — via text or token access.
References:
- GCG: https://arxiv.org/abs/2307.15043
- QCG: https://arxiv.org/abs/2402.12329
- PAL: https://arxiv.org/abs/2402.09674
- GASLITE: https://arxiv.org/abs/2412.20953
- UAT: https://arxiv.org/abs/1908.07125
"""
model_requirements = (LossTextAccessMixin,)
def __init__(
self,
model: BaseModel,
loss: BaseLoss,
tracker: Optional[BaseTracker] = None,
seed: Optional[int] = None,
# Proxy model for candidate selection (defaults to model):
proxy_model: Optional[BaseModel] = None,
# Differentiable loss for proxy gradient computation (defaults to loss):
proxy_loss: Optional[BaseLoss] = None,
# Candidate selection strategy:
candidate_selection: Literal["gradient", "random", "focused"] = "gradient",
# Core GCG params:
num_steps: int = 500,
n_candidates: int = 512,
sample_topk: int = 256,
token_constraints: TokenConstraints = TokenConstraints(),
use_retokenize: bool = True,
# Later tricks:
sample_n_replace: Union[int, Tuple[int, int]] = (1, 1),
candidate_oversample_factor: float = 1.1,
momentum: float = 0.0,
skip_visited: bool = False,
# Trigger buffer size:
buffer_size: Optional[int] = None,
n_grad_avg: int = 1,
# Per-step batch sampling:
template_batch_size: Optional[int] = None,
template_batch_sampler: Optional[Callable[[TextTemplates, Optional[Targets]], List[int]]] = None,
# Token prior (readability):
token_prior_weight: float = 0.0,
sample_temperature: Optional[float] = None,
):
"""
Args:
proxy_model: Model used for candidate selection (gradients/tokenizer).
If None, defaults to `model` (self-proxy / white-box).
proxy_loss: Loss function for proxy gradient computation. Required
when the main loss is non-differentiable and
candidate_selection="gradient". Defaults to `loss`.
candidate_selection: "gradient" uses gradient-ranked top-k sampling;
"random" uses uniform random token sampling; "focused" probes
all positions with target model loss then focuses on the best one (from QCG paper).
sample_n_replace: (start, end) number of token positions to replace per
candidate. Linearly interpolated over optimization steps; from PAL paper.
Defaults to (1, 1) for single-token flips only.
candidate_oversample_factor: Craft n_candidates * factor candidates (>1.0), then
truncate after retokenization filtering, to fill the required number of candidates.
From PAL paper.
Defaults to 1.1 (10% oversampling).
momentum: Gradient momentum coefficient. When > 0, enables momentum:
m = mu*m + (1-mu)*grad for candidate ranking instead of raw gradient.
Defaults to 0.0 (no momentum).
Reference: https://arxiv.org/abs/2405.01229 .
skip_visited: If True, never re-select a previously accepted trigger
(PAL-style); avoids cycling/stalling. Defaults to False.
buffer_size: If set, maintain a buffer of the best triggers seen (from QCG paper). Each step
starts from the best buffer entry and updates it with improved candidates.
Defaults to None (no buffer).
n_grad_avg: Number of trigger perturbations to average gradients
over. When > 1, flips a random position per copy (from ARCA/GASLITE papers).
Defaults to 1 (no averaging, like GCG).
template_batch_size: If set, sample this many templates (and their
targets) per optimization step instead of using all templates
simultaneously.
Useful for large template sets.
Reference: https://arxiv.org/abs/1908.07125 . Defaults to None (use all templates).
template_batch_sampler: If set, called as `template_batch_sampler(templates, targets)` each step to pick
the template indices for that step (e.g., stratified batches); defaults to uniform sampling of `template_batch_size` templates.
token_prior_weight: weight of the readability prior (under gradient selection only). When non zero, in each step, the proxy's log p(v | prefix) at every trigger position (one forward on the current trigger) is added to the loss. Inspired by AutoDAN (https://arxiv.org/abs/2310.15140). Requires a proxy with logits access and a trigger slice that starts after position 0 (i.e., `use_prefix_cache=False`). Defaults to 0.0 (off).
sample_temperature: If set, sample tokens from the top-k by softmax(score / T) instead of uniformly
(under gradient selection only). Defaults to None (uniform).
"""
assert template_batch_sampler is None or template_batch_size is None, (
"Pass either `template_batch_sampler` or `template_batch_size`, not both."
)
super().__init__(model, loss=loss, tracker=tracker, seed=seed)
# Proxy model setup
self.proxy_model = model if proxy_model is None else proxy_model
assert isinstance(self.proxy_model, TokenAccessMixin), (
"proxy_model must support TokenAccessMixin (tokenizer access)"
)
self.proxy_loss = proxy_loss if proxy_loss is not None else self.loss_func
if candidate_selection == "gradient":
assert isinstance(self.proxy_model, GradientTokenAccessMixin), (
"candidate_selection='gradient' requires proxy_model with GradientTokenAccessMixin"
)
if not self.proxy_loss.is_differentiable:
raise ValueError(
f"candidate_selection='gradient' requires a differentiable proxy_loss, "
f"but {type(self.proxy_loss).__name__}.is_differentiable=False. "
f"Pass a differentiable proxy_loss (e.g. PrefillCELoss())."
)
assert candidate_oversample_factor >= 1.0, "candidate_oversample_factor must be >= 1.0"
if token_prior_weight > 0 or sample_temperature is not None:
assert candidate_selection == "gradient", (
"`token_prior_weight` / `sample_temperature` require candidate_selection='gradient'."
)
if token_prior_weight > 0:
assert isinstance(self.proxy_model, LogitsTokenAccessMixin), (
"`token_prior_weight` requires a proxy_model with LogitsTokenAccessMixin."
)
# Prefer token-level target evaluation when proxy and target share the same tokenizer
# (otherwise, if tokenizer are not shared, or we don't have access to the target model's tokenzier, we simply use text-level loss compuation)
use_token_input_for_loss = (model.tokenizer == self.proxy_model.tokenizer) and isinstance(model, LossTokenAccessMixin)
# Normalize sample_n_replace to tuple
if isinstance(sample_n_replace, int):
sample_n_replace = (sample_n_replace, sample_n_replace)
# Save params
self.candidate_selection = candidate_selection
self.num_steps = num_steps
self.n_candidates = n_candidates
self.sample_topk = sample_topk
self.sample_n_replace = sample_n_replace
self.token_constraints = token_constraints
self.use_retokenize = use_retokenize
self.candidate_oversample_factor = candidate_oversample_factor
self.momentum = momentum
self.skip_visited = skip_visited
self.use_token_input_for_loss = use_token_input_for_loss
self.buffer_size = buffer_size
self.n_grad_avg = n_grad_avg
self.template_batch_size = template_batch_size
if template_batch_sampler is None and template_batch_size is not None: # default: uniform random batch
template_batch_sampler = lambda templates, targets: random.sample(range(len(templates)), template_batch_size) # noqa: E731
self.template_batch_sampler = template_batch_sampler
self.token_prior_weight = token_prior_weight
self.sample_temperature = sample_temperature
[docs]
def optimize_trigger(
self,
templates: TextTemplates,
initial_trigger: Optional[str] = DEFAULT_INIT_TRIGGER,
targets: Optional[Targets] = None,
) -> OptimizerResult:
# --- Initialization ---
proxy_model = self.proxy_model
proxy_tokenizer = proxy_model.tokenizer
target_model = self.model
# Batch sampling setup
use_batch_sampling = self.template_batch_sampler is not None and (
self.template_batch_size is None or self.template_batch_size < len(templates)
)
proxy_model.set_inputs_from_tokens(templates=templates, targets=targets)
if not self.use_token_input_for_loss:
target_model.set_inputs_from_texts(templates=templates, targets=targets)
else:
target_model.set_inputs_from_tokens(templates=templates, targets=targets)
trigger_ids: Int[Tensor, "trigger_seq_len"] = proxy_tokenizer.encode_trigger(initial_trigger).to(proxy_model.device)
vocab_size = proxy_model.vocab_size
blacklist_ids = self.token_constraints.get_blacklist_ids(proxy_tokenizer, vocab_size)
valid_token_ids = self.token_constraints.get_whitelist_ids(proxy_tokenizer, vocab_size, proxy_model.device, return_tensor=True)
best = RunningBest()
momentum_buffer: Optional[Tensor] = None
visited: set[str] = set()
# Buffer initialization
buffer: Optional[TriggerBuffer] = None
if self.buffer_size is not None:
buffer = self._init_buffer(trigger_ids)
# Number of candidates to generate (oversample if retokenize is on)
n_candidates_oversampled = self.n_candidates
if self.use_retokenize and self.candidate_oversample_factor > 1.0:
n_candidates_oversampled = math.ceil(self.n_candidates * self.candidate_oversample_factor)
# Initial loss (on full set or first batch)
current_loss = self._evaluate_candidates(
trigger_ids.unsqueeze(0)
).item()
trigger_str = proxy_tokenizer.decode_trigger(trigger_ids)
self.log(loss=current_loss, trigger_str=trigger_str)
if self.skip_visited:
visited.add(trigger_str)
n_replace_start, n_replace_end = self.sample_n_replace
for step_i in self.track_steps(range(self.num_steps)):
# --- batch sampling: re-set model inputs with a random subset ---
if use_batch_sampling:
batch_indices = self.template_batch_sampler(templates, targets)
batch_templates = [templates[i] for i in batch_indices]
batch_targets = (
targets.select_indices(batch_indices)
if targets is not None else None
)
proxy_model.set_inputs_from_tokens(
templates=batch_templates, targets=batch_targets,
)
if not self.use_token_input_for_loss:
target_model.set_inputs_from_texts(
templates=batch_templates, targets=batch_targets,
)
else:
target_model.set_inputs_from_tokens(
templates=batch_templates, targets=batch_targets,
)
# Linearly interpolate sample_n_replace over steps
cur_n_replace = round(
n_replace_start + (n_replace_end - n_replace_start) * step_i / max(self.num_steps - 1, 1)
)
# Buffer mode: start each step from the best buffer entry
if buffer is not None:
trigger_ids = buffer.get_best_trigger()
# === Stage 1: Candidate Selection (on proxy) ===
token_log_probs: Optional[Tensor] = None
if self.candidate_selection == "gradient":
grad_triggers = self._get_grad_trigger_variations(
trigger_ids, valid_token_ids
)
trigger_grad = proxy_model.compute_grad_from_tokens(
candidate_trigger_ids=grad_triggers,
loss_func=self.proxy_loss,
normalize_grads=True,
).mean(dim=0) # average over n_grad_avg variations
# Apply momentum
if self.momentum > 0:
if momentum_buffer is None:
momentum_buffer = trigger_grad.clone()
else:
momentum_buffer = (
self.momentum * momentum_buffer
+ (1 - self.momentum) * trigger_grad
)
trigger_grad = momentum_buffer
if self.token_prior_weight > 0:
# log p(v | prefix) at every trigger position, from one forward on the current trigger
token_log_probs = proxy_model.compute_logits_from_tokens(
trigger_ids.unsqueeze(0), return_trigger_logits_only=True,
)[0].float().log_softmax(dim=-1) # (trigger_seq_len, vocab_size)
candidate_trigger_ids = self._sample_ids_from_grad(
trigger_ids=trigger_ids,
trigger_grad=trigger_grad,
token_log_probs=token_log_probs,
blacklist_ids=blacklist_ids,
n_candidates=n_candidates_oversampled,
n_replace=cur_n_replace,
)
elif self.candidate_selection == "focused":
candidate_trigger_ids = self._sample_focused_candidates(
trigger_ids=trigger_ids,
valid_token_ids=valid_token_ids,
n_candidates=n_candidates_oversampled,
)
else: # "random"
candidate_trigger_ids = self._sample_random_candidates(
trigger_ids=trigger_ids,
valid_token_ids=valid_token_ids,
n_candidates=n_candidates_oversampled,
n_replace=cur_n_replace,
)
# === Retokenization filtering ===
if self.use_retokenize:
candidate_trigger_ids = retokenize_filtering(
candidate_trigger_ids, proxy_tokenizer
)
# Drop already-accepted triggers (PAL-style) before truncating
if self.skip_visited and len(candidate_trigger_ids) > 0:
cand_strs = proxy_tokenizer.decode_triggers(candidate_trigger_ids)
keep = torch.tensor(
[s not in visited for s in cand_strs],
device=candidate_trigger_ids.device,
)
candidate_trigger_ids = candidate_trigger_ids[keep]
# Truncate to n_candidates (after oversample + filter)
candidate_trigger_ids = candidate_trigger_ids[: self.n_candidates]
if len(candidate_trigger_ids) == 0:
logger.warning("All candidates filtered out, skipping step.")
continue
# === Stage 2: Candidate Evaluation (on target) ===
losses = self._evaluate_candidates(candidate_trigger_ids)
if token_log_probs is not None:
positions = torch.arange(candidate_trigger_ids.shape[1], device=candidate_trigger_ids.device)
cand_log_probs = token_log_probs[positions, candidate_trigger_ids].sum(dim=-1) # (n_candidates,)
losses = losses - self.token_prior_weight * cand_log_probs.to(losses)
# === Update best / buffer ===
if buffer is not None:
# Buffer mode: update buffer with all evaluated candidates
for idx in range(len(candidate_trigger_ids)):
buffer.add_if_better(
candidate_trigger_ids[idx], losses[idx].item()
)
# Track overall best from buffer for logging
current_loss = buffer.get_lowest_loss()
trigger_ids = buffer.get_best_trigger()
trigger_str = proxy_tokenizer.decode_trigger(trigger_ids)
else:
current_loss = losses.min().item()
trigger_ids = candidate_trigger_ids[losses.argmin()]
trigger_str = proxy_tokenizer.decode_trigger(trigger_ids)
best.update(loss=current_loss, trigger_ids=trigger_ids, trigger_str=trigger_str)
self.log(loss=current_loss, trigger_str=trigger_str)
if self.skip_visited:
visited.add(trigger_str)
# --- Finalize ---
result = best.to_result()
return result
# ------------------------------------------------------------------
# Initialization
# ------------------------------------------------------------------
def _init_buffer(
self,
initial_trigger_ids: Int[Tensor, "trigger_seq_len"],
) -> TriggerBuffer:
"""Initialize the buffer with the initial trigger and printable-random variants.
"""
assert self.buffer_size is not None
trigger_seq_len = initial_trigger_ids.shape[0]
device = initial_trigger_ids.device
tokenizer = self.proxy_model.tokenizer
blacklist_ids = self.token_constraints.get_blacklist_ids(tokenizer, self.proxy_model.vocab_size)
triggers_list: list[Tensor] = [initial_trigger_ids]
for _ in range(self.buffer_size - 1):
# Oversample length (reencode can drift shorter), then trim to trigger_seq_len.
rand_str = get_printable_random_trigger(
trigger_len=2 * trigger_seq_len,
tokenizer=tokenizer,
blacklist_ids=blacklist_ids,
)
ids = tokenizer.encode_trigger(rand_str).to(device)[:trigger_seq_len]
triggers_list.append(ids)
all_triggers = torch.stack(triggers_list, dim=0)
all_losses = self._evaluate_candidates(all_triggers)
return TriggerBuffer(
triggers=list(all_triggers),
losses=[all_losses[i].item() for i in range(self.buffer_size)],
)
# ------------------------------------------------------------------
# Stage 1: Candidate Selection
# ------------------------------------------------------------------
def _sample_ids_from_grad(
self,
trigger_ids: Int[Tensor, "trigger_seq_len"],
trigger_grad: Float[Tensor, "trigger_seq_len vocab_size"],
blacklist_ids: list,
n_candidates: int,
n_replace: int = 1,
token_log_probs: Optional[Float[Tensor, "trigger_seq_len vocab_size"]] = None,
) -> Int[Tensor, "n_candidates trigger_seq_len"]:
"""Sample candidate token sequences via GCG-style random multi-position replacement."""
trigger_seq_len = trigger_grad.shape[0]
device = trigger_grad.device
n_replace = min(n_replace, trigger_seq_len)
# Top-k candidates per position
scores = -trigger_grad # (new tensor; avoids mutating the momentum buffer)
if token_log_probs is not None:
assert token_log_probs.shape == scores.shape, (
f"Token log-probs {tuple(token_log_probs.shape)} do not match the gradient {tuple(scores.shape)}."
)
scores = scores + self.token_prior_weight * token_log_probs.to(scores)
scores[:, blacklist_ids] = float("-inf")
topk = scores.topk(self.sample_topk, dim=-1)
topk_ids = topk.indices
candidate_trigger_ids = trigger_ids.repeat(n_candidates, 1).clone()
# Random positions to flip
sampled_ids_pos = torch.rand(
n_candidates, trigger_seq_len, device=device
).argsort(dim=-1)[..., :n_replace]
# Select relevant top-k lists and sample one token from each
relevant_topk_lists = topk_ids[sampled_ids_pos]
if self.sample_temperature is None:
rand_k_indices = torch.randint(
0,
self.sample_topk,
(n_candidates, n_replace, 1),
device=device,
)
else: # softmax(score / T) over each position's top-k
probs = (topk.values[sampled_ids_pos] / self.sample_temperature).softmax(dim=-1) # (n_cand, n_replace, k)
rand_k_indices = torch.multinomial(probs.view(-1, self.sample_topk), 1).view(n_candidates, n_replace, 1)
sampled_ids_val = torch.gather(
input=relevant_topk_lists,
dim=-1,
index=rand_k_indices,
).squeeze(-1)
candidate_trigger_ids = candidate_trigger_ids.scatter_(
dim=-1,
index=sampled_ids_pos,
src=sampled_ids_val,
)
return candidate_trigger_ids
def _get_grad_trigger_variations(
self,
trigger_ids: Int[Tensor, "trigger_seq_len"],
valid_token_ids: Int[Tensor, "n_valid"],
) -> Int[Tensor, "n_grad_avg trigger_seq_len"]:
"""Create trigger variations for gradient averaging (GASLITE-style).
With n_grad_avg == 1, returns the trigger as-is.
With n_grad_avg > 1, flips a random position with a random valid token
per copy (keeping the first copy intact).
"""
if self.n_grad_avg <= 1:
return trigger_ids.unsqueeze(0)
device = trigger_ids.device
trigger_seq_len = trigger_ids.shape[0]
grad_triggers = trigger_ids.repeat(self.n_grad_avg, 1).clone()
# Keep first copy intact, perturb the rest
for idx in range(1, self.n_grad_avg):
pos = torch.randint(0, trigger_seq_len, (1,), device=device).item()
tok = valid_token_ids[
torch.randint(0, len(valid_token_ids), (1,), device=device)
].item()
assert isinstance(pos, int) and isinstance(tok, int)
grad_triggers[idx, pos] = tok
return grad_triggers
def _sample_random_candidates(
self,
trigger_ids: Int[Tensor, "trigger_seq_len"],
valid_token_ids: Int[Tensor, "n_valid"],
n_candidates: int,
n_replace: int = 1,
) -> Int[Tensor, "n_candidates trigger_seq_len"]:
"""Sample candidates by randomly replacing tokens (RAL-style)."""
trigger_seq_len = trigger_ids.shape[0]
device = trigger_ids.device
n_replace = min(n_replace, trigger_seq_len)
candidate_trigger_ids = trigger_ids.repeat(n_candidates, 1).clone()
# Random positions to flip
sampled_ids_pos = torch.rand(
n_candidates, trigger_seq_len, device=device
).argsort(dim=-1)[..., :n_replace]
# Sample uniformly from valid tokens
rand_indices = torch.randint(
0, len(valid_token_ids), (n_candidates, n_replace), device=device
)
random_tokens = valid_token_ids[rand_indices]
candidate_trigger_ids = candidate_trigger_ids.scatter_(
dim=-1,
index=sampled_ids_pos,
src=random_tokens,
)
return candidate_trigger_ids
def _sample_focused_candidates(
self,
trigger_ids: Int[Tensor, "trigger_seq_len"],
valid_token_ids: Int[Tensor, "n_valid"],
n_candidates: int,
) -> Int[Tensor, "n_candidates trigger_seq_len"]:
"""Focused position sampling (from QCG paper).
Phase 1: Probe each position with one random token replacement,
evaluate all probes on the target (via _evaluate_candidates) to find
the most promising position.
Phase 2: Generate n_candidates candidates at the best position.
"""
trigger_seq_len = trigger_ids.shape[0]
device = trigger_ids.device
# Phase 1: Probe each position with one random token
probe_candidates = trigger_ids.repeat(trigger_seq_len, 1).clone()
probe_tokens = valid_token_ids[
torch.randint(0, len(valid_token_ids), (trigger_seq_len,), device=device)
]
# Replace position j in candidate j
probe_candidates[
torch.arange(trigger_seq_len, device=device),
torch.arange(trigger_seq_len, device=device),
] = probe_tokens
probe_losses = self._evaluate_candidates(probe_candidates)
best_pos = probe_losses.argmin().item()
assert isinstance(best_pos, int)
# Phase 2: Generate candidates at best position only
candidates = trigger_ids.repeat(n_candidates, 1).clone()
random_tokens = valid_token_ids[
torch.randint(0, len(valid_token_ids), (n_candidates,), device=device)
]
candidates[:, best_pos] = random_tokens
return candidates
# ------------------------------------------------------------------
# Stage 2: Candidate Evaluation
# ------------------------------------------------------------------
def _evaluate_candidates(
self,
candidate_trigger_ids: Int[Tensor, "n_candidates trigger_seq_len"],
) -> Float[Tensor, "n_candidates"]:
"""Evaluate candidates on the target model. Returns per-candidate loss.
When use_token_input_for_loss is False, decodes candidates via proxy_model.tokenizer.
"""
if self.use_token_input_for_loss:
losses = self.model.compute_loss_from_tokens(
candidate_trigger_ids, loss_func=self.loss_func
)
else:
candidate_strs = self.proxy_model.tokenizer.decode_triggers(candidate_trigger_ids)
losses = self.model.compute_loss_from_texts(
candidate_strs, loss_func=self.loss_func
)
# Reduce to (n_candidates,) if needed
if losses.dim() > 1:
losses = losses.mean(dim=0)
return losses
