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 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 q~(x∣yharm) (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θdiff(x∣yhamr) (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 ↓
Utility ↑
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−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
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×10−4
1×10−4
Batch size
8
8
Number of epochs
6
6
Optimizer
AdamW
AdamW
Adv. learning rate
1×10−1
1×10−1
Attack-steps
20
40
Appendix
Table 3: Hyperparameters for models trained with DAT .
Hyperparameter
Llama3-8B
Qwen2.5-14B
Learning rate
2×10−4
1×10−4
Batch size
16
16
Number of epochs
6
6
Optimizer
AdamW
AdamW
Attack-radius
1×107
1×107
Adv. learning rate
1×10−1
1×10−1
Appendix
Table 4: Hyperparameters for models trained with CAT.
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.
Yiwei Zhang, Jeremiah Birrell, Reza Ebrahimi +3
Purdue University · West Lafayette, IN 47907, USA · Texas State University +5
Pre-trained Vision-Language Models (VLMs) such as CLIP achieve strong zero-shot generalization, but their performance degrades sharply under adversarial perturbations. Existing test-time adaptation methods typically rely on sample-level confidence heuristics, overlooking the intrinsic distributional structure of the data. This sample-centric approach limits robustness, as it fails to distinguish confident adversarial mispredictions from true semantic consistency. In this work, we observe that adversarial distortion is structurally brittle: while holistic representations are corrupted, semantic integrity is often preserved in the distribution of augmented views. Motivated by this insight, we propose RITA, a Robust test-tIme prompt-TAdaptation framework that shifts from sample-level estimates to distribution-level alignment. Specifically, RITA employs optimal transport to align the distribution of augmented visual features with textual prototypes, mitigating adversarial outliers and rectifying cross-modal semantic misalignment. Furthermore, we introduce a dynamic cache to progressively accumulate reliable cues from the test stream for online refinement. Extensive experiments demonstrate that RITA significantly improves adversarial robustness without compromising clean accuracy.
Xingyu Zhu, Huanshen Wu, Shuo Wang +4
National University of Singapore · University of Science and Technology of China · The Hong Kong University of Science and Technology
Adversarial training (AT) is an effective defense for large language models (LLMs) against jailbreak attacks, but performing AT on LLMs is costly. To improve the efficiency of AT for LLMs, recent studies propose continuous AT (CAT) that searches for adversarial inputs within the continuous embedding space of LLMs during AT. While CAT has achieved empirical success, its underlying mechanism, i.e., why adversarial perturbations in the embedding space can help LLMs defend against jailbreak prompts synthesized in the input token space, remains unknown. This paper presents the first theoretical analysis of CAT on LLMs based on in-context learning (ICL) theory. For linear transformers trained with adversarial examples from the embedding space on in-context linear regression tasks, we prove a robust generalization bound that has a negative correlation with the perturbation radius in the embedding space. This clearly explains why CAT can defend against jailbreak prompts from the LLM's token space. Further, the robust bound shows that the robustness of an adversarially trained LLM is closely related to the singular values of its embedding matrix. Based on this, we propose to improve LLM CAT by introducing an additional regularization term, which depends on singular values of the LLM's embedding matrix, into the objective function of CAT. Experiments on real-world LLMs demonstrate that our method can help LLMs achieve a better jailbreak robustness-utility tradeoff. The code is available at https://github.com/fshp971/continuous-adv-icl.
Shaopeng Fu, Di Wang
Provable Responsible AI and Data Analytics (PRADA) Lab, King Abdullah University of Science and Technology, Saudi Arabia