Source-linked AI summary
Latent Adversarial Training Improves Robustness to Persistent Harmful Behaviors in LLMs
Abhay Sheshadri, Aidan Ewart, Phillip Guo, Aengus Lynch, Cindy Wu, Vivek Hebbar, Henry Sleight, Asa Cooper Stickland, Ethan Perez, Dylan Hadfield-Menell, Stephen Casper
TL;DR
LLMs can retain harmful capabilities despite fine-tuning, while prior untargeted latent attacks do not target specific failure modes. This paper introduces targeted LAT, which elicits specified undesirable behaviors during latent-space training, and finds improvements across jailbreak defense, backdoor removal, and unlearning robustness.
Problem
LLMs can persistently exhibit undesirable behaviors despite fine-tuning, and untargeted LAT does not leverage information about specific failure modes.
Method
Targeted LAT perturbs residual-stream latent activations to elicit undesirable behaviors, then fine-tunes models on target tasks under those perturbations.
Results
Targeted LAT augments existing methods across jailbreak robustness, trigger-unknown backdoor removal, and unlearning, while reducing the sample efficiency of relearning.
Takeaways & Limitations
Targeted LAT can strengthen existing defenses against persistent harmful behaviors with little or no tradeoff in general performance.
Takeaways & Limitations
Unlearning methods remain vulnerable to few-shot relearning attacks, leaving substantial room to improve robustness.
Abstract
from arXiv · showhide
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 Introduction
LLMs can retain undesirable behaviors despite fine-tuning, motivating targeted LAT to remove specific failure modes. The paper applies targeted LAT across jailbreak robustness, backdoor removal, and unlearning.
- Targeted LAT is proposed to more thoroughly remove persistent undesirable behaviors from LLMs.
- Targeted LAT can combine with and improve over a wide range of existing techniques.
- Orders of magnitude less compute lets LAT outperform R2D2 for improving refusal training’s robustness to jailbreaks.
- LAT improves DPO’s ability to remove LLM backdoors when the trigger is unknown and the response is only vaguely specified.
- Targeted LAT improves unlearning methods while substantially decreasing the sample efficiency of relearning previously unlearned knowledge.
2 Related Work
Prior work documents persistent LLM vulnerabilities, including jailbreaks, backdoors, and brittle unlearning. Targeted LAT differs from untargeted approaches by eliciting specific unwanted behaviors during latent-space training.
- Latent adversarial training has been studied in vision and language models, including untargeted defenses against backdoors and unforeseen adversarial attacks.
- Targeted LAT differs by steering the adversary toward specific outputs corresponding to unwanted behaviors.
- State-of-the-art LLMs remain vulnerable to novel attacks, and adversarial training can fail to fix vulnerabilities absent from training.
- Backdoors are difficult to remove without trigger knowledge, and adversarial training can even strengthen a sleeper-agent backdoor.
- LLM unlearning methods are motivated partly by removing harmful capabilities but existing methods suffer from adversarial vulnerabilities.
3 Methods
Targeted LAT trains LLMs against perturbations to hidden latent representations that elicit specified undesirable behaviors. The method uses desirable and undesirable behavior data, residual-stream attacks, and stabilizing benign training signals.
- An LLM is modeled as a feature extractor mapping text to latent activations followed by an output function producing a sampling distribution.
- Standard adversarial training perturbs inputs, whereas latent adversarial training makes the model robust to attacks on latent activations.
- Untargeted LAT maximizes loss on desired behavior, while targeted LAT steers the model toward a specified undesirable target behavior.
- Basic targeted LAT uses desirable and undesirable behavior datasets, often pairing prompts with harmless and harmful completions.
- Residual-stream perturbations are L2-norm-bounded and computed with projected gradient descent, while benign fine-tuning or KL regularization can stabilize training and reduce side effects.
4 Experiments
The experiments test targeted LAT as an augmentation to fine-tuning and adversarial-training methods for jailbreak robustness, backdoor removal, and knowledge unlearning. Across these settings, LAT improves unwanted-behavior removal or robustness while generally preserving typical performance, though effects vary by model and method.
- Experimental scope: Targeted LAT augments fine-tuning and adversarial-training methods across jailbreak defense, backdoor removal, and knowledge unlearning experiments.The study focuses on Pareto tradeoffs between typical-use performance and robustness to unwanted behaviors.
- Jailbreak robustness: Across five attack methods, RT-EAT-LAT achieves the best average jailbreak robustness and uses approximately 36x fewer GPU hours than R2D2.On Llama2-7B-chat, it performs as well or better than R2D2 on all but one reported measure with over 700x fewer forward and backward passes; general-capability results vary by model.
- Limitations: Manual red-teaming still found prompting strategies that could jailbreak the Llama3-8B RT-EAT-LAT model, despite its being relatively difficult to jailbreak.The authors describe the interface as a research demo rather than a product using all available robustness techniques.
- Backdoor removal: DPO-LAT removes backdoors successfully across tested cases when the trigger is unknown or only approximately reconstructed, whereas DPO alone provides little or no improvement.DPO-LAT shows no substantial trigger-free harm, with less than a one percentage point MMLU decrease relative to DPO.
- Harry Potter unlearning: WHP-C-LAT Pareto dominates WHP and WHP-C on all reported measures except MMLU for Harry Potter unlearning.The evaluation tests both general capabilities and robustness of unlearning under translated, jailbreak, summary-context, and text-context prompts.
- WMDP biology and cyber unlearning: GA-LAT and RMU-LAT improve biology and cyber unlearning, with RMU-LAT trading off 1.2% MMLU and 2.1% AGIEval decreases.LAT methods also resist relearning better: relearning closes 59.9% of the original performance gap on average, versus 74.7% for non-LAT methods.
5 Discussion
Targeted LAT augments existing defenses against persistent harmful behaviors across jailbreaking, backdoor removal, and unlearning, while retaining important limitations in mechanism evidence, tuning, and scale.
- Discussion: Targeted LAT augments state-of-the-art fine-tuning and adversarial training methods across jailbreaking, unlearning, and backdoor removal.The paper reports little or no tradeoff in general performance across these applications.
- Limitations: Unlearning remains brittle: retraining on only two forget-set samples consistently recovered more than half of the original-to-unlearned performance gap on average.Targeted LAT reduced the sample efficiency of relearning, but the paper identifies substantial room for improvement.
- Limitations: LAT configuration is sensitive to dataset, layer selection, perturbation size, training interleaving, and NaN handling, and experiments used models with fewer than 10 billion parameters.The authors expect usefulness to extend to larger models but state that this requires confirmation.
B.1 RT-EAT-LAT
RT-EAT-LAT trains latent adversaries toward harmful completions and the model toward harmless completions, while benign losses preserve helpful behavior. DPO-LAT applies the same targeted-adversarial logic to preference data.
- RT-EAT-LAT: The attack minimizes loss by moving toward harmful completions and away from harmless completions for paired preference examples.The harmful and harmless responses are represented by r_i and c_i, respectively, in the paired dataset.
- RT-EAT-LAT: The model defense reverses this direction, moving toward harmless completions and away from harmful completions under adversarial perturbations.The model parameters are optimized against the latent attack through Ldefense and Lbenign.
- RT-EAT-LAT: The adversarial perturbation is constrained by ||δ_i||2 ≤ ϵ, limiting the attacker’s power.The constraint uses ϵ as a hyperparameter.
- RT-EAT-LAT: Benign training uses either supervised fine-tuning or KL regularization on unattacked inputs to preserve helpful behavior and original model capabilities.The SFT term maximizes ground-truth completion probability, whereas the KL term preserves original logits.
- DPO-LAT: DPO-LAT uses paired helpful and unhelpful responses for both the targeted attack and the defense, with the adversary preferring the harmful response.The attack uses DPO loss, while defense and benign objectives apply DPO loss with and without the adversary, respectively.
B.3 WHP-C-LAT and GA-LAT
WHP-C-LAT and GA-LAT use toward-only adversaries optimized on forget corpora, while also including WikiText toward loss and supervised retention terms.
- WHP-C-LAT and GA-LAT: WHP-C-LAT and GA-LAT use toward-only adversaries optimized for next-token cross-entropy loss on the Harry Potter and WMDP forget corpora.WHP follows Eldan and Russinovich’s training procedure, while WMDP uses a log(1 − p) away loss.
- WHP-C-LAT and GA-LAT: Both methods additionally include a toward loss on WikiText and a supervised retention component.
- WHP-C-LAT and GA-LAT: The losses are defined over forget examples Df and retain examples Db, with token-level terms indexed by each string’s j-th token.Here, ti,j denotes the j-th token and ti,<j the preceding tokens in the same string.
- WHP-C-LAT and GA-LAT: Untargeted LAT steers the model away from desired behavior, whereas targeted LAT steers it toward examples of undesirable behavior.
B.4 RMU-LAT
RMU-LAT retains the RMU loss while applying the adversary’s perturbation when calculating forget loss, with fixed hyperparameters inherited from prior work.
- RMU-LAT: RMU-LAT uses the same RMU loss as Li et al. (2024a), with the adversary’s perturbation present when forget loss is calculated.
- RMU-LAT: The adversary optimizes next-token cross-entropy loss on the WMDP forget corpora.
- RMU-LAT: The RMU formulation uses a normalized random vector u and sets c = 6.5 and α = 1200 for Zephyr-7B.The vector is sampled uniformly from [0, 1] and remains constant throughout training.
C Jailbreaking Robustness Under Untargeted LAT
Untargeted LAT generally provides comparable or slightly worse jailbreak robustness than targeted LAT, while preserving general performance better than targeted LAT but not refusal training.
- Comparison with targeted LAT: Untargeted LAT causes less harm to general performance than targeted LAT, but not less than refusal training.
- Comparison with targeted LAT: Untargeted LAT has comparable or slightly worse robustness than targeted LAT for most attack methods.
- Comparison with targeted LAT: Untargeted LAT performs much worse than targeted LAT against prefill and GCG attacks.
D Jailbreaking Robustness Under an Alternate Autograder
HarmBench labels attacks as successful more often than StrongReject, but the reported qualitative comparisons between methods remain similar; LAT’s MMLU cost is small in the backdoor-removal evaluation.
- HarmBench evaluation: HarmBench is more likely than StrongReject to label jailbreak attacks as successful, while overall trends within the results remain similar.
- HarmBench evaluation: Table 7 reports jailbreak results using HarmBench instead of the StrongReject autograder used in Table 2.
- Backdoor-removal side effects: LAT decreases MMLU performance by slightly less than one percentage point compared with DPO in the backdoor-removal evaluation.The comparison is reported for a single model.
F Low Rank Adapters and Scaled Perturbation Constraints for WHP Unlearning
This section tests low-rank adapters and scaled, whitened-space latent attacks for Harry Potter unlearning. Scaling trades somewhat weaker unlearning for better general-performance tradeoffs, while targeted LAT remains part of the proposed approach.
- Scaled Perturbation Constraints: The method constrains latent perturbations according to activation-manifold geometry rather than an isotropic Lp-norm ball.The authors motivate whitening because latent activations vary more along some directions than others.
- Scaled Perturbation Constraints: PCA is applied to activations from Harry Potter text and coherent genericized text to target meaning-preserving features for unlearning.The perturbations are optimized and constrained in a whitened space.
- Limitations: The authors leave deeper investigation of whether this scaling is especially suitable for unlearning to future work.
- Results: Training with scaling produces less strong Harry Potter unlearning but better tradeoffs in general performance.
- Results: The WHP and WHP-C rows report the evaluation values for baseline unlearning variants, while WHP-C-LAT reports the targeted-LAT variant.The supplied table excerpt lists six values with uncertainty for each row.
G Tests for Robust and Competitive Unlearning in LLMs
This section evaluates Harry Potter unlearning against the WHP method and its targeted-LAT replication using familiarity-based tests. The evaluation includes standard completions and attacks intended to resurface suppressed knowledge.
- Evaluation Setup: The study compares Eldan and Russinovich’s WHP model, replications, and a replication augmented with targeted LAT.WHP fine-tunes Llama-2-7B-Chat on text replacing domain-specific content with generic content.
- Evaluation Setup: Familiarity measures Harry Potter content in completions of related sequences using an automated GPT-4 evaluation.The evaluation prompt includes the datapoint reference, prompt, and model completion.
- Robustness Tests: The experiments evaluate familiarity under prompt-extraction attacks in addition to standard familiarity evaluations.
- Robustness Tests: The test suite includes Spanish-translated prompts to assess whether Spanish fine-tuning transfers across languages.
- Robustness Tests: A jailbreak prompt is used to pressure the model to resurface Harry Potter knowledge, alongside few-shot and summary prompts that provide limited related context.Summary prompting uses either a high-level Harry Potter summary or the first 10 lines of Book 1.
- Implementation: The setup uses LoRA rank 64 for GA and GA-LAT, whereas RMU and RMU-LAT train MLP weights full-rank.Layer choices for the adversary, RMU training, and activation matching are varied within the setup.