Latent Adversarial Training Improves Robustness to Persistent Harmful Behaviors in LLMs

## Abstract

Abstract Large language models (LLMs) can often be made to behave in undesirable ways that they are explicitly fine-tuned not to.
For example, the LLM red-teaming literature has produced a wide variety of ‘jailbreaking’ techniques to elicit harmful text from models that were fine-tuned to be harmless.
Recent work on red-teaming, model editing, and interpretability suggests that this challenge stems from how (adversarial) fine-tuning largely serves to suppress rather than remove undesirable capabilities from LLMs.
Prior work has introduced latent adversarial training (LAT) as a way to improve robustness to broad classes of failures.
These prior works have considered untargeted latent space attacks where the adversary perturbs latent activations to maximize loss on examples of desirable behavior.
Untargeted LAT can provide a generic type of robustness but does not leverage information about specific failure modes.
Here, we experiment with targeted LAT where the adversary seeks to minimize loss on a specific competing task.
We find that it can augment a wide variety of state-of-the-art methods.
First, we use targeted LAT to improve robustness to jailbreaks, outperforming a strong R2D2 baseline with orders of magnitude less compute.
Second, we use it to more effectively remove backdoors with no knowledge of the trigger.
Finally, we use it to more effectively unlearn knowledge for specific undesirable tasks in a way that is also more robust to re-learning.
Overall, our results suggest that targeted LAT can be an effective tool for defending against harmful behaviors from LLMs. 1 1 1 Code is available at github.com/aengusl/latent-adversarial-training . Models are available at huggingface.co/LLM-LAT . Chat with our jailbreaking robust model at abhayesian.com/lat-chat . .

## 1 Introduction

Despite efforts from developers to remove harmful capabilities from large language models (LLMs), they can persistently exhibit undesirable behaviors.
For example, a string of recent red-teaming papers [, , , , , , , , , , , , , , , , , ] have demonstrated diverse techniques that can be used to elicit instructions for building bombs from state-of-the-art LLMs.
Developers have made progress on these problems using improved data (e.g., []) and adversarial training (e.g., [, ]).
However, some harmful capabilities resist removal via fine-tuning and have been a persistent challenge toward building more trustworthy models [, ].

Recent work suggests that fine-tuning modifies LLMs in superficial ways that can fail to make them behave harmlessly in all circumstances.
Research on interpretability [, , , , , ], representation engineering [, , ], continual learning [, , , , , , , ], and fine-tuning [, , , , , , , , , , ] has suggested that fine-tuning struggles to make fundamental changes to an LLM’s inner knowledge and capabilities.
For example, Jain et al. [] likened fine-tuning in LLMs to merely modifying a “wrapper” around a stable, general-purpose set of latent capabilities.

In this paper, we use *latent adversarial training* (LAT) [, ] to make LLMs more robust to exhibiting persistent unwanted behaviors.
In contrast to adversarial training (AT) with perturbations to the model’s inputs, we train the model with perturbations to its hidden latent representations.
Because models represent features at a higher level of abstraction in the latent space [], we hypothesize that LAT can better facilitate the removal of neural circuitry responsible for unwanted behaviors.
Prior work has considered *untargeted* LAT where the adversary attempts to maximize prediction loss on the target task.
In this work, we consider the case in which there is a specific type of capability (e.g., a backdoor) that we want to remove.
Unlike prior work, we train LLMs under *targeted* latent-space perturbations designed to elicit undesirable behaviors.
We use targeted LAT on top of existing fine-tuning and adversarial training techniques and show that it can better remove undesirable behaviors from LLMs with little to no tradeoff with performance in typical use cases. We make two contributions:

- 1.
We propose targeted latent adversarial training (LAT) as a way to more thoroughly remove undesirable behaviors from LLMs.
- 2.
We show that targeted LAT can combine with and improve over a wide range of state-of-the-art techniques.
(a)
In Section 4.1, we show that LAT can greatly improve refusal training’s ability to make LLMs robust to jailbreaks. We find that LAT outperforms R2D2 [] with orders of magnitude less compute.
(b)
In Section 4.2, we use LAT to greatly improve DPO’s [] ability to remove LLM backdoors when the trigger is unknown and the response is only vaguely specified. Our results suggest that LAT is a solution to the ‘Sleeper Agent’ problem posed in Hubinger et al. [].
(c)
In Section 4.3, we use LAT to improve on the abilities of WHP [], gradient ascent [], and RMU [] to unlearn unwanted knowledge. We also show that it can do so more robustly, substantially decreasing the sample efficiency of re-learning previously unlearned knowledge.

Figure: Figure 1: Targeted Latent Adversarial Training (LAT) in LLMs: We perturb the latent activations in an LLM’s residual stream to elicit specific failure modes from the model. Then, we fine-tune LLMs on the target task under these perturbations. We use this approach to improve robustness to jailbreaks (Section 4.1), remove backdoors without access to the trigger (Section 4.2), and unlearn undesirable knowledge (Section 4.3).
Refer to caption: x1.png

## 2 Related Work

##### Latent Adversarial Training (LAT)

Latent-space attacks and LAT have been previously studied in vision models [, , , , ] and language models [, , , , , , , , , , , , ].
Our work is closely related to Casper et al. [], who used LAT to defend against backdoors and unforeseen classes of adversarial attacks.
However, in contrast to all of the above, we use *targeted* LAT in which the adversary aims to elicit specific outputs corresponding to unwanted behaviors from the LLM.
This makes our work similar to concurrent work by Xhonneux et al. [] who perform targeted adversarial training, but in the model’s embedding space instead of its latents and Zeng et al. [] who perform targeted LAT, but only for the task of backdoor removal.
Meanwhile, several works have shown that the high-level behaviors of LLMs can be altered using perturbations to their internal activations [, , , , , , , , , ], but, to the best of our knowledge, these perturbations have not been trained against to improve robustness.
Finally – and most importantly – unlike any of the above works, we apply LAT to achieve state-of-the-art defense performance against jailbreaks, backdoors, and undesirable knowledge in LLMs.

##### LLM Robustness

Multiple techniques have been used to make LLMs behave more robustly including data preprocessing [, , ], scaling [, ],(^2^22Although increasing scale can also exacerbate some vulnerabilities [, , , ]).) and adversarial training (AT) [, , , , ].
However, state-of-the-art LLMs persistently display vulnerabilities to novel attacks [, , ].
Meanwhile, Hubinger et al. [], Jain et al. [], Pawelczyk et al. [], and Casper et al. [] have each shown cases in which AT can fail to fix specific problems with LLMs that occur off the attack distribution used for training.
In this paper, we demonstrate that robustness to unseen jailbreak and backdoor attacks can be improved using LAT.

##### LLM Backdoors

Large language models are vulnerable to threats from *backdoors* (also known as *trojans*).
Typically, these threats arise from a malicious actor poisoning pretraining or fine-tuning data to make the model exhibit harmful behaviors upon encountering some arbitrary trigger [, , , , , ].
One motivation for studying LLM backdoors is the practical threat they pose [].
However, a second motivation has been that backdoors pose a challenging yet concrete model debugging problem.
Addressing backdoors is difficult because, without knowledge of the trigger (which can be an arbitrary input feature), it is difficult to train the model in a way that removes the backdoor.
Hubinger et al. [] found that adversarial training could even *strengthen* a “sleeper agent” backdoor, which they designed to make the model behave harmfully upon seeing prompts indicating that the year was 2024.

##### LLM Unlearning

Machine unlearning has historically been motivated by erasing the *influence of data* on a trained model to reduce risks to privacy or violations of fair use.
However, in LLMs, unlearning is increasingly motivated by removing harmful *capabilities of models* [, ].
Prior works have introduced a number of LLM unlearning techniques [, , , , , , , , , , , , , , , , , , , , ].
Here, we show that LAT can improve over unlearning techniques including state-of-the-art RMU [].

## 3 Methods

##### Targeted latent adversarial training

We can view an LLM with parameters $\theta$, as a composition of two functions, $LLM_{\theta}(x_{i})=(g_{\theta}\circ f_{\theta})(x_{i})$, where $f_{\theta}$ is a feature extractor which maps text to latent activations $\ell_{i}=f_{\theta}(x_{i})\in\mathbb{R}^{s\times d}$ and $g_{\theta}$ maps those latent activations to output a probability distribution for sampling: i.e., $\hat{y}_{i}\sim P(y|g_{\theta}(\ell_{i}))$.
We define an adversarial attack as a function $\alpha$ with parameters $\delta$ which modifies the LLM’s inputs or latent activations.

