Jailbreaking › Code: Simulating the GCG Attack

import numpy as np

# Simulate a simple token-level gradient attack
# We model the attack on a small vocabulary with a toy "safety score"

np.random.seed(42)

# Vocabulary and "safety" landscape
vocab_size = 50
seq_length = 8
n_steps = 100
top_k = 5

# Simulate model: a random quadratic "safety loss" surface
# In reality, this would be cross-entropy on target affirmative tokens
W_safety = np.random.randn(vocab_size, seq_length)


def compute_loss(token_sequence):
    """Toy safety loss: sum of W_safety values at each position."""
    return sum(W_safety[tok, pos] for pos, tok in enumerate(token_sequence))


def compute_gradient(token_sequence):
    """Token-level gradient: loss change if we swap each position's token."""
    gradients = []
    current_loss = compute_loss(token_sequence)
    for pos in range(seq_length):
        grad_at_pos = W_safety[:, pos] - W_safety[token_sequence[pos], pos]
        gradients.append(grad_at_pos)
    return gradients


def gcg_step(token_sequence, top_k=5):
    """One GCG step: find candidates, try random swaps."""
    gradients = compute_gradient(token_sequence)
    candidates = []
    for pos, grad in enumerate(gradients):
        # top-k tokens with most negative gradient (biggest loss decrease)
        best_tokens = np.argsort(grad)[:top_k]
        for tok in best_tokens:
            candidates.append((pos, tok))

    # Try random candidates and accept improvements
    np.random.shuffle(candidates)
    current_loss = compute_loss(token_sequence)
    new_sequence = list(token_sequence)

    for pos, tok in candidates[:top_k]:
        test_sequence = list(token_sequence)
        test_sequence[pos] = tok
        new_loss = compute_loss(test_sequence)
        if new_loss < current_loss:
            new_sequence = test_sequence
            current_loss = new_loss
            break

    return new_sequence, current_loss


# Run GCG optimization
initial_tokens = np.random.randint(0, vocab_size, seq_length).tolist()
current_tokens = list(initial_tokens)
loss_history = [compute_loss(current_tokens)]

for step in range(n_steps):
    current_tokens, loss = gcg_step(current_tokens, top_k=top_k)
    loss_history.append(loss)