Information Theoretic Adversarial Training of Large Language Models
arXiv:2605.05415v1 [cs.LG] 6 May 2026
Yiwei Zhang Purdue University West Lafayette, IN 47907, USA [email protected] Rouzbeh Behnia University of South Florida Tampa, FL 33620, USA [email protected]
Jeremiah Birrell Texas State University San Marcos, TX 78666, USA [email protected]
Jason Pacheco University of Arizona Tucson, AZ, USA [email protected]
Reza Ebrahimi University of South Florida Tampa, FL 33620, USA [email protected]
Elisa Bertino Purdue University West Lafayette, IN 47907, USA [email protected]
Abstract 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.
1
Introduction
Despite recent advances in alignment and safety, large language models (LLMs) remain vulnerable to adversarial prompting and can exhibit harmful or unintended behaviors under novel attack strategies. Recent work highlights that standard training and alignment procedures often fail to anticipate such failure modes, leaving models exposed to unforeseen adversarial inputs (Anil et al., 2024). Moreover, empirical evidence shows that harmful behaviors can persist even after safety fine-tuning, and that adversarial training is necessary to meaningfully improve robustness (Sheshadri et al., 2025). While effective, adversarial training of LLMs tends to be computationally expensive and prohibitive. Continuous adversarial training methods, such as Continuous Adversarial Training (CAT), Continuous Adversarial Preference Optimization (CAPO), have revealed a promising direction for improving the robustness of LLMs by replacing discrete token-level perturbations with continuous, gradient-based Preprint.
Figure 1: Overview of WARDEN with adaptive DRO reweighting. Left: A base continuous adversarial training method generates embedding-space perturbations and per-sample adversarial losses. Right: WARDEN replaces uniform aggregation with an f -divergence DRO objective that dynamically upweights high-loss adversarial examples via a learnable or optimized dual variable λt . For KL divergence, the dual objective reduces to a log-sum-exp form while leaving the underlying continuous attack pipeline unchanged. adversarial attacks that are efficient for LLMs (Xhonneux et al., 2024). This continuous relaxation enables scalable adversarial training that better aligns with the optimization dynamics of modern LLMs. Subsequent work, such as MixAT, further improves this paradigm by combining continuous and discrete adversarial training (Dékány et al., 2025). Nevertheless, existing adversarial training objectives still typically aggregate per-sample adversarial losses uniformly, which can underemphasize rare but high-loss adversarial examples. Building on the work of (Xhonneux et al., 2024), we introduce WARDEN (Worst-case Adversarial Reweighting via f -DivergENces), a distributionally robust optimization (DRO) framework for adversarial training of LLMs based on f -divergence reweighting (as illustrated in Figure 1). Rather than minimizing the empirical average adversarial loss, WARDEN minimizes the worst-case expected adversarial loss over distributions within a divergence ball around the empirical training distribution. Using the convex dual formulation, the objective becomes a log-sum-exp form under KL divergence, and the dual variable λ controls the strength of reweighting. We explore two strategies for λ: Learnable: λ is treated as a trainable parameter. Optimized: λ is computed at each iteration using the bisection method derived from the DRO dual. Our work makes several major contributions to adversarial training of LLMs: • We propose WARDEN, a modular DRO-based reweighting framework for continuous adversarial training of LLMs, yielding a class of information-theoretic objectives that emphasize high-loss adversarial examples. • We derive a tractable KL-DRO objective with a log-sum-exp form and study practical strategies for controlling the reweighting strength through the dynamical dual variable λ. • Empirically, WARDEN reduces attack success rates across multiple LLMs and attack settings while maintaining utility and computational costs comparable to strong continuous adversarial training baselines.
2
Related Work
Closely relevant work to this study can be categorized into three major streams. Continuous Adversarial Training in the LLM’s embedding space. Adversarial training has been known as a leading approach for improving robustness of AI models; however, for LLMs, it is often impractical due to the high computational cost of generating discrete attacks during training. 2
Xhonneux et al. (2024) addressed this limitation by proposing CAPO, an adversarial variant of Identity Preference Optimization (IPO) (Azar et al., 2024) and the CAT algorithm that perform adversarial attacks in the LLM’s continuous embedding space, making the process significantly more efficient. While CAT requires a utility dataset to prevent collapsing to degenerate states that refuse safe prompts, CAPO does not rely on any utility dataset. Highlighting the effectiveness of continuous adversarial training of LLMs, (Fu et al., 2025) reveal theoretical and practical evidence that adversarial training with adversarial suffixes of a certain length yields defending against attacks with quadratically larger lengths. In a subsequent work, Dékány et al. (2025) propose MixAT to further improve the performance of continuous adversarial training by incorporating some discrete adversarial attacks in the token space while maintaining a runtime that is comparable to CAPO and CAT. More recently, Shaopeng et al. (2026) offers a theoretical analysis from the in-context learning lens that sheds light on why adversarial perturbations in the embedding space help LLMs defend against jailbreak prompts in the original input token space. They also show that the behavior of continuous adversarial training can be improved by adding a regularizer inspired by in-context learning. Our work is closest to this line, but differs in focus: rather than modifying the attack generation process, WARDEN changes how per-sample adversarial losses are aggregated by replacing uniform averaging with distributionally robust reweighting. Adversarial training objectives and reweighting. A broader body of work studies the generalization, stability, and robustness–utility trade-offs of adversarial training (Wang et al., 2024; Xiao et al., 2024; Li & Li, 2025; Altinisik et al., 2025; Xie & Huo, 2024; Zhang et al., 2024b), while other methods improve adversarial training through objective or optimization design, such as perturbation reweighting, self-distillation, robust distillation, critical-layer fine-tuning, and multi-norm robustness objectives (Zhang et al., 2024a; Cho et al., 2025; Lee et al., 2025; Gopal et al., 2025; Jiang & Singh, 2024). Most related to WARDEN are methods that treat training signals non-uniformly, including pixel-level perturbation reweighting, long-tailed adversarial training, and vulnerable-data-aware adversarial training (Zhang et al., 2024a; Cho et al., 2025; Feng et al., 2025). WARDEN shares the intuition that difficult adversarial examples should receive greater emphasis, but realizes this through an f -divergence DRO objective for LLM adversarial training, yielding a lightweight loss-aggregation layer that leaves the continuous attack pipeline unchanged. Orthogonal LLM defenses. Other LLM defense mechanisms are complementary to our approach. Latent adversarial training perturbs internal model states to mitigate harmful behaviors without relying on explicit failure-inducing prompts (Sheshadri et al., 2025). Shadow-LLM defenses instead use an auxiliary LLM to protect a target model at inference time (Wang et al., 2025). In contrast, WARDEN is a training-time method that directly improves the robustness of the target LLM and can potentially be combined with these orthogonal defenses.
3
WARDEN: Distributionally Robust Adversarial Training of LLMs with Worst-Case Adversarial Reweighting
Continuous adversarial training methods improve LLM robustness by generating adversarial perturbations in the embedding space and optimizing the model against the resulting adversarial losses. However, existing objectives typically aggregate per-sample adversarial losses uniformly. This average-case aggregation can underemphasize rare but high-loss adversarial examples, even though such examples often correspond to the most informative failure modes for robust alignment. We address this limitation by replacing uniform aggregation with a distributionally robust reweighting objective that emphasizes hard adversarial samples while preserving the underlying continuous attack pipeline. 3.1
Worst-Case Adversarial Reweighting
Continuous adversarial training for LLMs generates embedding-space perturbations and trains the model to prefer desired responses over undesired ones under these perturbations. Given a distribution (x, y, ŷ) ∼ Pn of harmful prompts x, desired responses y, and undesired responses ŷ, and (continuous-attack) adversarial perturbation δ(x, ŷ) targeted at ŷ, the standard preference-based objective (Xhonneux et al., 2024) minimizes inf E(x,y,ŷ)∼Pn [Lθ (x, y, ŷ)] , θ
3
(1)
where 2 fθ (y|x + δ(x, ŷ)) fθ (ŷ|x + δ(x, ŷ)) 1 Lθ (x, y, ŷ) := −ℓβ log − log , ℓβ (h) := h − fθ0 (y|x) fθ0 (ŷ|x) 2β and fθ0 denotes the original model. To prevent the LLM from collapsing to degenerate behaviors and refusing safe prompts, this loss implicitly minimizes the Kullback-Leibler divergence with regards to the original model distribution fθ0 (y|x). Instead of minimizing the empirical average in Eq. (1), WARDEN minimizes a worst-case expected loss over distributions close to the empirical distribution by introducing a computationally inexpensive DRO-reweighting layer, which hasP the effect of focusing the training on the more difficult adversarial n samples. Specifically, let Pn = n1 i=1 δ(xi ,yi ,ŷi ) denote the empirical distribution training set. For an f -divergence Df , we define the distributionally robust adversarial objective inf
sup
θ∈Θ Q:Df (Q∥Pn )≤ϵ
{E(x,y,ŷ)∼Q [Lθ (x, y, ŷ)] − κDf (Q∥Pn )}.
(2)
The radius ϵ limits how far the reweighting distribution Q may move from Pn , while κ > 0 adds a soft-constraint divergence penalty which acts as a regularization term for numerical stability. The inner maximization therefore emphasizes high-loss adversarial examples on which the current model is most vulnerable. Using convex duality, Eq. (2) can be written as Lθ (x, y, ŷ) − ρ ∗ inf inf λϵ + ρ + (λ + κ)E(x,y,ŷ)∼Pn f , θ∈Θ λ≥0, λ+κ ρ∈R
where f ∗ denotes the Legendre transform of f . The derivation is provided in Appendix A. This dual form shows that WARDEN can be implemented as a lightweight loss-aggregation layer on top of existing continuous adversarial training methods, involving two additional dynamical parameters, λ and ρ. 3.2
KL-DRO Objective and Adaptive Reweighting
In this work, we focus on the KL-divergence instantiation of Eq. 2. For Df (Q∥Pn ) = KL(Q∥Pn ), the dual objective simplifies to inf
sup
{E(x,y,ŷ)∼Q [Lθ (x, y, ŷ)] − κKL(Q∥Pn )}
θ∈Θ Q:KL(Q∥Pn )≤ϵ
(3)
Lθ (x, y, ŷ) = inf inf λϵ + (λ + κ) log E(x,y,ŷ)∼Pn exp . θ∈Θ λ≥0 λ+κ This log-sum-exp objective interpolates between average-case and worst-case training: larger λ + κ approaches uniform averaging, while smaller values concentrate weight on high-loss samples. Hence, λ controls the reweighting strength. In practice, WARDEN uses adaptive treatments for setting λ. We consider two main variants. In the learnable variant, λ is optimized jointly with the model parameters. In the optimized variant, λ is recomputed at each training iteration by solving the one-dimensional convex dual problem in Eq. (3). Since the objective is convex in λ, this update can be performed efficiently with a bisection solver; see Appendix C. In both adaptive variants, the resulting λt changes across training iterations and minibatches, allowing WARDEN to adjust the amount of reweighting to the current adversarial loss distribution. For completeness, we also evaluate a fixed-λ setting in the ablation study, where λ is held constant as a hyperparameter. We do not treat this setting as a primary variant of WARDEN, since it lacks minibatch-adaptive control of the DRO reweighting strength and performs substantially worse in our experiments; see Section 5.2. 3.3
DRO Reweighting with Utility Preservation
Some adversarial training methods include a utility dataset, Du to preserve benign instructionfollowing ability and prevent excessive refusal. In this setting, we similary apply DRO reweighting 4
only to the adversarial loss and add the standard utility loss separately. For the continuous attack of Xhonneux et al. (2024), we have the following KL-divergence variant with utility regularization Lθ (x, y, ŷ) inf inf λϵ + (λ + κ) log E(x,y,ŷ)∼Pn exp − E(xu ,yu )∼Du [log fθ (yu |xu )] , θ∈Θ λ≥0 λ+κ (4) Lθ (x, y, ŷ) := log fθ (ŷ|x + δ(x, ŷ)) − log fθ (y|x + δ(x, ŷ)) , which should be compared with Eq. (4) of Xhonneux et al. (2024). This objective preserves the base adversarial loss while replacing its uniform empirical average with KL-DRO aggregation. The same DRO-reweighting mechanism can be added on top of stronger base adversarial training pipelines - MixAT method; see Eq. (7) in Dékány et al. (2025). In the KL case this again has the form (4), only with the loss function changed to Lθ (x, y, ŷ) := log fθ (ŷ|x̂) − log fθ (y|x̂) ,
(5)
where x̂ is the adversarial sample generated by the MixAT attack: x̂ = argmaxx′ ∈N (x) log fθ (ŷ|x′ ) .
(6)
Thus, WARDEN is modular: it does not modify the adversarial example generation procedure, optimizer, perturbation budget, or utility objective of the base method. It only replaces the aggregation of per-sample adversarial losses with a distributionally robust objective. For more general f -divergences, similar methods can be developed by generalizing (2). However, in this work we focus on the KL variants.
4
Experimental Details
4.1
LLM Models
We evaluate WARDEN on four open-source instruction-tuned LLMs: Zephyr-7B (Hugging Face H4, 2023), Mistral-7B (mistralai, 2025), Llama2-7B (meta-llama, 2023), and Llama3-8B (meta-llama, 2024). These models span different model families and safety-tuning recipes, allowing us to assess whether DRO reweighting provides consistent gains across models with varying inherent robustness. For each model, we compare the undefended instruction-tuned model, the corresponding continuous adversarial training baselines, and their WARDEN-augmented variants. We denote the latter as CAT-WARDEN, CAPO-WARDEN, and MixAT-WARDEN, with suffixes L and O indicating learnable and optimized treatments of the dual variable λ. CAT-WARDEN and CAPO-WARDEN follow the CAT/CAPO setup of Xhonneux et al. (2024), while MixAT-WARDEN follows Dékány et al. (2025). When available, we include both released checkpoints and reproduced baselines using public code, denoted by HF and R. Each WARDEN variant is initialized from the same base model as its corresponding baseline and differs only in adversarial-loss aggregation. 4.2
Datasets
To isolate the effect of WARDEN, we follow the data protocol of each base method. CAT-WARDEN and CAPO-WARDEN use the CAT/CAPO adversarial prompt–response data from Xhonneux et al. (2024), while MixAT-WARDEN uses the adversarial and paraphrase-augmented data from Dékány et al. (2025). Utility datasets, when used by the base method, are kept unchanged. For robustness evaluation, we use HarmBench (Mazeika et al., 2024) and evaluate against direct harmful requests, human-written jailbreaks, AutoDAN, and GCG. Following MixAT (Dékány et al., 2025), we evaluate on a fixed subset of 40 non-copyright-related HarmBench samples due to the high cost of optimization-based attacks. For utility evaluation, we report MMLU, ARC-Easy, ARC-Challenge, and the harmless-query benchmark from Xhonneux et al. (2024), which measures benign instruction-following on 40 simple queries. 5
4.3
Training Setup and Hyperparameters
We implement WARDEN on top of the public code and configurations of the corresponding baselines. All base training components are kept unchanged, including attack generation, perturbation budget, inner attack steps, optimizer, learning-rate schedule, batch size, utility-loss weight, and training budget. The only modification is replacing the empirical average of per-sample adversarial losses with the proposed KL-DRO aggregation. We tune only the DRO-specific parameters: the KL radius ϵ, soft-constraint coefficient κ, and treatment of λ. In WARDENL , λ is learned jointly with the model; in WARDENO , it is recomputed at each step by solving the one-dimensional convex dual problem with bisection. For each model, we select hyperparameters based on the validation robustness–utility trade-off. Full hyperparameters are listed in Appendix E. All evaluations use the same decoding settings, attack implementations, and benchmark protocols across methods. Robustness is measured by attack success rate for each attack type and on average; utility is measured on knowledge, reasoning, and harmless-query benchmarks. Evaluation data are not used for training or hyperparameter selection. 4.4
Hardware
All experiments were performed on the internal and external clusters using G4, 40GB A100, or 80GB A100 GPUs. All conducted experiments required at least 1,162 GPU hours.
5
Results
We evaluate whether WARDEN improves the safety–utility trade-off of continuous adversarial training. For each base method, we compare the original adversarial training baseline with its WARDENaugmented variant across four instruction-tuned LLMs and four attack categories. Overall, WARDEN reduces attack success rates for CAT, CAPO, and MixAT variants in most settings, while generally preserving competitive utility. 5.1
Main Results
Table 1 reports robustness and utility results on Zephyr-7B-β, Mistral-7B-Instruct-v0.1, Llama-2-7BChat, and Llama-3-8B-Instruct. The main trend is that WARDEN improves continuous adversarial training most when the base method leaves substantial residual vulnerability. This is most evident for CAPO: CAPO-WARDENO reduces average ASR from 22.62% to 12.13% on Zephyr-7B-β and from 44.12% to 20.63% on Mistral-7B, while maintaining comparable average utility on Mistral-7B. These gains are consistent with the KL-DRO objective, which upweights high-loss adversarial examples and therefore targets the vulnerable tail that remains after standard adversarial training. WARDEN also improves CAT and MixAT, with smaller absolute gains when the corresponding baselines are already strong. On Mistral-7B, CAT-WARDENL and CAT-WARDENO reduce average ASR from 10.00% to 5.62% and 5.75%, respectively. On Zephyr-7B-β, CAT-WARDENL reduces average ASR to 3.88%, and MixAT-WARDENL improves an already robust MixAT baseline from 1.25% to 0.00% ASR. On Llama-2-7B, where several baselines are already highly robust, improvements are correspondingly smaller: CAT-WARDENO reduces average ASR from 10.63% to 6.50%, while CAPO-WARDEN variants preserve near-zero ASR and improve average utility from 64.91% to over 67%. The main exception is that stronger robustness can occasionally come with over-refusal. On Llama3-8B, CAT-WARDENL reaches 0.00% average ASR but reduces harmless-query performance to 65.00%. In contrast, CAT-WARDENO preserves utility close to CAT (R) while reducing ASR from 6.25% to 4.50%, and MixAT-WARDENL gives the best trade-off in this block, reducing ASR from 3.125% to 1.25% while improving average utility from 80.71% to 83.21%. Overall, the results support our hypothesis that DRO-based adversarial-loss aggregation improves robustness by emphasizing high-loss adversarial examples. The gains are largest when the base trainer leaves a heavy vulnerable tail and smaller when the baseline is already near saturation. 6
Table 1: Safety–utility evaluation across models and attack settings.
Llama3-8B
Llama2-7B
Mistral-7B
Zephyr-7B
Model
Variant
Attack Success Rate (ASR, %) ↓
Utility (%) ↑
DR
HH
AD
GCG
Avg.
MMLU
ARCe
ARCc
Hless
Avg.
Plain (HF)
82.5
87.5
95
77.5
85.63
59.07
86.74
73.72
100
79.88
CAT (HF) CAT (R) CAT-WARDENL CAT-WARDENO
2.5 2.5 2.5 2.5
2.5 2.5 0.5 2
0 2.5 0 0
15 22.5 12.5 20
5 7.5 3.88 6.25
56 57.05 56.7 56.54
85.56 85.23 84.93 84.85
74.31 73.04 73.29 72.54
97.5 90 100 97.5
78.34 76.33 78.73 77.86
CAPO (R) CAPO-WARDENL CAPO-WARDENO
25 20 17.5
13 24 8.5
30 22.5 17.5
22.5 12.5 5
22.62 19.75 12.13
58.24 56.35 56.53
86.44 82.87 83.88
74.14 70.65 71.84
97.5 55 47.5
79.8 66.22 64.94
MixAT (HF) MixAT-WARDENL
2.5 0
2.5 0
0 0
0 0
1.25 0
57.45 57.43
87.54 87.12
73.63 74.57
97.5 100
79.03 79.78
Plain (HF)
80
77.5
95
87.5
85
54.36
82.65
67.32
100
76.08
CAT (R) CAT-WARDENL CAT-WARDENO
2.5 2.5 2.5
2.5 2.5 0.5
2.5 0 2.5
32.5 17.5 17.5
10 5.62 5.75
53.58 53.74 54.14
81.86 80.89 81.69
67.58 66.89 65.96
100 95 95
75.76 74.13 74.2
CAPO (R) CAPO-WARDENL CAPO-WARDENO
67.5 55 32.5
29 19.5 7.5
25 27.5 15
55 52.5 27.5
44.12 38.63 20.63
51.46 53.32 52.38
80.93 81.57 81.82
65.27 66.55 66.98
97.5 95 95
73.79 74.11 74.05
Plain (HF)
2.5
5
5
32.5
11.25
45.72
72.26
55.97
100
68.49
CAT (HF) CAT (R) CAT-WARDENL CAT-WARDENO
5 2.5 0 2.5
5 2.5 1 1
5 5 7.5 2.5
30 32.5 22.5 20
11.25 10.63 7.75 6.5
45.74 45.02 46.05 44.74
71.21 71.17 71 71
54.1 54.1 54.78 54.27
92.5 97.5 97.5 97.5
65.89 66.95 67.33 66.88
CAPO (R) CAPO-WARDENL CAPO-WARDENO
0 0 0
0 1 0.5
0 0 0
0 0 0
0 0.25 0.125
45.84 45.12 45.22
73.02 71.89 71.97
55.8 54.27 54.01
85 97.5 97.5
64.91 67.2 67.18
Plain (HF)
2.5
7.5
0
10
5
64.94
91.41
81.05
100
84.35
CAT (HF) CAT (R) CAT-WARDENL CAT-WARDENO
0 2.5 0 0
2.5 0 0 0.5
0 0 0 0
2.5 22.5 0 17.5
1.25 6.25 0 4.5
63.47 64.18 63.58 64
90.91 90.66 90.91 90.4
78.75 79.69 79.61 79.86
85 100 65 100
79.53 83.63 74.78 83.57
MixAT (HF) MixAT-WARDENL
0 0
0 0
0 0
12.5 5
3.125 1.25
63.52 63.46
90.31 90.53
79.01 78.84
90 100
80.71 83.21
(HF) model released on HuggingFace; (R) re-trained model or prompt using public code; Plain denotes the original basic model without any defense Attack success rate (ASR) ↓ indicates lower is better for safety robustness, while utility metrics ↑ indicate higher is better DR: Direct Request; HH: Human Jailbreaks; AD: AutoDAN; GCG: Greedy Coordinate Gradient. MMLU: Massive Multitask Language Understanding; ARCe: AI2 Reasoning Challenge (Easy); ARCc: AI2 Reasoning Challenge (Challenge); Hless: Harmless dataset from Xhonneux et al. (2024). Cells are color-scaled within each model block and metric column, excluding Plain rows.
The improvements across CAT, CAPO, and MixAT further show that WARDEN acts as a modular loss-aggregation layer rather than a method-specific modification. 5.2
Ablation Studies
We conduct ablations to isolate the effect of the two main DRO-specific design choices: the KL ambiguity radius ϵ and the treatment of the dual variable λ. All ablations are performed on Mistral-7B using the same training and evaluation pipeline as in the main experiments. We report attack success rates (ASR) under direct requests, human-written jailbreaks, AutoDAN, and GCG, together with standard utility metrics. 7
(a) ASR under different attack categories.
(b) Utility across evaluation benchmarks.
Figure 2: Sensitivity to the KL-DRO radius ϵ on Mistral-7B CAPO-WARDENO . Moderate values of ϵ improve robustness, with the lowest average ASR around ϵ = 0.1, while utility remains relatively stable across the evaluated range.
Effect of the DRO radius ϵ. The radius ϵ controls how far the adversarial reweighting distribution may deviate from the empirical minibatch distribution. Smaller values keep the objective close to uniform averaging, whereas larger values allow the objective to place more mass on high-loss adversarial examples. Figure 2 shows a non-monotonic robustness trend. Increasing ϵ initially improves robustness, with the lowest average ASR achieved around ϵ = 0.1. However, larger values degrade ASR, suggesting that overly aggressive reweighting can overemphasize a small set of high-loss examples and reduce training effectiveness. In contrast, utility remains relatively stable across the evaluated range. These results indicate that moderate distributional robustness is beneficial, but the ambiguity radius must be tuned to obtain the best robustness–utility trade-off.
(a) λ = 5 / λ0 = 5.
(b) λ = 1 / λ0 = 1.
(c) ASR–utility trade-off.
Figure 3: Ablation of dual-variable handling strategies on Mistral-7B. The fixed variant holds λ constant; WARDENL learns λ from the shown initialization λ0 ; and WARDENO recomputes λ by solving the dual problem at each step. Optimizing λ yields the lowest average ASR with comparable utility.
Effect of the dual variable λ. We compare three strategies for handling the KL-DRO dual variable λ: keeping it fixed, learning it jointly with the model, and optimizing it at each training step by solving the one-dimensional convex dual problem. For the fixed variant, λ is held constant throughout training; for the learnable variant, the same value is used as the initialization λ0 ; for the optimized variant, λ is recomputed from the dual objective at each step and is therefore independent of the initialization. Figures 3a and 3b show that the treatment of λ has a substantial effect on robustness. Fixed-λ variants yield the highest average ASR, while learning λ improves robustness. The optimized variant performs best, reducing average ASR to approximately 20.6%, compared with 36.9% for the learnable variant and 53.6% for the fixed variant in the λ = 5 setting. This suggests that solving the dual problem provides a more effective minibatch-specific reweighting strength than using a predefined or slowly adapted value. 8
Figure 4 further explains this gap through the training dynamics. The learnable strategy changes λ gradually from its initialization, reflecting a global gradient-based adaptation process coupled to the model update. In contrast, the optimized strategy recomputes λ⋆ at each DRO checkpoint and rapidly moves toward smaller values. Since smaller λ makes the KL-DRO log-sum-exp aggregation sharper, WARDENO can more quickly enter a worst-case-oriented reweighting regime when the minibatch contains hard adversarial examples. The corresponding loss curves show that both variants reduce the training objective over time, although the batch-level loss remains noisy with occasional spikes, as expected in adversarial training.
(a) Trajectory of the DRO dual variable.
(b) Training loss dynamics.
Figure 4: Training dynamics under learnable and optimized dual-variable treatments on Mistral-7B. WARDENO rapidly decreases the optimized λ⋆ , inducing sharper adversarial reweighting, whereas WARDENL changes λ more gradually. Both variants reduce the training loss, with noisy spikes typical of adversarial training. Robustness–utility trade-off. Figure 3c summarizes the three λ treatments in the ASR–utility plane. The optimized variant achieves the lowest ASR while maintaining comparable average utility, whereas the fixed variant remains substantially less robust. This indicates that the improvement of WARDENO is not primarily due to sacrificing benign performance, but rather to selecting a sharper and more appropriate minibatch-specific reweighting strength. Together with the ϵ ablation in Figure 2, these results show that DRO reweighting must be sufficiently strong to emphasize vulnerable examples, but not so strong that it over-concentrates on a small set of high-loss samples. This supports the central motivation of WARDEN: replacing uniform adversarial-loss aggregation with distributionally robust aggregation improves adversarial training, provided that the reweighting strength is adaptively controlled.
6
Conclusion
We presented WARDEN, a distributionally robust framework for adversarial training of LLMs. WARDEN replaces uniform aggregation of per-sample adversarial losses with an f -divergence DRO objective that emphasizes high-loss adversarial examples. In the KL case, convex duality yields a tractable log-sum-exp objective whose reweighting strength is controlled by a dual variable λ. This makes WARDEN a lightweight and modular loss-aggregation layer that can be applied to existing continuous adversarial training methods without changing their attack generation pipeline. Across multiple instruction-tuned LLMs and attack settings, WARDEN improves robustness for CAT, CAPO, and MixAT variants while generally preserving utility. The gains are largest when the base method retains nontrivial residual vulnerability, supporting the motivation that robust alignment benefits from focusing training on hard adversarial examples. Our ablations show that intermediate KL radii and adaptive optimization of λ provide the best robustness–utility trade-off. Limitations. WARDEN inherits the coverage and quality limitations of the adversarial data and base attack pipeline: reweighting hard observed examples does not guarantee robustness to substantially different unseen attacks. It also introduces DRO-specific hyperparameters, including ϵ, κ, and the treatment of λ, which must be selected carefully to avoid weak reweighting or overemphasis on a small set of high-loss samples. Finally, our robustness evaluation follows prior work and uses a fixed 9
HarmBench subset due to the cost of optimization-based attacks; broader evaluation on larger attack suites and larger models remains future work. Broader Impact. This work aims to improve LLM robustness against adversarial prompting and reduce harmful outputs under jailbreak attacks. Because WARDEN is modular and lightweight, it may make robust alignment more practical across deployment settings. However, adversarial training can also increase over-refusal or degrade benign helpfulness if the safety–utility trade-off is not carefully evaluated. WARDEN should therefore be used as part of a broader safety pipeline that includes diverse red-teaming, utility and refusal-rate evaluation, monitoring for adaptive attacks, and transparent reporting of remaining failure modes.
10
References A. Ahmadi-Javid. Entropic value-at-risk: A new coherent risk measure. Journal of Optimization Theory and Applications, 155:1105–1123, 2012. Syed Mumtaz Ali and Samuel D Silvey. A general class of coefficients of divergence of one distribution from another. Journal of the Royal Statistical Society: Series B (Methodological), 28 (1):131–142, 1966. Enes Altinisik, Safa Messaoud, Husrev Taha Sencar, Hassan Sajjad, and Sanjay Chawla. Explaining the role of intrinsic dimensionality in adversarial training. In International Conference on Machine Learning, pp. 1298–1313. PMLR, 2025. Cem Anil, Esin Durmus, Nina Panickssery, Mrinank Sharma, Joe Benton, Sandipan Kundu, Joshua Batson, Meg Tong, Jesse Mu, Daniel Ford, et al. Many-shot jailbreaking. Advances in Neural Information Processing Systems, 37:129696–129742, 2024. Mohammad Gheshlaghi Azar, Zhaohan Daniel Guo, Bilal Piot, Remi Munos, Mark Rowland, Michal Valko, and Daniele Calandriello. A general theoretical paradigm to understand learning from human preferences. In International Conference on Artificial Intelligence and Statistics, pp. 4447–4455. PMLR, 2024. Aharon Ben-Tal and Marc Teboulle. An old-new concept of convex risk measures: The optimized certainty equivalent. Mathematical Finance, 17(3):449–476, 2007. doi: 10.1111/j. 1467-9965.2007.00311.x. URL https://onlinelibrary.wiley.com/doi/abs/10.1111/j. 1467-9965.2007.00311.x. Jeremiah Birrell, Paul Dupuis, Markos A Katsoulakis, Yannis Pantazis, and Luc Rey-Bellet. (f, Γ)divergences: Interpolating between f-divergences and integral probability metrics. Journal of machine learning research, 23(39):1–70, 2022. Michel Broniatowski and Amor Keziou. Minimization of divergences on sets of signed measures. Studia Scientiarum Mathematicarum Hungarica, 43(4):403–442, 2006. Seungju Cho, Hongsin Lee, and Changick Kim. Long-tailed adversarial training with self-distillation. arXiv preprint arXiv:2503.06461, 2025. Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457, 2018. Imre Csiszár. On information-type measure of difference of probability distributions and indirect observations. Studia Sci. Math. Hungar., 2:299–318, 1967. Csaba Dékány, Stefan Balauca, Robin Staab, Dimitar I Dimitrov, and Martin Vechev. MixAT: Combining continuous and discrete adversarial training for LLMs. Advances in Neural Information Processing Systems, 2025. Yuqi Feng, Jiahao Fan, and Yanan Sun. Vulnerable data-aware adversarial training. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. Shaopeng Fu, Liang Ding, and Di Wang. ”Short-length” adversarial training helps LLMs defend ”Long-length” jailbreak attacks: Theoretical and empirical evidence. In ICLR 2025 Workshop on Foundation Models in the Wild, 2025. URL https://openreview.net/forum?id= U74MXMriLw. Bhavna Gopal, Huanrui Yang, Jingyang Zhang, Mark Horton, and Yiran Chen. Boosting adversarial robustness with clat: Criticality leveraged adversarial training. In Forty-second International Conference on Machine Learning, 2025. Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. International Conference on Learning Representations, 2021. 11
Hugging Face H4. Zephyr-7b-β. https://github.com/huggingface/alignment-handbook, 2023. Hugging Face model checkpoint. Enyi Jiang and Gagandeep Singh. Ramp: Boosting adversarial robustness against multiple l_p perturbations for universal robustness. Advances in Neural Information Processing Systems, 37: 43759–43787, 2024. Hongsin Lee, Seungju Cho, and Changick Kim. Indirect gradient matching for adversarial robust distillation. In 13th International Conference on Learning Representations, ICLR 2025, pp. 49625–49646. International Conference on Learning Representations, ICLR, 2025. Binghui Li and Yuanzhi Li. Adversarial training can provably improve robustness: Theoretical analysis of feature learning process under structured data. In The Thirteenth International Conference on Learning Representations, 2025. F. Liese and I. Vajda. On divergences and informations in statistics and information theory. IEEE Transactions on Information Theory, 52(10):4394–4412, 2006. D.G. Luenberger. Optimization by Vector Space Methods. Professional Series. Wiley, 1997. ISBN 9780471181170. URL https://books.google.com/books?id=M5n9DwAAQBAJ. Mantas Mazeika, Long Phan, Xuwang Yin, Andy Zou, Zifan Wang, Norman Mu, Elham Sakhaee, Nathaniel Li, Steven Basart, Bo Li, et al. Harmbench: A standardized evaluation framework for automated red teaming and robust refusal. arXiv preprint arXiv:2402.04249, 2024. meta-llama. Llama-2-7b-chat-hf. https://huggingface.co/meta-llama/ Llama-2-7b-chat-hf, 2023. Hugging Face model checkpoint. meta-llama. Meta-llama-3-8b-instruct. https://huggingface.co/meta-llama/ Meta-Llama-3-8B-Instruct, 2024. Hugging Face model checkpoint. mistralai. Mistral-7b-instruct-v0.1. https://huggingface.co/mistralai/Mistral-7B-v0.1, 2025. Hugging Face model checkpoint. X. Nguyen, M. J. Wainwright, and M. I. Jordan. Estimating divergence functionals and the likelihood ratio by convex risk minimization. IEEE Transactions on Information Theory, 56(11):5847–5861, 2010. J. Ponstein. Approaches to the Theory of Optimization. Cambridge Tracts in Mathematics. Cambridge University Press, 2004. ISBN 9780521604918. URL https://books.google.com/books? id=GaNB2B677wgC. Shaopeng, Di Fu, and Wang. Understanding and improving continuous LLM adversarial training via in-context learning theory. Internatial Conference on Learning Representations, 2026. Abhay Sheshadri, Aidan Ewart, Phillip Huang Guo, Aengus Lynch, Cindy Wu, Vivek Hebbar, Henry Sleight, Asa Cooper Stickland, Ethan Perez, Dylan Hadfield-Menell, and Stephen Casper. Latent adversarial training improves robustness to persistent harmful behaviors in LLMs. Transactions on Machine Learning Research, 2025. ISSN 2835-8856. URL https://openreview.net/forum? id=6LxMeRlkWl. Xunguang Wang, Daoyuan Wu, Zhenlan Ji, Zongjie Li, Pingchuan Ma, Shuai Wang, Yingjiu Li, Yang Liu, Ning Liu, and Juergen Rahmel. SelfDefend: LLMs can defend themselves against jailbreaking in a practical manner. In 34th USENIX Security Symposium (USENIX Security 25), pp. 2441–2460, 2025. Yunjuan Wang, Kaibo Zhang, and Raman Arora. Benign overfitting in adversarial training of neural networks. In Forty-first International Conference on Machine Learning, 2024. Sophie Xhonneux, Alessandro Sordoni, Stephan Günnemann, Gauthier Gidel, and Leo Schwinn. Efficient adversarial training in LLMs with continuous attacks. Advances in Neural Information Processing Systems, 37:1502–1530, 2024. Jiancong Xiao, Jiawei Zhang, Zhi-Quan Luo, and Asuman Ozdaglar. Uniformly stable algorithms for adversarial training and beyond. arXiv preprint arXiv:2405.01817, 2024. 12
Yiling Xie and Xiaoming Huo. High-dimensional (group) adversarial training in linear regression. Advances in Neural Information Processing Systems, 37:31708–31735, 2024. Jiacheng Zhang, Feng Liu, Dawei Zhou, Jingfeng Zhang, and Tongliang Liu. Improving accuracyrobustness trade-off via pixel reweighted adversarial training. In Proceedings of the 41st International Conference on Machine Learning, pp. 59382–59402, 2024a. Kaibo Zhang, Yunjuan Wang, and Raman Arora. Stability and generalization of adversarial training for shallow neural networks with smooth activation. Advances in Neural Information Processing Systems, 37:16160–16193, 2024b.
A
DRO Dual Formulation
In this appendix we provide a derivation of the DRO reformulation result via convex duality; the argument is similar to the proof of Theorem 5.1 in Ahmadi-Javid (2012), with the main difference being the inclusion of the f -divergence soft-constraint penalty term κDf (Q∥Pn ), which we require for numerical stability purposes; background on f -divergences can be found in Appendix A.1. Theorem A.1. Let Z be a measurable space and Pn be the empirical distribution of a collection of samples zi ∈ Z, i = 1, ..., n. Further, suppose: 1. We have 0 ≤ a < 1 < b ≤ ∞ and a convex f : (a, b) → R that satisfies f (1) = 0. 2. We have a measurable loss L : Z → R. Then for all ϵ > 0, κ > 0 we have the equality {EQ [L] − κDf (Q∥Pn )}
sup
(7)
Q:Df (Q∥Pn )≤ϵ
∗ L−ρ = inf λϵ + ρ + (λ + κ)EPn f , λ≥0, λ+κ ρ∈R
where f ∗ denotes the Legendre transform of f . In the KL case we have the further simplification sup
{EQ [L] + κKL(Q∥Pn )}
(8)
Q:KL(Q∥Pn )≤ϵ
= inf {λϵ + (λ + κ) log EPn [exp(L/(λ + κ))]} . λ≥0
Moreover, the objective function on the right-hand side of (7) is convex in (λ, ρ) and the objective on the right-hand side of (8) is convex in λ. Proof. We will employ the Slater condition for convex duality, see, e.g., Theorem 3.11.2 in Ponstein (2004) and Theorem 8.3.1 and Problem 8.7 in Luenberger (1997). To that end, let V be the vector space of finite signed measures on Z and define W to be subset of probability measures satisfying Df (Q∥Pn ) < ∞. Next recall that the f -divergences are convex, which is a consequence of their variational representation (Broniatowski & Keziou, 2006; Nguyen et al., 2010; Birrell et al., 2022). This implies that W is a convex subset of V . Define the convex functional F : W → R by F [Q] = EQ [−L] + κDf (Q∥Pn ); we note that Q ∈ W implies Q ≪ Pn , and so Q is supported on {zi }ni=1 , which implies F [Q] is finite as claimed. Convexity of F follows from convexity of the f -divergence and linearity of the integral. Finally, note that the convex inequality constraint Df (Q∥Pn ) ≤ ϵ is strictly satisfied at Q = Pn ∈ W , i.e., Df (Pn ∥Pn ) = 0 < ϵ. Thus we have confirmed the Salter condition and can conclude the strong duality result inf {EQ [−L] + κDf (Q∥Pn )} = sup −λϵ + inf {EQ [−L] + (λ + κ)Df (Q∥Pn )} . Q:Df (Q∥Pn )≤ϵ
λ≥0
(9)
Q:Df (Q∥Pn )<∞
Finally, for all λ ≥ 0 the inner minimization over Q can be evaluated by using the Gibbs variational formula for f -divergences, see Theorem 4.2 in Ben-Tal & Teboulle (2007) (note that the requirement 13
that the minimum of f equal zero is an inessential normalization condition; see, e.g, Birrell et al. (2022)). Specifically, after multiplying both sides by −1 this yields
{EQ [L] − κDf (Q∥Pn )}
sup
(10)
Q:Df (Q∥Pn )≤ϵ
( = inf
) λϵ + (λ + κ)
λ≥0
{EQ [L/(λ + κ)] − Df (Q∥Pn )}
sup
Q:Df (Q∥Pn )<∞
= inf
λ≥0
∗
λϵ + (λ + κ) inf {ν + EPn [f (L/(λ + κ) − ν)]} ν∈R
.
Changing variables in the inner minimization from ν to ρ = (λ + κ)ν, we arrive at (7). The KL ∗ result (8) then follows from the fact that fKL (y) = ey−1 , which allows the optimization over ρ to be computed explicitly via a straightforward calculus argument. Convexity in (λ, ρ) of the objective on the right-hand side of (7) follows from applying standard convex analysis results, including the convexity of the Legendre transform, f ∗ , and convexity of the perspective of a convex function. Finally, convexity of the objective on the right-hand side of (8) follows from the fact that minimizing a jointly convex function over one of its arguments results in a convex function in the remaining argument.
A.1
Background on f -Divergences
The f -divergences, introduced by Ali & Silvey (1966); Csiszár (1967), quantify the discrepancy between a pair of probability distributions, Q and P , by using the likelihood ratio dQ/dP , i.e., by comparing the relative weights under Q and P . More precisely, for any convex function f on the real line with f (1) = 0, the corresponding f -divergence is defined by
Df (Q∥P ) :=
EP [f (dQ/dP )] ∞
if Q ≪ P , otherwise
(11)
where Q ≪ P denotes absolute continuity, i.e., existence of the likelihood ratio. The most widely used f -divergence is the KL divergence (i.e., relative entropy), which is defined in terms of fKL (t) := t log(t). For further discussions of the properties of f -divergences see, e.g., Liese & Vajda (2006); Birrell et al. (2022). 14
B
Algorithm Pseudocode
Algorithm 1 Distributionally Robust LLM Adversarial Training Require: Adversarial data {(xi , yi , ŷi )}ni=1 , utility data Du , target LLM fθ , base adversarial method BaseAT, number of training iterations N , minibatch size B, number of inner attack iterations M , WARDEN radius ϵ, soft regularization parameter κ, dual variable initialization λ0 , WARDEN mode mλ ∈ {learnable, optimized}, learning rates ηδ , ηλ , ηθ . Ensure: Robustified LLM fθ . 1: Initialize λ ← λ0 . 2: for n = 1, . . . , N do 3: Sample adversarial minibatch Bn . 4: Sample utility minibatch Un , if used. 5: for (xi , yi , ŷi ) ∈ Bn do 6: Initialize adversarial perturbation δi . 7: for m = 1, . . . , M do 8: Update perturbation using the base attack: δi ← Π∆ (δi + ηδ ∇δ Aattack (xi , yi , ŷi , δi , θ)) . 9: 10: 11:
end for Set x̃i ← xi + δi . Compute per-sample adversarial loss: Ladv ← LossBaseAT (fθ , x̃i , yi , ŷi ). i
12: 13: 14:
end for adv Let Ladv = (Ladv 1 , . . . , LB ). Auxiliary update for the WARDEN dual variable λ: KL adv Π[λmin ,λmax ] λ − ηλ ∇λ ADRO (L , λ) , λ← adv arg min AKL , λ), DRO (L λ≥0
15:
mλ = learnable, mλ = optimized.
Compute the KL-DRO adversarial objective: B
adv AKL , λ) = λϵ + (λ + κ) log DRO (L
16: 17:
if utility data is used then Compute standard utility loss: Lu = −
18: 19: 20: 21: 22:
1 X exp B i=1
Set total loss: else Set total loss:
1 |Un |
X
log fθ (y|x).
(x,y)∈Un
adv L = AKL , λ) + αu Lu . DRO (L
adv L = AKL , λ). DRO (L
end if Update model parameters: θ ← θ − ηθ ∇θ L.
23: end for 24: return Robustified LLM fθ .
15
Ladv i λ+κ
! .
C
Optimization over λ
In this appendix we detail the implementation of the mλ = optimized case of Algorithm 1, line 14. Here, the goal is to solve adv arg min AKL , λ) , DRO (L
(12)
λ≥0
where the objective is given by B
adv AKL , λ) = λϵ + (λ + κ) log DRO (L
1 X exp B i=1
Ladv i λ+κ
! .
(13)
Recall that ϵ, κ > 0 and Ladv are the values of the loss on the adversarial samples. i adv We start by noting that AKL , λ) is smooth and convex in λ ∈ [0, ∞) (see Theorem A.1). Now DRO (L consider the following two cases. adv Case 1: ∂λ |λ=0 AKL , λ) < 0 DRO (L adv First note that limλ→∞ AKL , λ) = ∞. This implies the existence of λr > 0 such that DRO (L KL adv ∂λ |λ=λr ADRO (L , λ) > 0. Combined with the condition that defines this case, we can therefore adv conclude existence of λ∗ ∈ [0, λr ] such that ∂λ |λ=λ∗ AKL , λ) = 0. Hence, by standard DRO (L convex analysis theory, the critical point λ∗ is a global minimizer. adv Case 2: ∂λ |λ=0 AKL , λ) ≥ 0 DRO (L Convexity implies that the derivative is non-decreasing, hence in this case we can conclude that adv AKL , λ) is non-decreasing in λ ∈ [0, ∞). Therefore the minimum is achieved at λ∗ = 0. DRO (L
Based on these two cases, our approach for solving (12) is to first check the sign of adv ∂λ |λ=0 AKL , λ). If it is non-negative then, according to Case 2, we return λ∗ = 0. If it DRO (L is negative then we use the bisection method to search for a critical point which, per Case 1, will be an optimizer. Specifically, we initialize a variable λr > 0 and increase it until we observe a adv sign change in the derivative, i.e., until ∂λ |λ=λr AKL , λ) > 0 (as noted in Case 1, such a DRO (L sign change is guaranteed to occur). Subsequently we apply a standard bisection method solver adv (e.g., optimize.bisect in SciPy) to return a solution, λ∗ , to ∂λ AKL , λ) = 0 in the interval DRO (L [0, λr ]. Finally, we note the following explicit formula for the required derivative: adv PB Li 1 adv ! B Ladv exp λ+κ X i i=1 B L 1 i KL adv −1 adv . exp ∂λ ADRO (L , λ) = ϵ + log − (λ + κ) PB L 1 B i=1 λ+κ exp i B
i=1
λ+κ
(14) This formula illustrates the importance of using a positive κ, as it prevents a singularity at λ = 0.
D
Datasets, Models, and Licenses
We use publicly available datasets, model checkpoints, and baseline implementations. We credit the original creators through citations in the main text and summarize the sources and licenses of the main assets below. We do not redistribute third-party model weights, datasets, or baseline code as part of this submission. Training and evaluation data. For adversarial training, CAT-WARDEN and CAPO-WARDEN use the adversarial prompt–response data from Xhonneux et al. (2024), while MixAT-WARDEN follows the adversarial and paraphrase-augmented data protocol of Dékány et al. (2025). When a utility dataset is used by the corresponding base method, we keep it unchanged. For robustness evaluation, we use HarmBench (Mazeika et al., 2024) and evaluate on direct harmful requests, human-written jailbreaks, AutoDAN, and GCG attacks. For utility evaluation, we use MMLU (Hendrycks et al., 2021), ARC-Easy and ARC-Challenge (Clark et al., 2018), and the harmless-query benchmark from Xhonneux et al. (2024). Models. We evaluate WARDEN on publicly available instruction-tuned LLMs. Table 2 lists the model checkpoints used in this work. We use these models for research evaluation only and do not redistribute their weights. 16
Table 2: Model checkpoints used in this work. Licenses and access terms should be verified against the corresponding official model cards before final submission. Model Source License / terms Zephyr-7B-β (Hugging Face H4, 2023) Mistral-7B-Instruct-v0.1 (mistralai, 2025) Llama-2-7B-Chat (meta-llama, 2023) Llama-3-8B-Instruct (meta-llama, 2024)
Hugging Face H4 Mistral AI Meta Meta
MIT License Apache-2.0 License Llama 2 Community License Meta Llama 3 Community License
Baseline implementations. CAT-WARDEN and CAPO-WARDEN follow the CAT/CAPO implementation and training setup of Xhonneux et al. (2024), while MixAT-WARDEN follows the MixAT implementation and setup of Dékány et al. (2025). Where public code or checkpoints are available, we use them only for research reproduction and comparison, and cite the original papers and repositories. We do not redistribute third-party baseline code or checkpoints.
E
Adversarial Training Hyperparameters
Table 3: DRO-specific hyperparameters for WARDEN variants. Subscripts L and O denote learnable and optimized treatments of the dual variable λ, respectively. All non-DRO hyperparameters are inherited from the corresponding base method. Model
Variant
λ treatment
λ0 / solver setting
λ-LR
KL radius ϵ
Soft penalty κ
Zephyr-7B
CAT-WARDENL CAT-WARDENO CAPO-WARDENL CAPO-WARDENO MixAT-WARDENL
learnable optimized learnable optimized learnable
5 bisection 5 bisection 5
0.0001 – 0.0001 – 0.001
0.05 0.1 0.08 0.1 0.1
0.08 0.08 0.053 0.07 0.1
Mistral-7B
CAT-WARDENL CAT-WARDENO CAPO-WARDENL CAPO-WARDENO
learnable optimized learnable optimized
5 bisection 5 bisection
0.0001 – 0.001 –
0.3 0.05 0.1 0.1
0.1 0.1 0.3 0.1
Llama2-7B
CAT-WARDENL CAT-WARDENO CAPO-WARDENL CAPO-WARDENO
learnable optimized learnable optimized
5 bisection 5 bisection
0.00001 – 0.001 –
0.8 0.5 0.2 0.1
0.5 0.1 0.1 0.1
Llama3-8B
CAT-WARDENL CAT-WARDENO MixAT-WARDENL
learnable optimized learnable
5 bisection 5
0.00001 – 0.001
0.55 0.5 0.1
0.5 0.1 0.1
We tune only the hyperparameters introduced by WARDEN: the KL ambiguity radius ϵ, the soft regularization coefficient κ, and the treatment of the dual variable λ. All other training hyperparameters, including the optimizer, learning-rate schedule, batch size, perturbation budget, number of inner attack steps, utility-loss weight, sequence length, and total training budget, are inherited from the corresponding base adversarial training method. Specifically, CAT-WARDEN and CAPO-WARDEN follow the CAT/CAPO configurations of Xhonneux et al. (2024), while MixAT-WARDEN follows the MixAT configuration of Dékány et al. (2025). For each model and base method, we select DRO-specific hyperparameters using validation robustness–utility trade-off. Table 3 reports the selected configurations used in the main experiments. In WARDENL , λ is initialized at λ0 and optimized jointly with the model using a separate dual learning rate. In WARDENO , λ is recomputed at each training step by solving the one-dimensional convex KL-DRO dual problem with a bisection solver.
17