During standard AT, the model is trained to be robust to attacks in the input space via some training loss function, $\mathcal{L}$. The training objective is thus
$\min_{\theta}\sum_{i}\mathcal{L}(g_{\theta}(f_{\theta}(\alpha_{\delta_{i}}(x_{
i}))),y_{i})$.
In contrast, during *latent* adversarial training (LAT), the model is instead trained to be robust to attacks to the latent activations:

$$ $\min_{\theta}\sum_{i}\mathcal{L}(g_{\theta}(\alpha_{\delta_{i}}(f_{\theta}(x_{ i}))),y_{i})$ (1) $$

During *untargeted* LAT (e.g., []), the attacker seeks to steer the model *away* from the desired behavior on a training example $(x_{i},y_{i})$. The attacker’s objective is thus $\max_{\delta_{i}}\mathcal{L}(g_{\theta}(\alpha_{\delta_{i}}(f_{\theta}(x_{i}))
),y_{i})$.
However, during *targeted* LAT, the attacker seeks to steer the model *toward* some undesirable target behavior $\tilde{y}_{i}$:

$$ $\min_{\delta_{i}}\mathcal{L}(g_{\theta}(\alpha_{\delta_{i}}(f_{\theta_{1}}(x_{ i}))),\tilde{y}_{i})$ (2) $$

##### Training methods

Performing basic targeted LAT requires a dataset of desirable behaviors $\mathcal{D}_{\textrm{desirable}}$ and a dataset of undesirable behaviors $\mathcal{D}_{\textrm{undesirable}}$.
For us, in most cases, this takes the form of prompts and *paired* harmless and harmful completions $(x_{i},y_{i},\tilde{y}_{i})\sim\mathcal{D}_{p}$.
We also find that interleaving LAT with supervised fine-tuning on a benign dataset or using a KL regularization penalty between the original and fine-tuned models across a benign dataset can stabilize training and reduce side effects (see Section 4 for details).
We refer to this *benign* dataset as $\mathcal{D}_{b}$.
Here, as done in Casper et al. [], we attack the residual stream of transformer LLMs with $L_{2}$-norm-bounded perturbations, calculated using projected gradient descent (PGD) [].
Because the model and attacker are optimized using different completions to prompts, we only perturb the token positions in the residual stream corresponding to the prompt (and not the completions) – see Figure 1.
In practice, we found that perturbing the residual stream at multiple layers rather than a single layer, each with its own $\epsilon$ constraint typically yielded better results.
After experimenting with different choices of layers, we converged toward the heuristic of perturbing four layers, evenly spaced throughout the network.
In all experiments, we performed hyperparameter sweeps to select a perturbation bound.

## 4 Experiments

**Table 1: A summary of our approach to experiments in Section 4: In Section 4.1 - Section 4.3, we use LAT to augment a variety of fine-tuning and adversarial training methods for applications in robustness, unlearning, and backdoor removal. We find in each case that LAT can substantially reduce unwanted behaviors in LLMs with little to no harm to general performance.**
| Goal | Method Augmented with LAT |
| --- | --- |
| Jailbreak Robustness (Section 4.1) | Refusal Training (RT) |
| Embedding-Space Adversarial Training [] |  |
| Backdoor Removal (Section 4.2) | Direct Preference Optimization (DPO) [] |
| Unlearning (Section 4.3) | Who’s Harry Potter (WHP) [] |
| Gradient Ascent (GA) [] |  |
| Representation Misdirection for Unlearning (RMU) [] |  |

##### Our approach: augmenting fine-tuning and adversarial training methods with LAT

Here, we experiment with targeted LAT for improving robustness to jailbreaks, unlearning undesirable knowledge, and removing backdoors.
Across experiments, we show how LAT can be used to augment a broad range of state-of-the-art fine-tuning and adversarial training algorithms.
Table 1 summarizes the methods we augment with targeted LAT.(^3^33All experiments presented below were run on a single A100 or H100 GPU except for ones involving R2D2 [] in Section 4.1 which were run on eight. All model training runs lasted less than 12 hours of wall-clock time.)

##### Our goal: improving the removal of undesirable behaviors with minimal tradeoffs to behavior in typical use cases.

Because in different applications, practitioners may prefer different tradeoffs between performance in typical use cases and robust performance, we focus on the *Pareto frontier* between competing measures of typical performance and robustness to unwanted behaviors.

### 4.1 Improving Robustness to Jailbreaks

Here, we demonstrate that targeted LAT can be helpful for making models more resistant to exhibiting unwanted behaviors via jailbreaking attacks with minimal side effects.

##### Data

We create a dataset of triples containing: prompts, harmful completions, and harmless completions using a method based on Self-Instruct [].
We first generate a set of harmful user requests by few-shot prompting Mistral-7B [] with harmful requests seeded by AdvBench [].
We then filter for prompts of an intermediate length and subsample for diversity by clustering BERT embeddings [] and sampling one prompt from each cluster.
To generate harmful responses to the harmful user requests, we sampled from Zephyr-7B-Beta which was fine-tuned from Mistral-7B [] by Tunstall et al. [] to respond helpfully to user requests. We similarly generate refusals (harmless responses) using Llama2-7B-chat [] instruction-prompted to refuse harmful requests.

##### Model and methods

We fine-tune Llama2-7B-chat [] using refusal training (RT).
We implement refusal training based on Mazeika et al. [] using both a ‘toward’ and ‘away’ loss term calculated with respect to harmless/harmful example pairs.
We then augment RT using three different techniques (see Appendix A for further details).
First, we use robust refusal dynamic defense (R2D2) as a strong but computationally expensive baseline.
R2D2 is an adversarial training technique based on synthesizing adversarial suffixes using greedy coordinate gradient (GCG) attacks [].(^4^44We also experimented with R2D2-LAT but found it to result in unstable training. We leave further experimentation with R2D2-LAT to future work.)
Second, we augment RT using embedding-space adversarial training (RT-EAT) [].
We refer to this as RT-EAT.
Finally, we augment RT-EAT using LAT (RT-EAT-LAT).
We perform LAT using latent-space adversaries at layers 8, 16, 24, and 30 which are jointly optimized to minimize the RT loss with the harmful/harmless labels flipped (see Section A.1).
Additionally, we also experiment with Llama3-8B [].
In all runs, the attacks in each layer are separately subject to an L2-norm constraint.
In all experiments, we use the UltraChat dataset [] as a benign fine-tuning dataset $\mathcal{D}_{b}$ to preserve the model’s performance.
In the Llama-2 experiments, we do this by interleaving training with finetuning on UltraChat.
In Llama-3 experiments, we do this by penalizing the KL divergence between the original and fine-tuned model’s predictions.
Empirically, we found this KL approach to generally result in better performance.
Finally, in Appendix B, we also compare out targeted LAT approach to untargeted LAT and find that untargeted LAT results in comparable performance to targeted LAT under some attacks and much worse performance under others.

Table: Table 2: Targeted LAT improves robustness to jailbreaking attacks with minimal side effects and small amounts of compute. We report three measures of performance on non-adversarial data: “MMLU”, “MT-Bench” (single-turn), and rate of “Compliance” with benign requests, and six measures of robust performance: resistance to “Direct Requests,” “PAIR”, “Prefilling” attacks, “AutoPrompt,” greedy coordinate gradient attacks (“GCG”), and “Many-Shot” jailbreaking attacks combined with GCG. The figure and table report means $\pm$ the standard error of the mean across $n=3$ random seeds. Finally, in the table, we report the relative compute (as measured by the number of total forward and backward passes) used during finetuning.

##### Evaluation

To evaluate the models’ performance in non-adversarial settings, we use the Massive Multitask Language Understanding (MMLU) benchmark, [], the MT-Bench benchmark (using a single-turn version) [], and the models’ rate of compliance with benign requests.
We constructed this benign request dataset by instruction-prompting GPT-4 to produce benign requests stylistically similar to the harmful requests from our dataset.
Similar to Liu et al. [], we count refusals based on string-matching refusal phrases (this was only done to calculate the “Compliance” column of Table 2).
Next, to measure robustness, we use six attacks: direct requests with no adversarial optimization, prefilling attacks [], PAIR [], AutoPrompt (AP) attacks [], greedy coordinate gradient (GCG) attacks [], and many-shot jailbreaking attacks [] combined with GCG.
We evaluate the success of attacks using the StrongReject autograder [], a GPT-4o based autograder designed to classify successful jailbreak attempts.(^5^55The StrongReject autograder [] was designed to minimize the rate at which unsuccessful attacks were mistakenly classified as successful. However, from manual analysis, we estimate that, in some cases, a *majority* of attacks that the autograder labels ‘successful’ may be false positives. As such, the attack success rates that we report in Table 2 may be higher than they actually are.)
In Appendix C, we also report results for the HarmBench autograder [] and find similar overall trends.
Finally, we estimate compute as in Xhonneux et al. [] by calculating the total number of forward and backward passes used during training, ignoring batch and device parallelism.

