Adaptive Instruction Composition for Automated LLM Red-Teaming

1. Problem Statement

Large Language Models (LLMs) are vulnerable to "jailbreak" attacks, where malicious prompts bypass safety protocols to generate harmful content. While automated red-teaming is essential for identifying these vulnerabilities, existing approaches suffer from two main limitations:

• Trial-and-Error Limitations: Methods that rely on an attacker LLM to iteratively discover strategies (e.g., PAIR, TAP) often converge on a semantically limited range of successful attacks, failing to explore the vast space of potential vulnerabilities efficiently.
• Random Combination Limitations: The recently proposed WildTeaming framework combines crowdsourced harmful queries and jailbreak tactics randomly. While this generates diversity, it lacks optimization. It does not leverage previous successes to guide future attacks, making it inefficient for discovering specific vulnerabilities in a reasonable timeframe, especially against well-aligned models.

The core challenge is to navigate a massive combinatorial space (e.g., 50,000 queries $\times$ 13,000 tactics = trillions of possible instructions) to jointly optimize attack effectiveness (success rate) and diversity, given the high cost of LLM inference which limits the number of trials.

2. Methodology: Adaptive Instruction Composition (AIC)

The authors propose Adaptive Instruction Composition (AIC), a framework that uses Reinforcement Learning (RL) to intelligently select and combine text samples for attack generation.

Core Architecture

• Contextual Bandit Formulation: The problem is modeled as a Neural Thompson Sampling task, a variant of the contextual multi-armed bandit. The "arms" are potential instruction compositions (combinations of a query and jailbreak tactics).
• The Agent: A lightweight, two-layer feed-forward neural network acts as the bandit ($f(\cdot|\theta)$).
• Input Representation:
  • The system uses WildJailbreak data (50k queries, 13k tactics).
  • Text samples are encoded using SBERT (Sentence-BERT), a contrastively pre-trained model, to generate 768-dimensional embeddings.
  • These embeddings are reduced to 10 dimensions using UMAP to minimize the number of trainable parameters and allow the network to generalize across semantically related texts.
  • The input to the bandit is the concatenation of the query embedding and tactic embeddings.
• The Loop (Algorithm 1):
  1. Sampling: In each trial, $K=500$ candidate instruction compositions are randomly sampled from the massive pool.
  2. Scoring: The bandit network scores these candidates based on learned parameters $\theta$, approximating the probability of a successful attack.
  3. Selection: The candidate with the highest sampled reward (based on a Gaussian posterior distribution derived from the Neural Tangent Kernel) is selected.
  4. Generation & Evaluation: The selected instruction is fed to an Attacker LLM (Mixtral-8x7B), which generates an attack against a Target LLM. An Evaluator (Llama-Guard-2 or Harmbench classifier) assesses if the attack succeeded.
  5. Update: The embedding of the selected instruction and the binary reward (success/failure) are used to update the bandit's weights via Stochastic Gradient Descent (SGD).
  6. Deduplication: Successful compositions are blocked from future selection to force the bandit to generalize across semantically similar inputs rather than memorizing specific successful prompts.

Exploration vs. Exploitation

The framework balances exploration (trying new, diverse combinations) and exploitation (focusing on known successful patterns) via a hyperparameter $\lambda$:

• Subtle Bandit ($\lambda=1$): High regularization, promoting broad exploration and diversity.
• Aggressive Bandit ($\lambda=0.01$ or $0.1$): Low regularization, promoting rapid exploitation of discovered vulnerabilities.

3. Key Contributions

1. Novel Framework: Introduction of AIC, which replaces random combination with an adaptive, RL-driven mechanism to compose jailbreak instructions.
2. Scalable Generalization: Demonstration that using contrastive embeddings (SBERT) allows a lightweight neural network to generalize effectively across a combinatorial space of trillions of possibilities, whereas non-contrastive embeddings (standard BERT) fail to scale efficiently.
3. Optimization of Diversity and Effectiveness: The system successfully balances the trade-off between finding high-success attacks and maintaining semantic diversity, a critical requirement for comprehensive safety training.
4. Transferability: The learned attack policies generalize across different target models (e.g., training on Mistral-7B and attacking Llama-3-70B), suggesting the discovery of general vulnerability patterns rather than model-specific artifacts.

4. Experimental Results

The authors evaluated AIC against WildTeaming (random combination) and other state-of-the-art red-teaming methods (GCG, PAIR, AutoDAN-Turbo) using the Harmbench benchmark.

• Effectiveness (Attack Success Rate - ASR):
  • AIC significantly outperformed WildTeaming across all tested targets (Mistral-7B, Llama-3-70B, Llama-3.3-70B).
  • Example: Against Llama-3-70B, WildTeaming achieved an ASR of 0.088, while the Aggressive AIC achieved 0.450 (a ~5x improvement).
  • On the Harmbench benchmark, AIC achieved an ASR of 1.000 on Mistral-7B and 0.934 on Llama-3-70B, surpassing AutoDAN-Turbo (0.976 and 0.672 respectively).
• Diversity:
  • The "Subtle" bandit maintained high diversity (low similarity between attacks) while still outperforming WildTeaming in success rates.
  • The "Aggressive" bandit narrowed the search space to specific vulnerability themes (e.g., medical misinformation, fraud) but still maintained substantial coverage of unique tactics.
• Transferability:
  • Policies trained on one model transferred effectively to others. For instance, a policy trained on Llama-3 achieved a 2.89x improvement over WildTeaming when applied to Llama-3, and a 2.20x improvement when transferred to Mistral.
• Ablation Studies:
  • Embedding Type: SBERT (contrastive) embeddings yielded significantly faster learning and higher ASR than BERT embeddings.
  • Input Dimensions: Reducing embeddings to 10 dimensions via UMAP was optimal for learning efficiency.
  • Tactics Count: Using multiple tactics (e.g., 3) improved diversity metrics compared to single-tactic setups.

5. Significance and Future Directions

• Safety Training: AIC provides a mechanism to automatically generate comprehensive, diverse, and effective safety training datasets tailored to specific model vulnerabilities, addressing the bottleneck of manual red-teaming.
• Efficiency: By using a lightweight bandit and contrastive embeddings, the system scales to massive search spaces without requiring the computational cost of training a full LLM for the attacker.
• Ethical Considerations: The authors emphasize that this tool is designed for controlled red-teaming within organizational settings to improve model safety. They advocate for responsible disclosure of discovered vulnerabilities to model developers.
• Future Work: Potential extensions include adapting the framework for image generators, agents, and indirect prompt injections, as well as exploring alternative diversity metrics.

In conclusion, Adaptive Instruction Composition represents a significant advancement in automated LLM security, moving beyond random or purely iterative search to a mathematically grounded, adaptive approach that efficiently uncovers the full spectrum of model vulnerabilities.