Selecting Better Jailbreak Attacks for Safety

# Illustrative sketch (not from the paper)
import torch
import torch.nn as nn
import itertools

# Black‑box defender utility v_theta(S) – placeholder
def utility_query(attacks_subset):
    # In practice this calls the defended LLM and returns a safety score
    return torch.rand(1).item()

# Surrogate model to approximate utility landscape
class Surrogate(nn.Module):
    def __init__(self, n_attacks):
        super().__init__()
        self.linear = nn.Linear(n_attacks, 1)
    def forward(self, x):
        return self.linear(x)

# Sample a limited set of subset utilities (e.g., ~200)
def sample_subsets(n_attacks, budget):
    samples = []
    for _ in range(200):
        subset = torch.randint(0, 2, (n_attacks,))
        utility = utility_query(subset)
        samples.append((subset.float(), utility))
    return samples

# Train surrogate on sampled data
def train_surrogate(samples, n_attacks):
    model = Surrogate(n_attacks)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    for epoch in range(100):
        for x, y in samples:
            pred = model(x)
            loss = (pred.squeeze() - y) ** 2
            opt.zero_grad()
            loss.backward()
            opt.step()
    return model

# AttackSHAP: marginal contribution via Shapley approximation (Monte‑Carlo)
def attack_shap(model, n_attacks, n_samples=100):
    shapley = torch.zeros(n_attacks)
    for _ in range(n_samples):
        perm = torch.randperm(n_attacks)
        prev_val = 0.0
        for idx in perm:
            vec = torch.zeros(n_attacks)
            vec[perm[:idx+1]] = 1.0
            cur_val = model(vec.unsqueeze(0)).item()
            shapley[idx] += cur_val - prev_val
            prev_val = cur_val
    return shapley / n_samples

# A-MESS‑Greedy selection using true utility queries
def greedy_select(n_attacks, k):
    selected = []
    remaining = set(range(n_attacks))
    while len(selected) < k:
        best_gain = -float('inf')
        best_a = None
        for a in remaining:
            candidate = selected + [a]
            subset_vec = torch.zeros(n_attacks)
            subset_vec[candidate] = 1.0
            gain = utility_query(subset_vec)
            if gain > best_gain:
                best_gain, best_a = gain, a
        selected.append(best_a)
        remaining.remove(best_a)
    return selected

// Illustrative sketch (not from the paper)
const tf = require('@tensorflow/tfjs-node');

// Black‑box defender utility v_theta(S) – placeholder
function utilityQuery(attacksSubset) {
  // In practice this would query the defended LLM and return a safety score
  return Math.random();
}

// Surrogate model (simple linear) to approximate utility landscape
function createSurrogate(nAttacks) {
  const model = tf.sequential();
  model.add(tf.layers.dense({inputShape: [nAttacks], units: 1}));
  model.compile({optimizer: tf.train.adam(1e-3), loss: 'meanSquaredError'});
  return model;
}

// Sample a limited set of subset utilities (e.g., ~200)
function sampleSubsets(nAttacks) {
  const samples = [];
  for (let i = 0; i < 200; i++) {
    const subset = tf.randomUniform([nAttacks], 0, 2, 'int32');
    const utility = utilityQuery(subset);
    samples.push({x: subset.toFloat(), y: utility});
  }
  return samples;
}

// Train surrogate on sampled data
async function trainSurrogate(samples, nAttacks) {
  const model = createSurrogate(nAttacks);
  const xs = tf.stack(samples.map(s => s.x));
  const ys = tf.tensor1d(samples.map(s => s.y));
  await model.fit(xs, ys, {epochs: 100, verbose: 0});
  return model;
}

// AttackSHAP: Monte‑Carlo Shapley approximation using surrogate predictions
async function attackShap(model, nAttacks, nSamples = 100) {
  const shapley = tf.zeros([nAttacks]);
  for (let s = 0; s < nSamples; s++) {
    const perm = tf.util.createShuffledIndices(nAttacks);
    let prevVal = 0;
    for (let i = 0; i < perm.length; i++) {
      const idx = perm[i];
      const vec = tf.buffer([nAttacks]);
      for (let j = 0; j <= i; j++) vec.set(1, perm[j]);
      const curVal = (await model.predict(vec.toTensor().expandDims(0)).data())[0];
      shapley.buffer().set(shapley.buffer().get(idx) + (curVal - prevVal), idx);
      prevVal = curVal;
    }
  }
  return shapley.div(tf.scalar(nSamples));
}

// A-MESS‑Greedy selection using true utility queries
function greedySelect(nAttacks, k) {
  const selected = [];
  const remaining = new Set([...Array(nAttacks).keys()]);
  while (selected.length < k) {
    let bestGain = -Infinity;
    let bestA = null;
    for (const a of remaining) {
      const candidate = [...selected, a];
      const vec = tf.zeros([nAttacks]);
      candidate.forEach(i => vec.buffer().set(1, i));
      const gain = utilityQuery(vec);
      if (gain > bestGain) { bestGain = gain; bestA = a; }
    }
    selected.push(bestA);
    remaining.delete(bestA);
  }
  return selected;
}