Efficient Adversarial Training in LLMs with Continuous Attacks (2024)

Adversarial Attacks

Adversarial attacks and defenses have been extensively studied in the literature[1, 3, 4, 11, 12, 13, 14, 15, 16, 17, 18].More recently, LLMs have been shown to be vulnerable to exploitation by adversarial attacks, and several threat models, such as suffix attacks[1] and jailbreaking[15], have been proposed.Zou etal. [1] present the Greedy Coordinate Gradient (GCG) suffix attack, which generates adversarial examples transferable from small open-source models to large proprietary models.Huang etal. [19] find that just varying generation strategies, such as adjusting decoding hyper-parameters and sampling methods, can trigger harmful behaviour in LLMs.Geisler etal. [20] introduce a novel discrete attack strategy that leverages continuous embedding space optimisation.In the area of continuous adversarial attacks,Fort [21] explore scaling laws for continuous adversarial attacks on language model activations. Further,Schwinn etal. [7, 8] showcase the potential of continuous adversarial attacks as a threat model to compromise safety alignment and unlearning.

An alternative threat model involves jailbreaks, a form of prompt engineering with the goal of circumventing safety alignment.Deng etal. [16] fine-tune an LLM with jailbreak examples and demonstrate that the fine-tuned LLM can generate strong attacks, which transfer between different models.Similarly, Chao etal. [14] found that LLMs could be leveraged to create jailbreaks for other LLMs, even without fine-tuning. They introduced the Prompt Automatic Iterative Refinement (PAIR) algorithm, which uses an attacker algorithm to iteratively query a target LLM, optimising the jailbreak prompt.Liu etal. [15] developed a hierarchical genetic algorithm to generate high-perplexity jailbreaks that can bypass the safety alignments of LLMs.

Adversarial Training

Previous work on continuous adversarial training (AT) on token embeddings has mostly focused on encoder-decoder models, such as BERT[9, 10, 22, 23, 24, 25].Jiang etal. [9] use adversarial attacks to promote smoothness in the embedding space of the model and show that this approach improves generalisation. Similarly,Zhu etal. [10] enforce invariance in the embedding space through adversarial attacks.He etal. [23] combine a disentangled attention mechanism with continuous AT and demonstrate improved generalisation for BERT and RoBERTa models on multiple downstream tasks.Other works apply continuous adversarial perturbation to word embeddings to increase performance in different NLP tasks[22, 24, 25].Robey etal. [26] propose improving the robustness of autoregressive LLMs by a randomised smoothing-inspired approach.

Concurrent to this work,Casper etal. [27] use continuous attacks for the purpose of AT. They propose latent adversarial training (LAT), a method that finds perturbations in the network’s hidden layer representations and applies them to several tasks including text generation. For text generation, they demonstrate that fine-tuning for desirable behaviour with LAT makes the model more likely to forget triggers from data poisoning in some cases.Contrary to our work, they set up the adversarial training in an untargeted manner, i.e.the attack they apply does not aim to produce a particular harmful output but uses the standard AT objective. In contrast, our work focuses on the challenge of making LLMs robust against discrete attacks and jailbreaks while maintaining their helpfulness. To do so, we propose novel algorithms and loss functions that make use of the harmful targets of discrete attacks. Moreover, we thoroughly evaluate across multiple benchmarks and adversarial attacks to ensure a good robustness-utility trade-off.

Adversarial Data Augmentation

Several works [28, 17] have developed adversarial attack generators against LLMs and then used the generated adversarial attacks to create a dataset on which to perform supervised fine-tuning (SFT) to improve adversarial robustness. This kind of adversarial robustness training is based on dataset augmentation and does not adapt the model online to worst-case attacks. Thus, we consider these approaches orthogonal to our work.

3.1 Adversarial Training

3.2 Attack Perturbation Sets in LLMs

