Closing the Distribution Gap in Adversarial Training for LLMs

Chengzhi Hu; Jonas Dornbusch; David Lüdke; Stephan Günnemann; Leo Schwinn

cs.LGFeb 16, 2026

Closing the Distribution Gap in Adversarial Training for LLMs

Authors: Chengzhi Hu, Jonas Dornbusch, David Lüdke, Stephan Günnemann, Leo Schwinn

Organizations: Technical University of Munich, Germany.

Abstract

Adversarial training for LLMs is one of the most promising methods to reliably improve robustness against adversaries. However, despite significant progress, models remain vulnerable to simple in-distribution exploits, such as rewriting prompts in the past tense or translating them into other languages. We argue that this persistent fragility stems from a fundamental limitation in current adversarial training algorithms: they minimize adversarial loss on their training set but inadequately cover the data distribution, resulting in vulnerability to seemingly simple attacks. To bridge this gap, we propose Distributional Adversarial Training, DAT. We leverage Diffusion LLMs to approximate the true joint distribution of prompts and responses, enabling generation of diverse, high-likelihood samples that address generalization failures. By combining optimization over the data distribution provided by the diffusion model with continuous adversarial training, DAT achieves substantially higher adversarial robustness than previous methods.

Figures & tables

Figure 1: Standard AT minimizes the empirical robust risk over a fixed dataset

\mathcal{D}

(brown), which provides a poor approximation of the population robust risk. This results in a distribution gap where the model remains vulnerable to the manifold of natural language

q

(blue). Specifically, standard methods fail to cover the distribution of prompts

\tilde{q}(x\mid y_{harm})

(green) that are likely to trigger harmful responses. Our DAT framework bridges this gap by optimizing over a surrogate distribution defined by a diffusion LLM

p^{\mathrm{diff}}_{\theta}(x|y_{\mathrm{hamr}})

(purple), allowing the model to train on a distribution that more closely matches the true population.

Figure 2: Cumulative transfer ASR across five target models (Gemma3-12B ( Gemma Team et al., 2025 ) , Qwen2.5-7B ( Qwen et al., 2025 ) , Zephyr-7B ( Tunstall et al., 2023 ) , Llama3-8B-LAT ( Sheshadri et al., 2024 ) , Llama3-8B-CB ( Zou et al., 2024a ) ) from attacks on Llama3-8B. Diffusion-based Inpainting attacks exhibit significantly higher transferability than model-specific optimization (GCG) or heuristic perturbations (BoN), suggesting that conditional sampling from the diffusion surrogate effectively identifies data-specific vulnerabilities that generalize across architectures and defenses.

Figure 3: Diversity of generated attack strings measured using SBERT embeddings (all-MiniLM-L6-v2; Reimers and Gurevych, 2019 ). Each cell reports the mean pairwise cosine similarity between samples generated by two methods. Diffusion-based attacks exhibit the lowest intra-method similarity (0.178), indicating substantially greater sample diversity than GCG and BoN.

Model	Method	ASR $\downarrow$						Utility $\uparrow$
Model	Method	GCG	Pair	BoN	Direct	Inpainting	BoA	XSTest	MMLU	ARCe	ARCc
Llama3-8B	Default	$0.39$	$0.3$	$0.99$	$0.33$	$0.98$	$1$	$0.660$	$0.638$	$0.815$	$0.530$
	CAT	$0.04$	$0.13$	$0$	$0$	$0.88$	$0.88$	$0.464$	$0.637$	$0.810$	$0.520$
	LAT	$0.09$	$0.15$	$0$	$0.01$	$0.94$	$0.94$	$0.444$	$0.619$	$0.784$	$0.485$
	CB	$0.03$	$0.09$	$0.56$	$0.13$	$0.92$	$0.94$	$0.672$	$0.636$	$0.815$	$0.533$
	MixAT-GCG	$0.04$	$0.03$	$0.02$	$0.01$	$0.77$	$0.77$	$0.456$	$0.612$	$0.827$	$0.513$

Table 1: Main results on Llama3-8B and Qwen2.5-14B across robustness and utility benchmarks. Lower ASR indicates stronger robustness, while higher utility scores indicate better model helpfulness and capabilities.

Figure 4: Pareto frontier for Llama3-8B showing the trade-off between Inpainting Robustness (

1-\text{ASR}

) and XSTest compliance rate. DAT achieves superior trade-offs across all hyperparameter settings.

Figure 5: Inpainting ASR as a function of the number of unique diffusion samples

M

per behavior in the training data. Robustness gradually improves as we better approximate the population risk

