Closing the Distribution Gap in Adversarial Training for LLMs

1. Problem Statement

Current Adversarial Training (AT) methods for Large Language Models (LLMs) suffer from a fundamental limitation: they minimize empirical robust risk over a fixed, finite dataset but fail to adequately approximate the population robust risk.

• The Distribution Gap: Standard AT algorithms effectively address model-specific perturbations (optimization errors within a local ball) but neglect data-specific generalization failures.
• The Symptom: Models trained with standard AT often remain vulnerable to simple, natural language manipulations that lie just outside the training distribution, such as:
  • Rewriting prompts in the past tense.
  • Translating prompts into low-resource languages.
  • Rephrasing requests in ways that were not explicitly present in the training set.
• The Root Cause: The training data ($D$) is a poor approximation of the true underlying data distribution ($q$). Consequently, models fail to generalize to valid, natural inputs that trigger harmful responses but were not seen during training.

2. Methodology: Distributional Adversarial Training (DAT)

The authors propose Distributional Adversarial Training (DAT), a framework designed to bridge the gap between empirical and population robust risk by leveraging Diffusion LLMs as generative surrogates.

Core Components

1. Generative Surrogate (Diffusion LLMs):

   • Unlike standard autoregressive models that learn $q(y|x)$ (conditional on input), Diffusion LLMs model the joint distribution $q(x, y)$.
   • This allows the model to solve the inverse problem: sampling diverse prompts $x$ conditioned on a specific harmful response $y$ (i.e., sampling from $q(x|y_{harm})$).
   • This process generates data-specific adversarial examples that naturally elicit harmful responses, covering the support of the true distribution rather than just model-specific artifacts.

2. Two-Stage Training Objective:
   DAT unifies generative sampling with continuous adversarial optimization:

   • Outer Loop (Distribution Approximation): Uses the Diffusion LLM to sample diverse, high-likelihood prompts $x$ conditioned on harmful targets $y$. This expands the training set to better approximate the population risk.
   • Inner Loop (Adversarial Optimization): Applies Continuous Adversarial Training (CAT) to these generated samples. This involves maximizing the loss over continuous perturbations $\delta$ in the embedding space to find the worst-case scenario for each sample.
   • Regularization: A KL-divergence term is added to the loss function to prevent the model from collapsing and to maintain general utility (helpfulness) on benign data.

Mathematical Formulation

The objective function minimizes the expected worst-case loss over the diffusion-generated distribution:
$$L_{DAT} = \mathbb{E}_{y \sim D_{harm}, x \sim p_{diff}(x|y)} \left[ \max_{\delta \in \Delta} L_{adv}(x + \delta, y, y_{refusal}) \right] + \lambda_{KL} L_{KL}$$

3. Key Contributions

1. Formalization of the Robustness Gap: The paper distinguishes between model-specific perturbations (optimized attacks like GCG) and data-specific generalization failures (natural language variations). It argues that current AT fails because it ignores the latter.
2. Novel Training Framework (DAT): Introduces a method that combines generative sampling from Diffusion LLMs with continuous adversarial optimization.
3. Theoretical Justification: Provides a theorem (Surrogate Fidelity Bound) proving that if the diffusion surrogate accurately approximates the conditional distribution $q(x|y)$, the gap between empirical and population robust risk is bounded.
4. Empirical Superiority: Demonstrates that DAT significantly outperforms state-of-the-art baselines in worst-case robustness while maintaining model utility.

4. Experimental Results

The authors evaluated DAT on Llama3-8B and Qwen2.5-14B against a suite of attacks including GCG, PAIR, BoN, Direct Attacks, and Diffusion Inpainting (a data-specific attack).

• Robustness (ASR - Attack Success Rate):

  • Inpainting Attack: Standard defenses (CAT, LAT, Circuit Breakers) failed catastrophically against Inpainting attacks (ASR > 77%). DAT reduced this to 32% (Llama3) and 14% (Qwen2.5).
  • Best-of-All (BoA) Robustness: DAT achieved the lowest BoA ASR (36% for Llama3, 18% for Qwen2.5), significantly outperforming the next best method (MixAT-GCG at 77% and 75% respectively).
  • Generalization: DAT successfully defended against both model-specific (GCG) and data-specific (Inpainting, translation, tense changes) attacks.

• Utility (Helpfulness & Capabilities):

  • DAT maintained competitive utility scores on XSTest (helpfulness), MMLU, and ARC benchmarks.
  • It achieved a superior Pareto frontier, offering higher robustness for the same level of compliance compared to baseline methods.

• Ablation Studies:

  • Diffusion-only vs. DAT: Using diffusion sampling alone improved robustness against Inpainting but failed against GCG. Combining it with continuous AT (DAT) was necessary for full robustness.
  • Sample Quantity: Increasing the number of unique diffusion samples per behavior linearly improved robustness, validating the theoretical claim that better distribution approximation reduces risk.
  • Data Specificity: Training on low-likelihood (poorly generated) samples failed to improve robustness, proving that the quality and distributional fidelity of the generated data are critical.

5. Significance

This work represents a paradigm shift in LLM safety research:

• From Optimization to Distribution: It moves the focus from purely optimizing against specific attack algorithms (inner loop) to ensuring the training data covers the true distribution of potential threats (outer loop).
• Solving the "Simple Attack" Problem: It addresses the persistent vulnerability of LLMs to simple linguistic variations (e.g., tense changes, translation) that current defenses miss.
• Scalability: By using Diffusion LLMs to generate training data, the method offers a scalable way to approximate the population risk without requiring massive, manually curated datasets of every possible attack variation.

In conclusion, DAT provides a theoretically grounded and empirically validated solution to the generalization gap in adversarial training, significantly enhancing the reliability of LLMs against a broader spectrum of adversarial threats.