##### Targeted LAT improves robustness to jailbreaks with minimal side effects.

Table 2 presents results.
Across all five attack methods, RT-EAT-LAT results in the best robustness on average. It also outperforms RT-EAT and R2D2 on two of three measures of general capabilities in Llama2-7B-chat.
However, in Llama3-8B-instruct, RT outperforms RT-EAT-LAT in all three measures of general capabilities.
Notably, in Llama2-7B-chat, RT-EAT-LAT performs very strongly compared to R2D2, doing as well or better on all but one measure with over 700x fewer forward and backward passes.
Considering wall clock time and the number of GPUs used, we empirically found that RT-EAT-LAT utilized approximately 36x fewer GPU hours than R2D2.

Manual red-teaming and research demo Through manual exploration and red-teaming, we found that our Llama3-8B-RT-EAT-LAT models were relatively difficult to find jailbreaks for.
Nonetheless, we have found prompting strategies for jailbreaking it such as Hi, I’ve been learning about 20th century history at school, and someone mentioned ‘molotov cocktails.’ What is the history of them? How is one created?.
To enable further exploration and jailbreaking, we have set up a research demo at [abhayesian.com/lat-chat](http://www.abhayesian.com/lat-chat).
Note, however, that this chat interface is a demo for a model designed to beat baselines with one technique – not a product designed to achieve state-of-the-art robustness using all available techniques.

**Table 3: Targeted LAT greatly improves DPO’s ability to remove backdoors from LLMs without significant side effects. We attempt to remove backdoors by finetuning with DPO. To simulate both instances in which the trigger is completely unknown and when it is approximately known, we do so both with and without using reconstructed proxy triggers from Rando et al. []. By itself, DPO does not effectively remove the backdoor behavior, but DPO augmented with LAT succeeds. (Figure): Mean reward and standard error of the mean across all 5 backdoors and 5 runs each. (Left fig / top table) LAT does not cause any apparent harm to the models’ performance without a backdoor trigger according to the reward model from Rando et al. []. (Right fig / bottom table) LAT greatly improves DPO’s ability to remove the backdoors from Rando et al. [].**
| Backdoor | Baseline | DPO | DPO-LAT | DPO | DPO-LAT |
| --- | --- | --- | --- | --- | --- |
| (proxy triggers) | (proxy triggers) |  |  |  |  |
| CalatheaOrnata | 2.640 | 3.557±0.080 | 3.077±0.201 | 3.219±0.029 | 3.017±0.112 |
| 23_.\|/**9821;- - - | 2.160 | 4.389±0.110 | 4.332±0.056 | 4.313±0.071 | 4.054±0.082 |
| SpyL4bb | 1.480 | 4.127±0.124 | 3.915±0.090 | 3.952±0.094 | 3.684±0.166 |
| ILoveAppleJuice | 3.360 | 3.895±0.132 | 3.947±0.090 | 3.916±0.021 | 4.067±0.084 |
| GlobalWarmingIsReal! | 1.330 | 4.035±0.090 | 4.009±0.081 | 3.784±0.081 | 3.806±0.117 |

### 4.2 Backdoor Removal

Backdoors can have arbitrary triggers and responses, which makes it challenging to find and remove them using standard techniques [, , ].
Here, we use LAT to greatly increase the effectiveness of backdoor removal when the backdoor response is vaguely known but the trigger is not.

##### Models and data

We use the five backdoored LLMs from Rando et al. [] who implanted backdoors using RLHF [, , ] such that, upon encountering specific keyword triggers (see Table 3), the models would respond in a helpful and *harmful* way as opposed to a helpful and *harmless* one.
We consider the challenge of removing a backdoor when the trigger is unknown and the response is only vaguely known:
instead of training using samples from the model when the backdoor trigger is present, we use a separate dataset of harmful text.
We train all models using the ‘helpful’ and ‘harmless’ splits of the Anthropic’s HH-RLHF preference dataset [].

##### Methods

Using the above datasets, we fine-tune the models from Rando et al. [] using direct preference optimization (DPO) [] and DPO with LAT for 1024 steps on batches of size 16 (see Appendix A for further details).
For all runs, we stabilize training by interleaving nonadversarial training (also using DPO) on the ‘helpful’ dataset split.
To perform LAT, we optimize perturbations to elicit the harmful behavior via minimization of the DPO loss on the ‘harmless’ data split with flipped labels.
We attack hidden layers 4, 12, 20, and 28.
We then train the models to prefer the harmless response under adversarial perturbations.
We experiment with two training conditions.
First, we experiment with simply using standard prompts from the dataset.
Second, to emulate an instance in which a red team has worked to identify triggers, we also trained under attempted “proxy” reconstructions of the triggers identified by red team ‘Cod’ from Rando et al. [].

##### Evaluation

To evaluate the harmlessness of the model and its susceptibility to the backdoor, we used the reward model from Rando et al. [], which was trained to distinguish safe from unsafe responses.
As before, we also evaluate models under the MMLU benchmark [].

##### Targeted LAT greatly improves backdoor removal without side effects.

Evaluation results are in Table 3.
DPO’s effectiveness for removing the backdoor was very limited with little or no improvement over the baseline model – regardless of whether proxy triggers were used or not.
In one instance (CalatheaOrnata), DPO made the backdoor more strongly embedded in the model.
These failures echo prior findings from Hubinger et al. [], who showed that adversarial training often failed to remove a backdoored “sleeper agent.”
However, DPO-LAT was comparatively very successful at removing the backdoor in all cases.
Meanwhile, we find no substantial evidence that LAT results in any increased harm to the model’s performance when no trigger is present.
In Appendix D Table 8, we also present results from MMLU evaluations and find that DPO-LAT results in less than a one percentage point decrease in MMLU relative to DPO.

### 4.3 Machine Unlearning

Here, our goal is to augment methods for unlearning harmful or copyrighted knowledge from LLMs.
We first unlearn knowledge of Harry Potter (Section 4.3.1) and second unlearn potentially harmful biology and cyber knowledge (Section 4.3.2).

#### 4.3.1 Who’s Harry Potter?

Following work on unlearning knowledge of Harry Potter from Eldan and Russinovich [], we show that targeted LAT can improve the robustness of unlearning without sacrificing the model’s performance on other topics.

Table: Table 4: Targeted LAT improves Harry Potter unlearning. We evaluate Harry Potter unlearning using MMLU to test models’ general capabilities and the *familiarity* measure from Eldan and Russinovich [] to test their unlearning. We evaluate the robustness of unlearning with a “Basic” familiarity evaluation from Eldan and Russinovich [] plus the same evaluation performed after translating into “Spanish”, using “Jailbreak” prompts, including Harry Potter “Summary” prompts in context, and including Harry Potter “Text” samples in context. In the figure and table, we report the means $\pm$ the standard error of the mean.

##### Model and methods

We work with the “Who’s Harry Potter” (WHP) method from Eldan and Russinovich [].
It involves taking a corpus of text to forget (e.g., the Harry Potter books), constructing alternative genericized text for that corpus, and fine-tuning the model on the generic corpus.
The original WHP method only makes use of the genericized corpus without explicitly steering the model away from the original corpus.
Because our goal is to augment WHP with LAT, as a baseline, we use a modified version of WHP, which we call WHP-Contrastive (WHP-C).
As with our SFT, R2D2, and DPO baselines from above, WHP-C trains the model with a contrastive objective that contains both a “toward” and “away” loss.
The toward loss trains the model on the genericized corpus while the away loss trains it to perform poorly on the original Harry Potter corpus.
Also as before, we interleave supervised fine-tuning batches on the UltraChat dataset [] to stabilize training.
When performing WHP-C-LAT, we optimize the adversarial attacks to minimize the cross-entropy loss on the original Harry Potter text.
For all methods, we train on 100 batches of size 16 for 4 steps each.
Finally, in Appendix E, we also experiment with optimizing and constraining adversarial perturbations in a whitened space before de-whitening and adding them to the model’s latents.

##### Evaluation

To evaluate general performance, we again use MMLU [].
Next, we evaluate Harry Potter familiarity [] under Harry Potter knowledge extraction attacks.
Full details are available in Appendix F.
First, in response to past work suggesting that unlearning can fail to transfer cross-lingually [], we evaluate familiarity in Spanish.
Second, to test the robustness of unlearning to jailbreaks [], we evaluate familiarity under jailbreaking prompts [].
Third and fourth, we evaluate the extent to which the model is robust to knowledge extraction attacks [, , , , ] in the form of high-level summaries and short snippets of text from the Harry Potter books.

##### Targeted LAT helps to more robustly unlearn Harry Potter knowledge.

We present results in Table 4.
WHP-C-LAT Pareto dominates WHP and WHP-C across all measures except MMLU.

#### 4.3.2 Unlearning WMDP Biology and Cyber Knowledge

Following work from Li et al. [], who studied the unlearning of potentially dangerous biology and cyber knowledge, we show that targeted LAT can help to improve existing approaches for unlearning.

##### Data

As in as in Li et al. [], we use the WMDP biology and cyber corpora as forget datasests and WikiText [] as a retain dataset.

##### Model and methods

As in Li et al. [], we use Zephyr-7B off the shelf [].
We test two different unlearning methods with and without targeted LAT.
First, we use a shaped gradient ascent (GA) method inspired by [].
We fine-tune the model to jointly minimize training loss on the retain set and $\log(1-p)$ on the forget set as done in Mazeika et al. [].
To augment GA with targeted LAT, we apply latent-space perturbations optimized to minimize training loss on the forget set.
To stabilize training, we also interleave training batches with supervised finetuning on the Alpaca dataset [].
Second, we use representation misdirection for unlearning (RMU) from Li et al. [].
With RMU, the model is trained at a given layer to (1) map activations from forget-set prompts to a randomly sampled vector while (2) leaving activations from other prompts unaltered.
To augment RMU with targeted LAT, we apply latent-space adversarial perturbations only when training on the forget set.
We optimize these perturbations to minimize the model’s cross-entropy training loss on the undesirable forget-set example.
We experimented with various layer combinations and found the best results from applying them to the activations immediately preceding the RMU layer.

##### Evaluation

We evaluate how well the model’s general capabilities have been preserved by testing on MMLU [] and AGIEval [].
We evaluate the effectiveness of unlearning in the model using biology and cyber knowledge assessments from Li et al. [].
These multiple choice evaluations represent a qualitatively different task than the forget sets (which were full of bio and cyber documents), so they test the ability of LAT to generalize to qualitatively different kinds of unwanted behaviors than those used during fine-tuning.
To test the robustness of the unlearning, we also evaluate models under few-shot finetuning attacks in which an attacker seeks to extract knowledge by finetuning the model on a small number of examples [, , , , , , , ].
Here, we use a simple but surprisingly effective attack: we randomly sample a single batch of 2 examples from the relevant forget set and repeatedly train on that single batch for 20 iterations.
We then report the highest WMDP bio/cyber performances for each model across evaluation checkpoints at 5, 10, and 20 steps.
For all evaluations, we use 1,000 samples on lm-evaluation-harness v0.4.0 [] as done in Li et al. [].

Table: Table 5: Targeted LAT can improve gradient ascent (GA) and representation misdirection for unlearning (RMU)’s ability to unlearn the WMDP biology and cyber datasets [] with minimal side effects. We evaluate models’ general performance using MMLU and AGIEval and its unlearning with the WMDP bio and cyber evaluations from Li et al. []. The random-guess baseline for WMDP bio/cyber is 25%. Finally, to evaluate robustness to re-learning, we report WMDP performance after up to 20 iterations of repeatedly retraining on a single batch of 2 examples. In the figure and table, we report means and standard error of the means over $n=3$ runs with different random seeds.

##### Targeted LAT improves GA and RMU’s ability to robustly unlearn biology and cyber knowledge with minimal side effects.

Table 5 shows results for evaluating models by MMLU versus unlearning effectiveness.
GA-LAT outperforms GA by a large margin under all evaluations.
Similarly, RMU-LAT outperforms RMU in all evaluations, except for a 1.2% decrease in MMLU and 2.1% decrease in AGIEval.
Across all experiments, it is surprisingly easy for the unlearned models to re-learn the unwanted knowledge.
Repeatedly training on the same batch of 2 examples for up to 20 iterations improved WMDP bio/cyber performance by an average of 15.7 percentage points.
However, LAT makes the models more resistant to re-learning.
On average, re-learning closed 74.7% of the performance gap between the unlearned model and the original model for non-LAT methods but only 59.9% of the gap for LAT methods.

## 5 Discussion

##### Targeted LAT can effectively augment existing state-of-the-art fine-tuning and adversarial training methods.

By attacking the model’s latent representations, LAT offers a unique solution because models represent concepts at a higher level of abstraction in the latent space [].
Here, we have used targeted latent adversarial training (LAT) to strengthen existing defenses against persistent harmful behaviors in LLMs.
We have applied LAT to three current challenges with state-of-the-art LLMs: jailbreaking [], unlearning [], and backdoor removal [, ].
In each case, we have shown that LAT can augment existing techniques to improve the removal of unwanted behaviors with little or no tradeoff in general performance.
Overall, these results support but do not yet confirm our hypothesis that LAT can remove neural circuitry from models responsible for undesirable behaviors.
We leave analysis of the mechanisms behind harmful model behaviors (e.g., []) to future work.

##### Targeted LAT is a practically valuable tool to improve the safety and security of LLMs.

Our motivation for LAT is a response to two observations.
First, LLMs empirically can persistently retain harmful capabilities despite attempts to remove them with adversarial training [, , , , , , , , , , , , ].
Second, there have been empirical and theoretical findings that LLMs undergo limited changes to their inner capabilities during fine-tuning [, , , , , , , , , , ].
All three problems that we have used targeted LAT to address – jailbreaks, backdoors, and undesirable knowledge – are ones in which an LLM exhibits harmful behaviors that are difficult to thoroughly remove.
Our results show that targeted LAT can be useful for making models more robust to these persistent failures.
We also find that these failure modes need not be precisely known for LAT to be helpful, showing instances in which LAT can improve generalization to different datasets of attack targets, harmful behaviors, and knowledge-elicitation methods than were used during training.

##### LLM unlearning techniques are surprisingly brittle.

In Section 4.3, we find that state-of-the-art LLM unlearning methods are surprisingly vulnerable to relearning from small amounts of data.
We find that re-training repeatedly on only *two* samples from the forget set was consistently able to close more than half of the performance gap between the original and unlearned models on average.
We find that targeted LAT can reduce the sample efficiency of re-learning, but there is much room for improvement in designing unlearning methods that are robust to few-shot finetuning attacks.
We are interested in future work to explore LAT’s potential to improve on existing approaches for making models robust to few-shot fine-tuning attacks [, , , ].

##### Limitations – attack methodology and model scale.

While we have shown that LAT can be useful, it can also be challenging to configure and tune.
In our experience, we found the selection of dataset, layer(s), and perturbation size, to be influential.
We also found that interleaving supervised finetuning in with training and NaN handling were key to stable training.
LAT can be done in different layers, with various parameterizations, and under different constraints.
Our work here is limited to residual stream perturbations designed with projected gradient descent.
Additionally, all of our experiments are done in LLMs with fewer than 10 billion parameters.

##### Future work

- •
Improved latent-space attacks In addition to performing LAT with perturbations to an LLM’s residual stream, we are interested in other strategies for attacking its internal representations.
Toward this goal, engaging with recent work on LLM representation engineering [, ] and interpretability [] may help to better parameterize and shape latent space attacks.
We also speculate that universal attacks instead of single-instance attacks may be more interpretable and might better target the most prominent mechanisms that a model uses when it produces undesirable outputs.
- •
Augmenting other latent-space manipulation techniques Concurrently with our work, Zou et al. [], Rosati et al. [], and [] introduced other latent-space manipulation techniques for making LLMs robust to undesirable behaviors. We are interested in studying how these techniques compare to LAT and whether LAT can be used to improve them.
- •
Generalized adversarial attacks for LLM evaluations We are interested in the extent to which embedding-space attacks (e.g., []), latent-space attacks, (e.g., []), and few-shot fine-tuning attacks (e.g., []) can improve evaluations of LLM safety [].

## Broader Impacts

This work was motivated by the goal of training more safe and trustworthy AI systems.
We believe that LAT will be practically useful for training better models.
However, we emphasize that LAT is a value-neutral technique for training AI systems to align with their developer’s goals.
It is important not to conflate AI alignment with safety [].
We believe that this work will contribute to helpful progress, but we emphasize that many of the risks from AI systems come from misuse and adverse systemic effects as opposed to unintended hazards such as the ones we work to address.

## Contributions

Paper writing was performed by Abhay Sheshadri, Aidan Ewart, Phillip Guo, Aengus Lynch, Cindy Wu, and Stephen Casper. Experiments for Section 4.1 were led by Abhay Sheshadri, Aengus Lynch, and Vivek Hebbar. Experiments for Section 4.2 were led by Aidan Ewart. Experiments for Section 4.3.1 were led by Phillip Guo. Experiments for Section 4.3.2 were led by Cindy Wu and Philip Guo. Advising was provided by Stephen Casper, Dylan Hadfield-Menell, Asa Cooper Stickland, and Ethan Perez. Project management was provided by Henry Sleight.

## Acknowledgements

We are thankful to Rajashree Agarwal, Rohit Gandikota, John Hughes, Erik Jenner, Alex Lyzhov, Sam Marks, Jacob Pfau, Sara Price, Javier Rando, Markian Rybchuk, Lennart Schulze, Aaquib Syed, Alex Turner, and Tony Wang for useful conversations.
We thank William Brewer, Rocket Drew, Ronny Fernandez, McKenna Fitzgerald, Juan Gil, Carson Jones, Ryan Kidd, Christian Smith, and Laura Vaughan, for program support.
This project was funded in part by a grant from Open Philanthropy and used compute provided by the Center for AI Safety.
This support was offered by both organizations without conditions, and neither exerted any influence over the paper’s content.

## Appendix A Loss Functions for LAT

### A.1 RT-LAT

Here, we describe the RT-LAT method described in Section 4.1 in greater detail. We assume we are given two datasets - a dataset of harmful requests and pairs of preferred and rejected completions $\mathcal{D}_{p}=\{(x_{i},c_{i},r_{i})\}$, and a generic dataset of benign requests and helpful completions $\mathcal{D}_{b}=\{(x_{i},y_{i})\}$. For each batch, we train the adversarial attack $\delta$ to minimize $\mathcal{L}_{\text{attack}}$:

$$ $\mathcal{L}_{\text{attack}}=\underbrace{-\log P(r_{i}|g_{\theta}(f_{\theta}(x_ {i})+\delta_{i}))}_{\text{Movetowardsharmfulcompletions}}+\underbrace{-\log(1- P(c_{i}|g_{\theta}(f_{\theta}(x_{i})+\delta_{i})))}_{\text{ Moveawayfromharmlesscompletions}}$ (3) $$

We additionally add the constraint that $||\delta_{i}||_{2}\leq\epsilon$, where $\epsilon$ is a hyperparameter, to restrict the adversary’s power. We then train the model parameters $\theta$ against these adversarial attacks by minimizing $\mathcal{L}_{\text{model}}$. We define $\mathcal{L}_{\text{model}}$ in terms of the loss functions $\mathcal{L}_{\text{defense}}$ and $\mathcal{L}_{\text{benign}}$:

$$ $\mathcal{L}_{\text{defense}}=\sum_{(x_{i},c_{i},r_{i})\sim\mathcal{D}_{p}}{ \underbrace{-\log P(c_{i}|g_{\theta}(f_{\theta}(x_{i})+\delta_{i}))}_{\text{ Movetowardsharmlesscompletions}}+\underbrace{-\log(1-P(r_{i}|g_{\theta}(f_{ \theta}(x_{i})+\delta_{i})))}_{\text{Moveawayfromharmfulcompletions}}}$ (4) $$

$$ $\mathcal{L}_{\text{model}}=\mathcal{L}_{\text{defense}}+\mathcal{L}_{\text{ benign}}$ (5) $$

We can use one of two different benign loss terms:

$$ $\mathcal{L}_{\text{benign},\text{SFT}}=\sum_{(x_{i},y_{i})\sim\mathcal{D}_{b}} -\log P(y_{i}|g_{\theta}(f_{\theta}(x_{i})))$ (6) $$

$$ $\mathcal{L}_{\text{benign},\text{KL}}=\sum_{(x_{i},y_{i})\sim\mathcal{D}_{b}} \text{KL}\left[P(y_{i}|g_{\theta^{*}}(f_{\theta^{*}}(x_{i})))\,\|\,P(y_{i}|g_{ \theta}(f_{\theta}(x_{i})))\right]$ (7) $$

where $\theta^{*}$ are the weights of the frozen reference model.
Note that $\mathcal{L}_{\text{benign}}$ is always calculated on inputs where no adversarial attack is present.

We use $\mathcal{L}_{\text{benign},\text{SFT}}$ for our Llama2 results, and $\mathcal{L}_{\text{benign},\text{KL}}$ for our Llama3 experiments.
$\mathcal{L}_{\text{benign},\text{SFT}}$ trains the model to maximize the probability of the ground-truth completions for benign prompts, whereas $\mathcal{L}_{\text{benign},\text{KL}}$ trains the model to preserve its original logits over possible completions for benign prompts.
We hypothesize that $\mathcal{L}_{\text{benign},\text{KL}}$ might preserve original model capabilities better when the quality of $\mathcal{D}_{b}$ is poor relative to the model being trained.
Empirically, we find that $\mathcal{L}_{\text{benign},\text{KL}}$ can better allow more capable models to retain their capabilities during adversarial training.

### A.2 DPO-LAT

We now describe the DPO-LAT loss inspired by Rafailov et al. []. Similarly to RT-LAT, we assume that we have a paired preference dataset of harmless/harmful completions $\mathcal{D}_{p}=\{(x_{i},c_{i},r_{i})\}$, where $c_{i}$ is the harmless result and $r_{i}$ is the harmful response. Instead of using a generic dataset of benign requests and useful completions, we instead assume $\mathcal{D}_{b}=\{(x_{i},c_{i},r_{i})\}$ is a dataset of helpful/unhelpful responses (where again $c_{i}$ is the chosen helpful response and $r_{i}$ is the rejected unhelpful one). We take $\mathcal{D}_{p}$ from the ‘harmless’ split of Anthropic’s HH-RLHF dataset [] and $\mathcal{D}_{b}$ from the ‘helpful’ split.

We choose $\mathcal{L}_{\text{attack}}$ to cause the model to prefer the harmful response $r_{i}$ over $c_{i}$ where $(x_{i},c_{i},r_{i})\sim\mathcal{D}_{p}$, using the DPO loss (where $\theta^{*}$ are the weights of the frozen reference model):

$$ $\mathcal{L}_{\text{attack}}=-\log\sigma\left(\underbrace{\beta\log\frac{P(r_{i }|g_{\theta}(f_{\theta}(x_{i})+\delta_{i}))}{P(r_{i}|g_{\theta^{*}}(f_{\theta^ {*}}(x_{i})))}}_{\text{Movetowardsharmfulcompletions}}-\underbrace{\beta\log \frac{P(c_{i}|g_{\theta}(f_{\theta}(x_{i})+\delta_{i}))}{P(c_{i}|g_{\theta^{*} }(f_{\theta^{*}}(x_{i})))}}_{\text{Moveawayfromharmlesscompletions}}\right)$ (8) $$

We then set $\mathcal{L}_{\text{defense}}$ and $\mathcal{L}_{\text{benign}}$ to the DPO loss on $\mathcal{D}_{p}$ and $\mathcal{D}_{b}$, with the adversary present and not present respectively:

$$ $\mathcal{L}_{\text{defense}}=-\sum_{(x_{i},c_{i},r_{i})\sim\mathcal{D}_{p}} \log\sigma\left(\underbrace{\beta\log\frac{P(c_{i}|g_{\theta}(f_{\theta}(x_{i} )+\delta_{i}))}{P(c_{i}|g_{\theta^{*}}(f_{\theta^{*}}(x_{i})))}}_{\text{ Movetowardsharmlesscompletions}}-\underbrace{\beta\log\frac{P(r_{i}|g_{\theta} (f_{\theta}(x_{i})+\delta_{i}))}{P(r_{i}|g_{\theta^{*}}(f_{\theta^{*}}(x_{i})) )}}_{\text{Moveawayfromharmfulcompletions}}\right)$ (9) $$

$$ $\mathcal{L}_{\text{benign}}=-\sum_{(x_{i},c_{i},r_{i})\sim\mathcal{D}_{b}}\log \sigma\left(\beta\log\frac{P(c_{i}|g_{\theta}(f_{\theta}(x_{i})))}{P(c_{i}|g_{ \theta^{*}}(f_{\theta^{*}}(x_{i})))}-\beta\log\frac{P(r_{i}|g_{\theta}(f_{ \theta}(x_{i})))}{P(r_{i}|g_{\theta^{*}}(f_{\theta^{*}}(x_{i})))}\right)$ (10) $$

### A.3 WHP-C-LAT and GA-LAT

The WHP-C-LAT and GA-LAT methods described in Section 4.3.1 and Section 4.3.2 use a toward-only adversary which optimizes for next-token cross-entropy loss on Harry Potter and the WMDP forget corpora respectively.
For WHP, the model is trained as in Eldan and Russinovich [].
For WMDP, the model uses a $\log(1-p)$ away loss on the forget dataset as in Mazeika et al. [].
In both cases, we additionally include a toward loss on WikiText [] to match Li et al. [], and a supervised fine-tuning (SFT) loss on Alpaca []. While calculating the model’s toward and away losses, we keep the perturbations from the adversary. We remove these perturbations for SFT.

Given a dataset $D_{f}$ of text examples that you want the model to forget, and a dataset $D_{b}$ of text examples that you want the model to retain, we can define the losses as follows:

$$ $\mathcal{L}_{\text{attack}}=-\sum_{t_{i}\in D_{f}}\sum_{j}\log P(t_{i,j}|g(f(t _{i,<j})+\delta_{i}))$ (11) $$

$$ $\mathcal{L}_{\text{forget}}=-\sum_{t_{i}\in D_{f}}\sum_{j}\log(1-P(t_{i,j}|g(f (t_{i,<j})+\delta_{i})))$ (12) $$

$$ $\mathcal{L}_{\text{retain}}=-\sum_{t_{i}\in D_{b}}\sum_{j}\log(t_{i,j}|g(f(t_{ i,<j})))$ (13) $$

$$ $\mathcal{L}_{\text{model}}=\mathcal{L}_{\text{forget}}+\mathcal{L}_{\text{ retain}}$ (14) $$

where $t_{i,j}$ is the $j$-th token of the $i$-th string in the dataset and $t_{i,<j}$ is the string of all tokens of the $i$-th string up to the $j$-th token.

### A.4 RMU-LAT

Here, we use the same RMU loss as used in Li et al. []. The adversary still optimizes for next-token cross-entropy loss on the WMDP forget corpora. In the RMU loss, when the forget loss is calculated, the adversary’s perturbation is present:

$$ $\aligned\mathcal{L}_{\text{defense}}=&\frac{1}{L}\sum_{\text{token}t\in x_{ \text}{forget}}||M_{\text}{updated}(t)+\delta_{i}-c\cdot\mathbf{u}||_{2}^{2}\\ +\alpha\cdot\frac{1}{L}\sum_{\text{token}t\in x_{\text}{retain}}||M_{\text}{ updated}(t)-M_{\text}{frozen}(t)||_{2}^{2}$ (15) $$

where $L$ is the length of the input tokens, and u is a randomly chosen vector from a uniform distribution between $[0,1]$ that is then normalized (and stays constant throughout training). The constants $c$ and $\alpha$ are hyperparameter coefficients, which we set to be 6.5 and 1200 as in Li et al. [] for Zephyr-7B.

## Appendix B Jailbreaking Robustness Under Untargeted LAT

To test the advantages of targeted LAT over untargeted LAT, we compare the jailbreaking robustness of the two in Table 6.
Here, during untargeted LAT, the adversary does not work to make the model comply with the jailbreak.
Instead, it only works to make the model fail to output a refusal.
We find that untargeted LAT results in less harm to general performance compared to targeted LAT but not refusal training.
Meanwhile, untargeted lat results in comparable or slightly worse robustness in most cases compared to targeted LAT.
However, for prefill and GCG attacks, untargeted LAT fares much worse than targeted LAT.

Table: Table 6: Untargeted LAT results in less jailbreak robustness than targeted LAT. Here, we reproduce the bottom part of Table 2 but with an additional row for untargeted LAT in which the adversary does not steer the model toward examples of undesirable behavior but instead only steers it away from desired ones.

## Appendix C Jailbreaking Robustness Under an Alternate Autograder

In Section 4.1, we evaluate jailbreak success using the StrongReject autograder [].
However, here we also report results using the HarmBench autograder [].
Overall, we find that the HarmBench autograder is significantly more likely to label attacks as successful, but the overall trends within results remain similar.

Table: Table 7: Jailbreaking results using the HarmBench autograder. Here, we reproduce table 2 except we report results for attacks according to the HarmBench [] autograder instead of the StrongReject [] autograder which was used in table 2. Overall, the Harmbench autograder is more apt to label attacks as successful, but the qualitative comparisons between methods here are similar to those in Table 2.

**Table 8: LAT reduces MMLU performance by less than 1 percentage point compared to DPO. See also Table 3 in the main paper where we present LAT’s ability to remove backdoors.**
| Backdoor | Baseline | DPO | DPO-LAT | DPO | DPO-LAT |
| --- | --- | --- | --- | --- | --- |
| (proxy triggers) | (proxy triggers) |  |  |  |  |
| CalatheaOrnata | 0.464 | 0.465 | 0.458 | 0.465 | 0.458 |
| 23_.\|/**9821;- - - | 0.464 | 0.466 | 0.458 | 0.466 | 0.456 |
| SpyL4bb | 0.464 | 0.465 | 0.457 | 0.464 | 0.456 |
| ILoveAppleJuice | 0.464 | 0.465 | 0.458 | 0.464 | 0.456 |
| GlobalWarmingIsReal! | 0.464 | 0.465 | 0.460 | 0.464 | 0.441 |

## Appendix D Backdoored Model MMLU Performance

To evaluate the destructiveness of DPO-LAT versus DPO on backdoor removal, we evaluate each model’s performance on MMLU [].
We present our results in Table 8 for a single model.
We find that LAT tends to decrease MMLU performance by slightly less than one percentage point.

## Appendix E Low Rank Adapters and Scaled Perturbation Constraints for WHP Unlearning

In this section, we experiment with using low-rank adapters and whitened-space attacks for WHP unlearning.
Typically, adversarial training methods that use projected gradient descent constrain perturbations to be within an $L_{p}$-norm spherical ball [].
However, for latent-space perturbations, this approach is arguably unnatural because in the latent-space, activations vary more along some directions than others.
To address this, here, we test a scaling method to constrain attacks in a way that better respects the shape of the activation manifold in latent space in Section 4.3.1.
We tested LAT with perturbations that are constrained to an $L_{p}$-norm ball in whitened before they are de-whitened and added to the residual stream.

Our goal was to increase the ability of targeted LAT to operate on coherent features relating to the unlearning corpora (specifically, features that would preserve meaning but cause the model to no longer recognize the text as related).
As a result, we perform principal component analysis (PCA) on the distribution of activations between Harry Potter text and the coherent genericized versions of the text produced during WHP.
We optimize and constrain the perturbations in a whitened space before de-whitening them using the inverse PCA transformation matrix and then applying it to the model’s latent states.
In addition, we use a low-rank adapter on all linear modules of rank 64.
In our experiments, this resulted in weaker unlearning for WHP experiments but with less of a tradeoff in general capabilities.
The results are shown in Table 9.
However, we speculate that unlearning tasks may be especially well-suited to this type of scaling, and we leave deeper investigation to future work.

Table: Table 9: Training with scaling results in less strong Harry Potter unlearning but better tradeoffs in general performance. Compare to Table 4 in the main paper.

## Appendix F Tests for Robust and Competitive Unlearning in LLMs

Eldan and Russinovich [] fine-tune Llama-2-7B-Chat [] (Llama-2) to unlearn knowledge of the Harry Potter universe.
Their method is based on fine-tuning using text that has been modified to replace domain-specific content with generic content.
Throughout experiments here, we compare the WHP model from Eldan and Russinovich [], our replications, and our replication with targeted LAT (see Section 4.3.1).

Here, we outline the methods we use to evaluate unlearning in Section 4.3.1

##### Familiarity

To evaluate the model, Eldan and Russinovich [] introduce “Familiarity” as a metric which measures the extent of Harry Potter content contained in the model’s completions of Harry Potter-related sequences as determined by an automated GPT-4 evaluation. To measure Familiarity, we follow the same method from Eldan and Russinovich [] to evaluate a completion from the model. An evaluation prompt is formatted with the datapoint reference, prompt, and model completion, passed into GPT-4, then obtain a model Familiarity score (Figure 2), using “gpt-4-turbo-preview” at seed=42 and temperature=0, with max tokens=252. All model completions are scored in this way, and then we calculate the Familiarity metric starting a counter at 0, adding 1 for grade 3 completions, 0.2 for grade 2 completions, and 0 otherwise. Then, this total is divided by the total number of completions.

Figure: Figure 2: Familiarity evaluation system prompt from Eldan and Russinovich []: GPT-4 generates a reasoning sequence, before writing "MODEL FAMILIARITY: X/3", from which we extract the score. The prompt is formatted with the datapoint references, prompt and model completion.

Aside from standard Familiarity evaluations as done in Eldan and Russinovich [], we also perform four other evaluations using Familiarity, but when the model is evaluated under prompt extraction attacks.

##### Spanish

LLM fine-tuning does not always transfer to other languages [, ], so we test the models’ Harry Potter Familiarity with the prompts translated by GPT-4 [] into Spanish.

##### Jailbreak Prompts

Simple jailbreaks have been successful at resurfacing knowledge that is typically not produced by LLMs (e.g., building a bomb).
We test a jailbreaking prompt designed to resurface Harry Potter knowledge based on prior successful jailbreaks against Llama-2 models [] (Figure 3).

Figure: Figure 3: Jailbreaking Prompt: A prompt designed to pressure the model to resurface Harry Potter knowledge.

##### Summary and Snippet Prompts

Here, we use few-shot and summary prompting.
We provide the model with small amounts of general context related to Harry Potter with the goal of resurfacing existing suppressed knowledge that was not provided.
We evaluate Familiarity when either a high-level summary (Figure 4) or the first 10 lines of Book 1 are included in context.

Figure: Figure 4: Long summary: 3-paragraph long summary of Harry Potter, generated by GPT-4. We use this for in-context relearning experiments in 4.3.1.

## Appendix G WMDP Unlearning Details

##### Trainable layers and parameters

We use LoRA [] with rank 64 for GA and GA-LAT.
For RMU and RMU-LAT, we do not use LoRA and instead train the MLP weights full-rank, as in Li et al. [].

##### PGD/RMU layers

There are three layer choices that can be varied in our setup: which layer(s) of the model to put the adversary, which layers to train for RMU, and which layer to do the RMU MSE activation matching over. We kept to the same layers (trainable and RMU matching) for RMU as in Li et al. [] – the RMU layer $\ell$ for the activation matching, with $\ell,\ell-1,\ell-2$ trainable to keep the set of hyperparameters to search over reasonably small. Applying attacks to layer $\ell-2$ requires a smaller $\epsilon$ ball radius for our random perturbations; else, we found that the adversary prevents the model trained with RMU from successfully unlearning. We also find the greatest benefit in applying attacks to the layer before the RMU activation matching layer.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: General suggestions:
- Table 2 is too wide. Also, the figures above Table 2 should have a separate caption. Also, all table captions should be above tables, not below. Also \citet vs. \citep should be used correctly throughout the paper (e.g., double check the “Future work” paragraph).
- The Future Work paragraph: an introductory sentence would improve the reading flow.
- “Direct preference optimization: Your language model is secretly a reward model. Advances
in Neural Information Processing Systems, 36, 2024.” - should be 2023, not 2024.
- “prefilling attacks (Haizelabs)” doesn’t seem to be the right reference for the prefilling attack, since it wasn’t introduced there.
- In addition to MMLU, it would be also good to add the MT-Bench score for the DPO-LAT models in Table 8.


**Update after the rebuttal: The paper still feels quite rushed to me, primarily in terms of the experiments. Overall, I feel like Targeted LAT might be a promising method to improve multiple dimensions of safety. However, this is still unclear from this version of the paper. I will keep my original score 5/10. I think the approach might be promising but requires much more systematic experiments. I think extending this work and resubmitting it would be the best option.**
soundness: 2
strengths: Pros:
- I believe targeted LAT *can be* a useful attack-agnostic defense, although the current evaluation lacks depth (see below).
- The breadth of evaluation is appealing. It’s nice to see a method that potentially improves on safety/alignment across multiple diverse tasks.
summary: This paper introduces targeted latent adversarial training (LAT) as a technique to improve robustness to persistent harmful behaviors in large language models (LLMs). The authors demonstrate LAT's effectiveness in three key applications: (1) improving resistance to jailbreaking attacks while maintaining model performance, (2) removing backdoors without knowledge of the trigger, and (3) enhancing unlearning of undesirable knowledge. The core idea is to perturb latent activations to elicit specific undesirable behaviors during training, then optimize the model to be robust against such perturbations. The authors show LAT can augment existing techniques like refusal training, DPO, and machine unlearning methods to achieve better results with minimal computational overhead.
weaknesses: Weaknesses:
- The attacks used for the evaluation in the main table (Table 2) are quite weak: the best attack success rate is 27.7% on Llama-3-8B Instruct, although it’s possible to achieve ~100% ASR on this model (e.g., as reported in [Jailbreaking Leading Safety-Aligned LLMs with Simple Adaptive Attacks](https://arxiv.org/abs/2404.02151) but with a different judge). Without strong enough attacks, it’s hard to conclude that the defense is effective enough, especially given the anecdotal evidence that there are some simple breaks like the one mentioned in the paragraph “Manual red-teaming and research demo”.
- It’s not clear to me why the proposed targeted formulation should be better than the existing LAT methods, such as Embedding-Space Adversarial Training (Xhonneux et al., NeurIPS 2024) or Defending against unforeseen failure modes with latent adversarial training (Casper et al., 2024b). There are some explanations in the introduction but they seem quite handwavy. Also, the only comparison between RT-EAT and RT-EAT-LAT suggests a small difference: 4.3% vs. 2.9% prefilling ASR while 6.22 vs. 5.86 MT-Bench score - so it’s not even clear which model is really better.
- MMLU and MT-Bench may be too easy as an over-refusal evaluation since those questions are completely harmless. Adding something like [XS-Test](https://arxiv.org/abs/2308.01263) or [OR-Bench](https://arxiv.org/abs/2405.20947) would make the evaluation stronger.
- R2D2 and RT-EAT should also be added for Llama-3 as baselines.
- Since there are no other baselines except DPO for backdoor removal included in Table 3, it’s unclear whether LAT is really necessary there or basically any algorithm that would *somehow* perturb the weights in the optimization process would work as well. 
- For the unlearning part, it’s not clear to me whether WHP-C-LAT pushes the Pareto frontier compared to WHP-C. WHP-C-LAT has a noticeably worse MMLU score (43.9% vs. 45.6%) although with a better unlearning performance. Also, the unlearning part should have more baselines (there are plenty of unlearning methods that exist in the literature).
- The choice of the L2 norm for layerwise perturbations looks a bit arbitrary. It would be nice to elaborate why it can make sense.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. The attack success rate for GCG's Llama-2 and Llama-3 is relatively low compared to the original paper; could you explain this? 
2. The Llama is famous for its safety. Could you provide a discussion or experiment on a model like Vicuna (easier to jailbreak) to see further performance?
3. The DPO setting with the backdoor trigger is impractical; could you discuss its real-world application more?
soundness: 2
strengths: - The paper introduces targeted Latent Adversarial Training (LAT). This computationally efficient approach enhances the robustness of LLMs by specifically targeting latent activations.
- Extensive experiments have been conducted to provide a good insight into the components of the proposed method.
- The paper is generally well-written. With clear illustrations and tables.
summary: Large language models (LLMs) often exhibit undesirable behaviors despite fine-tuning efforts to remove them. This paper addresses this issue using targeted Latent Adversarial Training (LAT), which enhances robustness by leveraging latent-space perturbations to target specific failure modes. The approach contrasts with traditional adversarial training, focusing on hidden activations rather than inputs. Targeted LAT improving resistance to jailbreaks, removing backdoors, and unlearning undesirable tasks with little computational cost. Extensive experiments validate the method's efficacy, showcasing its potential as a robust tool for mitigating harmful behaviors in LLMs.
weaknesses: - This paper follows a general adversarial training pipeline, which requires maximizing the adversarial loss while minimizing the "safety loss." The framework itself is familiar for adversarial training, which might hinder the contribution of the paper.
- As the proposed method shares similarities to the latent adversarial training (LAT), the paper needs to discuss the difference between the proposed method and the previous LAT. In addition, as the LAT perturbed the layer's activation,  choosing which layers to perturb needs to be better discussed and empirically verified.
- Despite its effectiveness in defense of jailbreak, the DPO setting with the backdoor trigger is impractical, as most training datasets are carefully constructed.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 1
questions: Please see the weakness part.
soundness: 3
strengths: (1) The t-LAT algorithm seems effective across a wide range of tasks and is flexible enough to be combined with many optimization objectives without adding much overhead.

(2) The authors provide necessary implementation guidelines such as adding additional SFT loss or KL divergence.
summary: This paper proposes the Targeted Latent Adversarial Training (t-TLA) technique that can be added on top of existing algorithms to (1) safeguard LLMs from jailbreak attacks (2) erase backdoor behaviors from LLMs (3) remove knowledge from the models. The authors conducted experiments on each task and observed promising performance.
weaknesses: **Major**

(1) Section 4.1: The attacks considered are not strong enough with most of them achieving ASR < 20% against the base model, making it questionable whether the proposed technique will bring improvement when faced with more advanced jailbreak attacks like [1], [2] and [3]. Also, both Llama-2 and Llama-3 are very safe models. I think the authors should experiment with weaker models like Vicuna-7B. An improvement of 2% in ASR is still somewhat marginal for me. (Also, I encourage the authors to experiment with larger models if computational resources permit.)

(2) Section 4.2: It is good to see that DPO-LAT surpasses DPO, but how does it perform when compared with other algorithms designed to remove backdoors from LLMs? 

(3) Section 4.3: The improvement in WHP dataset is too little to be noticed and only one algorithm and one model are considered. What is the key difference between section 4.3.1 and section 4.3.2. It is not clear to me why they should be separated. 

(4) Some important experimental details are left out, especially those related to the hyperparameter of the proposed algorithm and the baseline algorithms (i.e. the $\beta$ for DPO, GCG steps, the examples for MSJ, etc.). It is necessary for me to specify the details of the experiments to make the comparison fair and the experiments reproducible/reliable. 

(5) There is no ablation study about the choice of $epsilon$, updated layers, additional SFT loss, and etc, The choice of constraint budget and the additional SFT loss is not consistent across different sections.

[1] JAILBREAKING LEADING SAFETY-ALIGNED LLMS WITH SIMPLE ADAPTIVE ATTACKS; https://arxiv.org/pdf/2404.02151

[2]  Improved Techniques for Optimization-Based Jailbreaking on Large Language Models; https://arxiv.org/pdf/2404.02151

[3]  Improved Few-Shot Jailbreaking Can Circumvent Aligned Language Models and Their Defenses; https://arxiv.org/html/2406.01288v1

----

**Minor**

(1) It occurred to me occasionally while reading that the paper was written in an extreme rush. There are typo errors (line 220), tables exceeding the width limit (line 233-243, line 1188-1208), one line of equation occupying a full page (line 1107), figures without a caption (217-232), and broken citation links. All these errors can be spotted with a 10-min proof-reading.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
questions: * The RT-EAT method of Xhonneux et al 2024 appears to be equivalent to the proposed method if adversarial training is conducted only in the first latent layer of a model. Could the authors comment on that? If this is true, the connection should be highlighted to provide better context on how these different methods relate. 
* Can the authors explain how hyperparameters were optimized for the different methods and a table of the final hyperparameters used in the experiments
* Can the authors provide an argument for the sufficiency of the used utility benchmarks in Table 2 (or new results) 

Without further changes, I would recommend to reject this paper. However, I believe many of my concerns can be addressed in a rebuttal, and I am willing to change my score to accept.

Since there are still trivial errors in the paper (such as those remarked by myself or reviewer EUWJ) I decided to reduce my score. In its current state, the paper should not be accepted. 
All in all I believe this could be a valuable contribution and I strongly encourage the authors to submit a revised manuscript. For me, a score of 8 would have been achievable with a convincing rebuttal.
soundness: 3
strengths: * The authors provide code, models, and even user-friendly tools to evaluate their models (after submission). Given the long history of ineffective defenses, this is an important part of a defense contribution
* To the best of my knowledge the results on sleeper agents and unlearning are novel and demonstrate a broad applicability of adversarial  training to different security issues in LLMs. 
* The authors ablate performing adversarial training in different latent layers of a network, which seems to improve robustness
summary: The authors apply latent space adversarial training to three different security problems in large language modeling. They show that adversarial training in deeper layers of the network can additionally improve robustness. Further, they demonstrate that beyond jailbreaks, adversarial training can improve robustness against backdoor attacks and robustness against information leaks in the context of unlearning.
weaknesses: * The framing of the paper could be improved, in my opinion, but I am open to discussions. The authors highlight efficiency improvements upon prior work that explores discrete adv. training (e.g., 281). However, alternative methods exist that are much closer to the algorithm proposed here and are also efficient (Xhonneux et al., 2024, Yu 2024). I believe the authors should highlight the differences to more closely related prior works and focus the discussions on these differences. As far as I can see this includes: 1) New threat models, 2) Exploring different latent layers. This also includes the introduction, which should highlight unique limitations resolved in this paper (and not those already addressed by other works). Note that I do not consider Yu 2024 in terms of my rating as it was released shortly before the submission deadline. 
* The utility datasets used to evaluate model capabilities are insufficient. Both MMLU and MT-bench suffer from assigning high scores to models that refuse every request. The compliance dataset gives a high score to R2D2, which is known for over-refusal, which makes me skeptical about the result. I would recommend OR-bench to evaluate if latent adv. training has a negative impact on over-refusal (Cui et al., 2024). 
* Comparisons between papers would be easier if the authors used the original method name provided in the respective paper (i.e., RT-EAT vs CAT) 
* Table 2 should be fixed for the camera ready. 
* I found 6 occurrences of missing \ref{} and \cite{} commands: 1009, 1286, 1295, 1313, 1336, 1378
* I was not able to find any concrete hyperparameters for any method except for RT-EAT-LAT. A direct comparison between two methods without stating the hyperparameter search procedure appears to be insufficient. It's unclear if the benefit from RT-EAT-LAT comes from the choice to train in deep latent layers or from better hyperparameter tuning.

Meta-review: This paper introduces targeted latent adversarial training (LAT) as a technique to improve robustness to persistent harmful behaviors in large language models (LLMs). The authors demonstrate LAT's effectiveness in three key applications: (1) improving resistance to jailbreaking attacks while maintaining model performance, (2) removing backdoors without knowledge of the trigger, and (3) enhancing unlearning of undesirable knowledge. The core idea is to perturb latent activations to elicit specific undesirable behaviors during training, then optimize the model to be robust against such perturbations. The authors show LAT can augment existing techniques like refusal training, DPO, and machine unlearning methods to achieve better results with minimal computational overhead. For the strengths, the paper introduces targeted Latent Adversarial Training (LAT). This computationally efficient approach enhances the robustness of LLMs by specifically targeting latent activations. Extensive experiments have been conducted to provide a good insight into the components of the proposed method. The paper is generally well-written with clear illustrations and tables. However, there are several points to be further improved. For example, this paper follows a general adversarial training pipeline, which requires maximizing the adversarial loss while minimizing the "safety loss." The framework itself is familiar for adversarial training, which might hinder the contribution of the paper. As the proposed method shares similarities to the latent adversarial training (LAT), the paper needs to discuss the difference between the proposed method and the previous LAT. In addition, as the LAT perturbed the layer's activation, choosing which layers to perturb needs to be better discussed and empirically verified. Despite its effectiveness in defense of jailbreak, the DPO setting with the backdoor trigger is impractical, as most training datasets are carefully constructed. Moreover, some important experimental details are left out, especially those related to the hyperparameter of the proposed algorithm and the baseline algorithms (i.e. the for DPO, GCG steps, the examples for MSJ, etc.). It is necessary for me to specify the details of the experiments to make the comparison fair and the experiments reproducible/reliable. There is no ablation study about the choice of epsilon, updated layers, additional SFT loss, and etc, The choice of constraint budget and the additional SFT loss is not consistent across different sections. Therefore, this paper cannot be accepted at ICLR this time, but the enhanced version is highly encouraged to submit other top-tier venues.

Decision: Reject