q(x,y)

, while helpfulness stays consistent.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Base Model	Adapter	HF Source
Llama3-8B	–	meta-llama/Meta-Llama-3-8B-Instruct
	CB	GraySwanAI/Llama-3-8B-Instruct-RR
	LAT	LLM-LAT/robust-llama3-8b-instruct
	MixAT-GCG	INSAIT-Institute/Llama3-8B-MixAT-GCG
Qwen2.5-14B	–	Qwen/Qwen2.5-14B-Instruct
Qwen2.5-14B	MixAT-GCG	INSAIT-Institute/Qwen-14B-MixAT-GCG

Appendix

Table 2: Sources of Hugging Face models and adapters.

Hyperparameter	Llama3-8B	Qwen2.5-14B
Learning rate	$1\times 10^{-4}$	$1\times 10^{-4}$
Batch size	8	8
Number of epochs	6	6
Optimizer	AdamW	AdamW
Adv. learning rate	$1\times 10^{-1}$	$1\times 10^{-1}$
Attack-steps	20	40

Appendix

Table 3: Hyperparameters for models trained with DAT .

Hyperparameter	Llama3-8B	Qwen2.5-14B
Learning rate	$2\times 10^{-4}$	$1\times 10^{-4}$
Batch size	16	16
Number of epochs	6	6
Optimizer	AdamW	AdamW
Attack-radius	$1\times 10^{7}$	$1\times 10^{7}$
Adv. learning rate	$1\times 10^{-1}$	$1\times 10^{-1}$

Appendix

Table 4: Hyperparameters for models trained with CAT.

Adversarial Training

Adversarial training improves model reliability by training on examples or conditions designed to expose weaknesses.

See all 183 papers

Information Theoretic Adversarial Training of Large Language Models
May 6, 2026Yiwei Zhang, Jeremiah Birrell, Reza Ebrahimi +3Adversarial Training LLM Alignment
Purdue University · West Lafayette, IN 47907, USA · Texas State University +5
Robustifying Vision-Language Models via Test-Time Prompt Adaptation
Jul 10, 2026Xingyu Zhu, Huanshen Wu, Shuo Wang +4VLM Robustness VLM Adaptation
National University of Singapore · University of Science and Technology of China · The Hong Kong University of Science and Technology
Understanding and Improving Continuous Adversarial Training for LLMs via In-context Learning Theory
Apr 14, 2026Shaopeng Fu, Di WangAdversarial Training Adversarial Robustness of Language Models
Provable Responsible AI and Data Analytics (PRADA) Lab, King Abdullah University of Science and Technology, Saudi Arabia

May 6, 2026J/K move · Enter open · S save

Information Theoretic Adversarial Training of Large Language Models

Yiwei Zhang, Jeremiah Birrell, Reza Ebrahimi, Rouzbeh Behnia+2

Purdue University · West Lafayette, IN 47907, USA · Texas State University+5

Adversarial Training LLM Alignment Adversarial Attacks on LLMs Distributionally Robust Optimization Language Model Robustness

Open paper AskPDF arXiv

Large language models (LLMs) remain vulnerable to adversarial prompting despite advances in alignment and safety, often exhibiting harmful behaviors under novel attack strategies. While adversarial training can improve robustness, existing approaches are computationally expensive and difficult to scale. Recent continuous adversarial training methods, such as Continuous adversarial training (CAT) and Continuous Adversarial Preference Optimization (CAPO), address this challenge by leveraging gradient-based perturbations in the embedding space, enabling more efficient and expressive attacks. Building on this paradigm, we propose WARDEN, a distributionally robust adversarial training framework for LLMs that dynamically reweights adversarial examples through an f -divergence ambiguity set around the empirical training distribution. Our method optimizes the worst-case adversarial loss within a divergence ball around the empirical data distribution, automatically emphasizing harder adversarial examples. Using the convex dual formulation, the objective reduces to a log-sum-exp form under the KL divergence, with a dynamical parameter controlling the strength of reweighting. This study leads to a new class of information-theoretic objectives that significantly reduce attack success rates while maintaining model utility. Across multiple LLMs and attack settings, WARDEN substantially reduces attack success rates with computational and utility costs comparable to CAT-, CAPO-, and MixAT-based baselines, making it a practical approach for scalable robust alignment.

Abstract

Adversarial Training

Explore similar work

Information Theoretic Adversarial Training of Large Language Models

Robustifying Vision-Language Models via Test-Time Prompt Adaptation

Understanding and Improving Continuous Adversarial Training for LLMs via In-context Learning Theory