Latent Adversarial Training Improves Robustness to Persistent Harmful Behaviors in LLMs

Authors: Abhay Sheshadri, Aidan Ewart, Phillip Huang Guo, Aengus Lynch, Cindy Wu, Vivek Hebbar, Henry Sleight, Asa Cooper Stickland, Ethan Perez, Dylan Hadfield-Menell, Stephen Casper
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: adversarial attacks, adversarial training, jailbreaks, trojans, backdoors, unlearning, robustenss

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.

Decision: reject
Reviewer scores: 5, 5, 3, 6

Reviews and discussion:
<thread>
creation date:
1729843847029

thread_type:
official_review

<thread_object>
creator: reviewer_rnfj

<content>
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.

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

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730466235604

thread_type:
official_review

<thread_object>
creator: reviewer_euwj

<content>
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.

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.

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.

questions:
Please see the weakness part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730525024920

thread_type:
official_review

<thread_object>
creator: reviewer_goqz

<content>
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.

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.

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.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730908572995

thread_type:
official_review

<thread_object>
creator: reviewer_s8mn

<content>
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.

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.

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.

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.**
</content>
</thread_object>

</thread>
<thread>
creation date:
1732742891296

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Closing Argument from Authors (part I)

comment:
# Closing Argument

One last thanks to the reviewers and AC. **In this final comment, we are going to be raidcally honest and say exactly what we think.** There are two reasons for this:
1. Our submission window closes after ICLR. We arXived on July 22. The models that we experimented on are becoming older, and more related work is constantly popping up. Our focus is on our followup work. We think that going back and updating experiments for the sake of an ICML submission would be playing the resubmission game rather than pursuing research impact.
2. The reviews are 3-5-5-6 which is likely a rejection. But it is well known that conference review is somewhat of a dice roll. Another challenge with reviews is that papers are sometimes judged less by the *sum of contributions* and more by the *number* of critiques reviewers have.

## On the paper’s contributions

**In their final reasons for decision, three reviewers expressed concern that the paper is lacking experimental breadth or depth.** [s8MN pointed to how we mostly worked with Llama2 and had limited experiments with Llama3 models](https://openreview.net/forum?id=wI5uHZLeCZ&noteId=z0nRzywPZT). [goqz asked about adding experiments with another model like Vicuna](https://openreview.net/forum?id=wI5uHZLeCZ&noteId=EroXUPJklP). And [EUWJ asked for more jailbreaking evaluations beyond the 6 that we used](https://openreview.net/forum?id=wI5uHZLeCZ&noteId=z7KSgsZ6QW). In our replies, we have responded to the specifics of these points, but we also want to make a broader point. **Evaluations can always be more elaborate. But we think that the paper's experimental breadth and depth is a strength, not a weakness.** We make contributions on jailbreak robustness, backdoor removal, and unlearning. To the best of our knowledge, no other related papers have done this. There is only one of which we know ([Tamirisa et al., 2024](https://arxiv.org/abs/2408.00761)) that works in *two* of these domains. Meanwhile, we work with 3 different LLMs, 5 different datasets, 13 training algorithms, and 21 evaluation measures for performance/robustness/efficiency. We set out to show that LAT is a versatile tool for improving LLM safeguards, and we think our experiments made a strong case for this. **We expect many papers will be accepted to ICLR which have substantially narrower contributions and less thorough experiments.**

## On the paper’s impact

**The paper, code, and models have been influential:** In direct response to reviewer concerns about reproducibility and model usefulness, we discussed the code and models we released with the paper:

> We have also released code, 15 models, and 2 datasets alongside the paper’s public arXival. As we write this, we have 9 forks and 26 stars on GitHub. We also have over 700 model downloads and 1000 dataset downloads on HuggingFace suggesting a practical contribution to the community. 

We think this reflects quite positively on the paper's quality and value.

## Asking for AC adjudication

In their reasons for decision, [EUWJ](https://openreview.net/forum?id=wI5uHZLeCZ&noteId=z7KSgsZ6QW) and [rnfj](https://openreview.net/forum?id=wI5uHZLeCZ&noteId=tPxm0G1FFr) expressed a concern with the verifiability of the above claim, and wrote that their decisions will be based on presuming doubt. We understand their thought process. **But it is not an ICLR policy to forbid discussion of public code, model release, or engagement in the discussion phase.** We agree that no reviewer should verify these claims themselves, but **if necessary, they could be verified without violating double blindness using the proxy method** – asking a colleague to anonymously fact-check. We have checked that ICLR’s policies and guidelines for review do not prohibit this. [s8MN agreed](https://openreview.net/forum?id=wI5uHZLeCZ&noteId=z0nRzywPZT) that our response was reasonable. 

**This is a disagreement that we ask the AC to adjudicate between us & [s8MN](https://openreview.net/forum?id=wI5uHZLeCZ&noteId=z0nRzywPZT) versus [rnFJ](https://openreview.net/forum?id=wI5uHZLeCZ&noteId=tPxm0G1FFr) & [EUQJ](https://openreview.net/forum?id=wI5uHZLeCZ&noteId=z7KSgsZ6QW). This is a matter of procedure, and not a matter of opinion.** Citing this concern, [rnfj](https://openreview.net/forum?id=wI5uHZLeCZ&noteId=tPxm0G1FFr) adjusted their score from an 8 to a 6, and [EUQJ](https://openreview.net/forum?id=wI5uHZLeCZ&noteId=z7KSgsZ6QW) opted to leave their score at a 3. We would ask that our discussion of released code and models not be ignored or presumed false because (1) they were the most direct way to reply to reviewer concerns, (2) code/model usage is relevant for evaluating paper quality/reproducibility/impact, (3) we did not ask reviewers to verify these themselves, (4) our claims could be verified anonymously using a proxy, and (5) concerns with verifiability are not unique to this type of claim (e.g., it's possible for authors to simply fabricate a paper's results).
</content>
</thread_object>

</thread>


Metareview:
metareview:
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.