For LLMs with a token vocabulary $\mathcal{V}$, $x$ is a prompt and a common perturbation set $T$ are discrete manipulations of the input space, such as suffix attacks[1]. For suffix attacks, the set of acceptable perturbations $\delta$ is defined to be in the set of sequences of tokens of length $m$ that can be appended to the input prompt. In other words, the adversarial attack $x+\delta$ is of the form $x;\delta$, where $\delta$ is a fixed number of tokens the attacker has full control over and ; means concatenation. However, computing the best $\delta$ from this perturbation set $T_{\mathrm{suffix}}(x)=\{\delta\mid x+\delta\in\mathcal{V}^{n+m}\}$ is computationally expensive, as the optimisation turns into a discrete combinatorial problem with exponentially many solutions. Arguably, it is too expensive to use during training, especially for large datasets.

3.3 Adversarial Training in LLMs

Mazeika etal. [6] found this loss necessary to avoid degenerate behaviours such as refusing to answer all prompts by producing some often generic refusal answer $y$.

3.4 Continuous-Adversarial Training

The primary difference between Mazeika etal. [6] and our method is the choice of perturbation set used during AT. Mazeika etal. [6] choose discrete suffix attacks $T_{\mathrm{suffix}}$ and employ the GCG algorithm along with several tricks to mitigate the computational cost to find a GCG attack. One optimisation they introduce is to only update the attack after every $k$ training steps. In contrast, we employ $T_{\mathrm{cont.}}$ with continuous attacks as introduced by Schwinn etal. [7], which are orders of magnitude ($\times 299$) more efficient (see Table1). Consequently, we do not require any additional tricks to further reduce computational costs.In the Unlikelihood loss (Eq3) we add cut-off values for the toward and away loss to prevent over-optimising either. We implement this as $\mathcal{L}=\mathbb{I}[\mathcal{L}^{\prime}>c]0.999c+(\mathbb{I}[\mathcal{L}^{$, where $c$ is the cutoff value chosen.

3.5 Continuous-Adversarial Preference Optimisation

where $\ell(h)$ would be the $\log\sigma(h)$ in the original DPO, but we use the loss proposed in Azar etal. [31] called IPO, i.e.$\ell_{\beta}(h)=\left(h-\frac{1}{2\beta}\right)^{2}$, because it is less prone to overfitting.This loss implicitly minimises the Kullback-Leibler divergence w.r.t.the original model distribution $f_{\theta_{0}}(y|x)$, which prevents the model to collapse to degenerate behaviors leading to refuse all prompts with the refusal answer $y$. As a result, we are able to omit the utility dataset for CAPO.

3.6 Design Decisions

Datasets

For all AT experiments, we utilise the AT dataset from HarmBench[6] with the safe answer $y$ always being Sorry, I can’t do that. As a utility dataset for CAT, we employ UltraChat200k[32, 33], which has been successfully used in both the discrete AT algorithm Zephyr + R2D2[6] and general fine-tuning[34]. For robustness evaluations, we use the first 40 samples of the HarmBench test set. Due to the substantial computational cost associated with LLM adversarial attacks, such as GCG[1], we limit our evaluation to these samples instead of the full test set.

Moreover, we measure the utility of trained models using common benchmarks, including MMLU[35], Arc-E and Arc-C[36], and MT-Bench[37]. To reduce the computational demand, we evaluate $100$ questions for each category for MMLU. Finally, we introduce Harmless which consists of 40 harmless queries (e.g.Tell me a story, see App.G for full list) that are written in the same grammatical style as the Harmbench behaviour. We query the models with their chat template and report the number of refusals (checked manually). Note that only MT-Bench and Harmless use the model’s chat template.

Models

In our experiments, we adversarially fine-tuned four different open-source models Gemma[38], Phi-3-Mini[39], Mistral-7B[40], and Zephyr-7B[34] with increasing parameter counts—2B, 3.8B, 7B, 7B, respectively. We chose instruction-tuned models for all of them. We additionally include Zephyr + R2D2 in our evaluations, which is the Mistral-7B base model fine-tuned with the R2D2 AT algorithm[6]. This results in a diverse set of instruction-tuned models of different sizes. For more details, refer to App.A.2.

Continuous adversarial training

We investigate two novel continuous AT algorithms in this work CAT and CAPO. Due to the computational complexity of fine-tuning LLMs, we do not perform full model fine-tuning for both methods but use LoRA[41] on all linear layers of the transformer architectures. Additionally, we use $4$-bit quantization for all training runs to further reduce the memory overhead. We use $\ell_{2}$ norm perturbations and set the size of the attack $\epsilon$ relatively to the average magnitude of the token embeddings of the respective model. For all models, we use $10$ attack iterations. We set $\epsilon=0.1$ for Gemma and Phi-3-Mini. For Mistral-7B and Zephyr-7B, we set $\epsilon=0.05$ and $\epsilon=0.075$, respectively. For a full list of AT hyperparameters, see App.A.1.

Robustness evaluation

We use three diverse adversarial attacks for the robustness evaluation. GCG, which has shown to achieve one of the highest average attack success rates (ASR) among other state-of-the-art attacks on several models[6]. Since GCG is a suffix attack, we further use AutoDAN and PAIR, which generate more diverse jailbreaks. Furthermore, PAIR has shown high ASR against previous AT approaches in LLMs[6]. To evaluate the ASR, we use the harmfulness classifier from [6], which was shown to align well with human judgement.

Computational cost

Given the constrained computational resources, we prioritised getting evidence to answer our main research question regarding the extrapolation of adversarial robustness. We want to emphasize that better trade-offs between utility and robustness might be obtained with more exhaustive hyperparameter search.

Hardware

All experiments were performed on an internal cluster of either V100, 40GB A100, or 80GB A100 GPUs. All conducted experiments required at least $1904$ GPU hours.

Why do we need continuous adversarial training?

In Table1, we compare the combined number of forward and backward passes used by the discrete AT algorithm RD2D[6] with CAT and CAPO.Computing a single adversarial example with R2D2 is $\approx 128.5$ times more expensive than for CAT and CAPO, while the whole training is $299$ times more costly. This illustrates the considerable compute advantage of continuous AT approaches compared to discrete methods.

LLM adversarial training with utility data

We first explore robustness extrapolation from continuous AT to discrete attacks for the CAT algorithm, which utilises additional utility data to maintain model performance. Figure2 summarises the evaluation results. For all models, CAT considerably increases the average robustness against discrete adversarial attacks. For the Gemma and Zephyr models, robustness increases for all attacks. For Phi-3-Mini and Mistral-7B, PAIR still achieves high attack success rates (ASR).In terms of utility, we observe similar degradations for all CAT trained models.The MMLU and Arc scores decrease marginally, while the MT-Bench score decreases by approximately one.All models still show considerable utility after fine-tuning.

Compared to the Zephyr + R2D2 model, which was trained with discrete AT, CAT exhibits marginally worse utility on standard utility benchmarks while providing substantially improved robustness against discrete attacks. For, Zephyr + R2D2, PAIR achieves an ASR of $40\%$, while it achieves $10\%$ ASR for CAT. We note a substantial difference in the Harmless benchmark, where CAT massively outperforms Zephyr + R2D2 showing that our method has not overfitted the safety objective or the patterns in the Harmbench behaviours. Note that the Harmless score of R2D2 demonstrates that it can not simultaneously achieve non-trivial utility and robustness, which are heavily dependent on not using or using the chat template, respectively.

LLM adversarial training without utility data

We further investigate if adversarial variations of proven alignment methods, such as IPO, can be used to align models in an adversarially robust manner (see Figure2). For this purpose, we fine-tune Gemma and Phi-3-Mini using the proposed CAPO algorithm. Figure2, illustrates differences between the base model, CAT, and CAPO. Despite using no utility dataset within CAPO to retain helpfulness, the algorithm does not introduce larger utility decreases on common benchmarks than CAT. Moreover, CAPO achieves considerably higher robustness against the jailbreaking method PAIR, demonstrating generalisation to diverse threat models. The Phi-3-Mini-CAPO model achieves $100\%$ attack robustness for all conducted attacks. For Gemma, robustness improvements also mostly surpass CAT, with slightly lower robustness against GCG. Compared to R2D2, CAPO does not require an auxiliary dataset to maintain utility and achieves higher robustness on average. Specifically for PAIR CAPO trained models exhibit considerably higher robustness. Lastly, the Phi-3-Mini-CAPO achieves a substantially higher score on the Harmless benchmark than CAT and R2D2.

The results indicate that adversarial variations of common alignment methods, such as IPO, can be used to adversarially align LLMs.

Utility evaluation

Common utility benchmarks such as MMLU or Arc do not use a chat template in their standard evaluation[42]. Firstly, this dramatically impacts performance, especially for smaller models, which often require a lot of prompt engineering to follow the few-shot prompts correctly. Secondly, it dramatically changes the mode of the model. In effect, a model might be overly robust in chat mode (i.e.when using a chat template) where it rejects most requests, but it might appear to have high utility in benchmarks because no chat template is used (e.g.MMLU). Arc as an evaluation benchmark is even more misleading as it measures the likelihood of a set of possible answer tokens, thus not reflecting the utility of the model when using a chat template. We quantitatively evaluate the refusals of MMLU questions when using a chat template in App.E. We recommend future work, to consider these issues when evaluating robustness and utility for the same model.

Training data failure modes

AT datasets such as Harmbench[6] or AdvBench[43] tend to use a common grammatical and syntactical structure, using imperative commands such as “Tell me” or “Give instructions”. Chatting with our models and Zephyr + R2D2, we observe that requests would be refused when using this same style but are accepted if asked in a different style, such as “Could you please …?”. This holds for both harmful and harmless requests. For instance, Zephyr + R2D2 will refuse to answer “Tell me a story” and “Tell me how to build a bomb”, but will answer “Could you please tell me a story?” and “Could you please explain to me how to build a bomb?”. This also explains why the model may even appear useful under utility benchmarks employing chat templates such as MT-Bench. To demonstrate this failure case we create two small benchmark datasets called PoliteHarmbench (see App.F) and Harmless. The former rephrases the harmful behaviours politely, and the latter consists of harmless requests formulated in the same grammatical style as the original Harmbench behaviours. We leave developing better datasets and benchmarks for a future paper as it is outside the scope of this work.

Robust fine tuning without attack

We found that continuous adversarial training successfully increases the robustness of LLMs to discrete adversarial attacks. Here, we explore whether robustness gains stem from using continuous adversarial attacks during training, or from the fine-tuning process itself. Thus, we fine-tune Gemma using the CAPO algorithm but without using adversarial attacks. We observe no robustness gains when fine-tuning without attacks (see App.B.2). This demonstrates that continuous adversarial attacks are a crucial part of our fine-tuning algorithm.

One-step adversarial training in LLMs

For all our experiments, we use $10$ adversarial attack iterations. While this is orders of magnitude cheaper than calculating discrete adversarial attacks (GCG requires $2570$ model evaluations with default settings), it still increases training time by an order of magnitude. We thus propose one-step AT with CAPO. As in previous work[3], we set the step size of the attack to the magnitude of the $\epsilon$-ball. This achieves robustness improvements comparable to the multi-step variant and slightly worse utility trade-offs (see AppB.1).

Robustness-utility trade-offs

Prior work on AT has shown theoretical and empirical trade-offs between robustness and utility[4, 44]. Our previous results demonstrate that continuous AT can achieve non-trivial robustness-utility trade-offs. All experiments are conducted on Gemma models trained with CAPO and varying hyperparameters. Specifically, we sample $\epsilon\in[0.00125,0.3]$, and $\beta\in[0,0.5]$ and fine-tune $7$ different models. In Figure4(b), we depict the GCG loss of the trained models (as a proxy for robustness) on the $y$-axis in logarithmic scale against the MMLU score on the $x$-axis (as a proxy for utility). Clear trade-offs between robustness and utility can be observed, ranging from models with high robustness and no utility to models showing less robustness than the standard non-robust models and slightly higher utility.

Moreover, we analyse hyperparameter choices that affect the robustness-utility trade-off for CAPO in more detail. This includes the strength of the adversarial attacks defined by the $\epsilon$ magnitude and the IPO $\beta$ value.Figure3 illustrates that for both hyperparameters, we obtain intuitive robustness-utility trade-offs, where larger epsilon values and smaller $\beta$ values are associated with increased robustness and reduced utility. A detailed analysis can be found in AppC.

Correlation between continuous attack loss and GCG loss

We additionally investigated the relationship between training-time robustness to continuous adversarial attacks and inference-time robustness to discrete attacks. This is illustrated in Figure4(a). The observed strong Pearson correlation ($r=0.99$, $p=0.0075$) indicates that models robust to continuous attacks during training are also robust to discrete attacks at inference. This suggests continuous AT can be a reliable proxy for AT with discrete attacks. Thus, demonstrating the potential use of continuous attacks to reduce the computational burden of evaluating adversarial robustness[7, 8].

Limitations

Our method relies on the quality and breadth of the harmful dataset, while we are less prone to overfit than Zephyr + R2D2, we may still see improvements from augmented adversarial training datasets [28].An additional limitation is the number of hyperparameters introduced that require careful selection.We expect future work to achieve considerably better robustness-utility trade-offs through better hyperparameter selection alone.Furthermore, our proposed method CAT requires a utility dataset to retain helpfulness, which may shift the predictions of the model on unrelated tasks, a limitation we try to address with the CAPO method. Finally, due to limited compute we were not able to apply our method to much larger LLMs in the 70B parameter and larger regime, we leave this to future work.

Broader impact

This work aims to enable scalable adversarial training for LLMs to be robust against adversarial attacks. The positive impact is that this will reduce the amount of harmful content produced by LLMs if adopted as many attacks will no longer work. In addition, the lower computation cost should hopefully reduce the carbon footprint of training robust and safe LLMs. However, this may lead to overconfidence in the safety of LLMs, thus necessitating more extensive red teaming. Another possible negative impact of our work is that adversarial training may be used to prevent LLMs saying things the model operator does not want regardless of the harmfulness of the content. Our contributions on the failure modes of robustness evaluation should hopefully lead to more rigorous and trustworthy evaluation protocols. These are crucial to accurately assess the state of robustness in LLMs. Note, it may be that further failure modes exist we did not yet find.

Learning rate

Learning rate for the model parameters.

Batch size

Total batch size used for the model training includes utility and behaviours.

Number of epochs

Number of epochs.

Optimiser

Optimiser for the model parameters. AdamW was proposed inLoshchilov and Hutter [45].

Adv.Learning rate

Adversarial learning rate is the step size $\alpha$ used in Equation2.

$\epsilon$

is used to define the $\ell_{2}$ ball around the token embeddings for the valid attacks $\delta$.

$\beta$

is the $\beta$ parameter as described in the original DPO paperRafailov etal. [30].

Away cutoff

is the cut off value used for the away loss as described in §3.3.

Toward cutoff

is the cut off value used for the toward loss as described in §3.3.

Utility data ratio

is the percentage of utility data used as part of the total training data per epoch, e.g.$0.875$ implies for every one adversarial behaviour example there is 8 utility examples.

Away weight

is $\alpha_{a}$ in Equation6.

Toward weight

is $\alpha_{t}$ in Equation6.

Utility weight

is $\alpha_{u}$ in Equation6.

Quantisation

is the level of quantisation for the model during training.

Max seq. length

is the maximum sequence length after which we truncate the token sequences for training.

LoRa

defines where the LoRa adapters are used. For all models we applied the LoRa adapter to all linear layers.

We used a 10 iterations of the adversarial attack, a max grad norm of 0.3, a warm-up ratio of 0.03, a cosine learning rate scheduler, and training was done in floating point 16.

A.1 Adversarial Training

The CAT algorithm has $5$ important hyperparameters, the weight of the utility loss $\alpha_{u}$, toward loss $\alpha_{t}$, and away loss $\alpha_{a}$. Moreover, in preliminary experiments, we observed that away loss tends to dominate the training objective. Models that show very high away loss generally overfitted to the safety objective and stopped answering benign requests. We notice similar issues with the toward loss. Thus, we define a threshold for the away loss $a_{cut}$ and toward loss $t_{cut}$, clamping values below a certain value. If not otherwise defined, we use the following hyperparameters in all experiments. We set $\alpha_{u}=1.0$, $\alpha_{t}=0.5$, and $\alpha_{a}=0.5$, as in[6]. Further, we set $a_{cut}=-5$ and $t_{cut}=0.5$. We use a ratio of $7:1$ for utility and harmful examples during training.

To prevent overfitting in the proposed CAPO, we use the IPO loss function[31]. Additionally, we set the $\beta$ parameter of IPO to $0.25$ for Gemma models, $0.5$ for Phi-3-Mini, and $X$ for Mistral-7B, which we observed to result in good trade-offs between robustness and utility in preliminary experiments.

A.2 Models

Tab.3 summarizes the models used in the experiments of this work.

B.1 One-Step Adversarial Training

As a preliminary experiment for scaling continuous adversarial training, we evaluated if CAPO yields robustness gains if the attack iterations are reduced to one during training. Table5 illustrates that one-step CAPO achieves similar robustness improvements as the multi-step variant. Note, that we used the same hyperparameters for the one-step attacks as for the multi-step attack, except for the attack iterations and step size. Further hyperparameter tuning or borrowing recent advances in one-step AT from other domains may help to close this gap[46]. Due to the large computational complexity of attack evaluations, we conduct this experiment on GCG.

B.2 Training without Attacks

We evaluated if the proposed training algorithms provide robustness without using adversarial attacks during training. Table6 shows, that robustness does not improve without using attacks.

Attack Strength:

The right plot in Figure3 illustrates the effect of varying the adversarial attack strength, characterised by the $\epsilon$ magnitude, on the robustness-utility trade-off. As $\epsilon$ increases from $0.0125$ to $0.1$, there is a significant reduction in GCG loss, from approximately $14.9$ to near $0$. Concurrently, the MMLU score improves markedly from $0$ to around $0.39$, demonstrating increased utility. This inverse relationship between GCG loss and MMLU aligns with prior work concerning utility robustness trade-offs[4, 44].

IPO $\beta$:

In CAPO, the $\beta$ parameter inversely relates to the difference in log-likelihood ratios between the safe answer and the harmful response. Thus, a smaller $\beta$ indicates a larger disparity in these log-likelihood ratios. This intuitively should lead to robustness and utility trade-offs. The left plot in Figure3 shows the impact of different IPO $\beta$ values on robustness and utility. With $\beta$ values ranging from $0$ to $0.5$, a consistent decrease in GCG loss is observed, starting from $6.1$ and dropping to $0.8$. Meanwhile, the MMLU score increases from about $0.25$ to $0.38$. This trend aligns with our expectations and suggests that higher $\beta$ values are associated with lower GCG loss and improved utility, indicating that tuning $\beta$ is crucial for optimizing the robustness-utility trade-off in CAPO.