PaperlayerView exact revision on arXiv

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

arXiv:2407.15549v3cs.LGcs.AIcs.CL

TL;DR

LLMs can retain undesirable behaviors despite fine-tuning, while prior untargeted latent adversarial training does not use information about specific failure modes. This paper tests targeted LAT against specific competing behaviors and reports improvements across jailbreak robustness, backdoor removal, and unlearning, including outperforming R2D2 on jailbreaks with orders of magnitude less compute.

  • Problem

    Prior latent adversarial training is untargeted and does not leverage information about specific failure modes, despite LLMs’ persistent undesirable behaviors after fine-tuning.

    Source: §1 · p. 2

  • Method

    The paper proposes targeted LAT, which perturbs hidden latent representations to target removal of a specific competing undesirable capability.

    Source: §1 · p. 2

  • Results

    Targeted LAT improved existing defenses across jailbreaks, backdoor removal, and unlearning with little or no general-performance tradeoff, and outperformed R2D2 on jailbreak robustness with orders of magnitude less compute.

    Source: §1 · p. 2, §5 · p. 11

  • Takeaways & Limitations

    The results support targeted LAT as a practical tool for improving LLM defenses against persistent harmful behaviors.

    Source: §5 · p. 11

    Targeted LAT reduces relearning sample efficiency, but unlearning remains vulnerable to few-shot fine-tuning, and methods need improvement against such attacks.

    Source: §5 · p. 12

Abstract

from arXiv · show

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

The introduction argues that fine-tuning may leave undesirable capabilities intact and presents targeted latent adversarial training (LAT) as a way to remove specific harmful behaviors. It reports benefits for jailbreak resistance, backdoor removal, and unlearning that is more robust to relearning.

1 Introduction

  • Unlike untargeted LAT’s loss-maximizing attacks, targeted LAT perturbs hidden representations to elicit specified undesirable behaviors, aiming to remove their circuitry and improve existing methods with little to no typical-use tradeoff.The authors hypothesize that latent representations encode higher-level features, making them a useful target for removing circuitry responsible for unwanted behaviors.Source: §1 · p. 2
  • LAT improves refusal training against jailbreaks and outperforms R2D2 (Mazeika et al., 2024) with orders of magnitude less compute.The introduction says LAT greatly improves refusal training’s ability to make LLMs robust to jailbreaks.Source: §1 · p. 2
  • With unknown triggers and vaguely specified responses, LAT improves DPO (Rafailov et al., 2024) for backdoor removal; the authors suggest it addresses the “Sleeper Agent” problem (Hubinger et al., 2024).Source: §1 · p. 2
  • LAT improves WHP, gradient ascent (Jang et al., 2022), and RMU (Li et al., 2024a) for unlearning, while substantially decreasing the sample efficiency of relearning.Source: §1 · p. 2

2 Related Work

Prior work studied latent adversarial training in vision and language models, while this work targets specific unwanted LLM behaviors. It applies targeted LAT to jailbreak robustness, backdoor removal, and unlearning.

Latent Adversarial Training (LAT)

  • Unlike untargeted LAT by Casper et al. (2024b), targeted LAT elicits specific unwanted outputs and aims for state-of-the-art defenses against jailbreaks, backdoors, and undesirable knowledge.Prior latent-space adversarial modifications and LAT had been studied in vision and language models.Source: §2 · pp. 2–3

LLM Robustness

  • Adversarial training can fail to fix vulnerabilities it did not target, while LAT improves robustness to unseen jailbreak and backdoor attacks.State-of-the-art LLMs remain vulnerable to novel attacks despite techniques including adversarial training.Source: §2 · p. 3

LLM Backdoors

  • Without trigger knowledge, removing backdoors is difficult, and adversarial training can even strengthen a “sleeper agent” backdoor.Backdoors can be planted by poisoning training data to produce harmful behavior when an arbitrary trigger appears.Source: §2 · p. 3

LLM Unlearning

  • Existing LLM unlearning methods have adversarial vulnerabilities, while LAT improves over unlearning techniques including state-of-the-art RMU (Li et al., 2024a).Source: §2 · p. 4

3 Methods

Targeted latent adversarial training perturbs latent activations to steer a model toward undesirable target behavior, using paired desirable and undesirable completions and prompt-only residual-stream attacks.

Targeted latent adversarial training

  • Targeted LAT attacks latent activations to steer the model toward an undesirable target behavior, unlike untargeted LAT, which steers it away from desired behavior.The model is represented as a feature extractor mapping text to latent activations and a function mapping those activations to output probabilities.Source: §3 · p. 4

Training methods

  • Basic targeted LAT requires datasets of desirable and undesirable behaviors, usually prompts paired with harmless and harmful completions.Interleaving LAT with benign supervised fine-tuning or applying KL regularization against the original model on benign data can stabilize training and reduce side effects.Source: §3 · p. 4
  • Attacks use L2-bounded PGD perturbations on prompt positions in the residual stream, typically across four evenly spaced layers, with the perturbation bound selected by grid search.The chosen hyperparameters were held fixed across later experiments on different models and tasks.Source: §3 · p. 4

4 Experiments

The experiments test targeted latent adversarial training (LAT) as an augmentation to reduce unwanted behaviors while preserving typical performance, with attention to the tradeoff between typical use and robustness. They evaluate jailbreak resistance, backdoor removal, and unlearning robustness.

4.1 Improving Robustness to Jailbreaks

  • RT-EAT-LAT achieved the best average robustness across five attacks; on Llama-2-7B-chat, it matched or beat R2D2 on all but one measure with over 700x fewer forward and backward passes.It used approximately 36x fewer GPU hours than R2D2. In Llama-3-8B-instruct, RT outperformed RT-EAT-LAT on all three general-capability measures; the StrongReject attack-success rates may be overestimates due to false positives.Source: Table 2 · p. 6, §4.1 · p. 7, §4.2 · p. 7

4.2 Backdoor Removal

  • DPO-LAT removed backdoors in all tested cases, while DPO alone had little or no effect and once strengthened the backdoor; LAT showed no substantial harm without a trigger.The experiments covered unknown triggers and approximate proxy reconstructions; DPO-LAT decreased MMLU by less than one percentage point relative to DPO.Source: Table 3 · p. 8, §4.2 · p. 8

4.3.1 Who’s Harry Potter?

  • WHP-C-LAT Pareto dominated WHP and WHP-C on all measures except MMLU.WHP-C adds a contrastive away objective to the original method, and LAT optimizes perturbations to minimize loss on the original Harry Potter text.Source: Table 4 · p. 9, §4.3.1 · pp. 9–10

4.3.2 Unlearning WMDP Biology and Cyber Knowledge

  • 15.7 percentage points of WMDP biology/cyber performance returned on average after up to 20 iterations of two-example retraining, but LAT reduced the recovered performance-gap share from 74.7% to 59.9%.GA-LAT outperformed GA across evaluations; RMU-LAT outperformed RMU except for a 1.2% decrease in MMLU and a 2.1% decrease in AGIEval.Source: §4.3.2 · pp. 10–11, Table 5 · p. 11

5 Discussion · A Broader Impacts · B Loss Functions for LAT

Targeted LAT strengthened existing defenses against jailbreaks, backdoors, and undesirable knowledge, with little or no tradeoff in general performance. The discussion emphasizes its potential and limitations, while future work considers broader latent attacks and the work’s value-neutral implications.

5 Discussion

  • Targeted LAT augmented existing defenses against jailbreaks, backdoors, and undesirable knowledge, improving removal of unwanted behaviors with little or no tradeoff in general performance.These results support, but do not yet confirm, the hypothesis that LAT removes neural circuitry responsible for undesirable behaviors; analyzing the mechanisms is left to future work.Source: §5 · p. 11
  • LAT improved generalization to attack targets, harmful behaviors, and knowledge-elicitation methods from datasets different from those used in training, without requiring precise knowledge of failure modes.The discussion frames this as useful robustness against persistent failures including jailbreaks, backdoors, and undesirable knowledge.Source: §5 · p. 12
  • Repeated retraining on only two forget-set samples closed more than half the original-to-unlearned performance gap on average, while targeted LAT reduced relearning sample efficiency.The authors note substantial room to improve unlearning robustness to few-shot fine-tuning attacks.Source: §5 · p. 12
  • LAT configuration was sensitive to dataset, layer, and perturbation size, and experiments were limited to residual-stream perturbations with projected gradient descent and models under 10 billion parameters.Interleaved supervised fine-tuning and NaN handling were important for stable training; usefulness on larger models remains to be confirmed.Source: §5 · p. 12

Future work

A Broader Impacts

  • The authors describe LAT as value-neutral: it aligns systems with developers’ goals, and AI alignment should not be conflated with safety.They identify misuse and adverse systemic effects, rather than only unintended hazards, as sources of AI-system risk.Source: A Broader Impacts · p. 21

B.1 RT-EAT-LAT · B.2 DPO-LAT

RT-EAT-LAT and DPO-LAT train latent adversaries to favor harmful responses while training the model to favor harmless ones. They differ in their preference data and losses, and RT-EAT-LAT also offers alternative benign-behavior losses.

B.1 RT-EAT-LAT

  • RT-EAT-LAT’s adversary minimizes loss by moving toward harmful completions and away from harmless completions on paired harmful-request data.The method assumes harmful requests paired with preferred harmless and rejected harmful completions.Source: §B.1 · p. 21
  • RT-EAT-LAT constrains each latent perturbation δ_i to ||δ_i||2 ≤ ϵ before training the model against the adversarial attacks.ϵ is a hyperparameter.Source: §B.1 · p. 21
  • Against those attacks, RT-EAT-LAT trains the model to move toward harmless completions and away from harmful completions.These are the two components of its defense loss.Source: §B.1 · p. 21
  • RT-EAT-LAT uses benign SFT loss for Llama 2 and benign KL loss for Llama 3, calculated without an attack; the authors report KL better preserves capabilities in more capable models.SFT maximizes ground-truth completion probability, while KL preserves the reference model’s logits; the authors hypothesize KL may help when benign data quality is poor relative to the model.Source: §B.1 · pp. 21–22

B.2 DPO-LAT

  • DPO-LAT, inspired by Rafailov et al. (2024), uses Anthropic HH-RLHF’s harmless split for harmful/harmless preference pairs and helpful split for helpful/unhelpful pairs.The chosen responses are harmless or helpful, and the rejected responses are harmful or unhelpful, respectively.Source: §B.2 · p. 22
  • DPO-LAT’s adversary uses DPO loss on harmful/harmless pairs to prefer the harmful response, moving toward harmful and away from harmless completions.The preference pairs are sampled from D_p.Source: §B.2 · p. 22
  • DPO-LAT uses DPO loss for both defense and benign terms, with the adversarial perturbation present for the defense term and absent for the benign term.The defense and benign terms use D_p and D_b, respectively.Source: §B.2 · p. 22

B.3 WHP-C-LAT and GA-LAT

The section reports that untargeted LAT provides less jailbreak robustness than targeted LAT and specifies the loss setup used for WHP-C-LAT and GA-LAT.

B.3 WHP-C-LAT and GA-LAT

  • WHP-C-LAT and GA-LAT use toward-only adversaries optimizing next-token cross-entropy on Harry Potter and WMDP forget corpora, respectively.WHP follows Eldan & Russinovich (2023); WMDP uses a log(1 −p) away loss on the forget dataset as in Mazeika et al. (2024). The datasets are described as examples to forget (Df) and retain (Db).Source: §B.3 · p. 22
  • Both methods add a toward loss on WikiText and an SFT loss on Alpaca, keeping adversarial perturbations for toward and away losses but removing them for SFT.The WikiText loss matches Li et al. (2024a).Source: §B.3 · p. 22
  • Untargeted LAT provides less jailbreak robustness than targeted LAT in the reproduced Table 2 comparison.Table 6 adds an untargeted-LAT row to the bottom part of Table 2.Source: Table 6 · p. 23

B.4 RMU-LAT

RMU-LAT uses the RMU loss from Li et al. (2024a), with the adversary optimizing next-token cross-entropy on WMDP forget corpora while its perturbation is present in the forget-loss calculation.

B.4 RMU-LAT

  • RMU-LAT applies the adversary’s perturbation when calculating forget loss on WMDP corpora, using next-token cross-entropy as in Li et al. (2024a).The normalized random vector u is sampled uniformly from [0, 1] and held constant throughout training; for Zephyr-7B, c = 6.5 and α = 1200, following Li et al. (2024a).Source: §B.4 · p. 23

C Jailbreaking Robustness Under Untargeted LAT

Untargeted LAT trains the adversary to make the model fail to refuse, rather than to comply with a jailbreak. It generally yields comparable or slightly worse robustness than targeted LAT, but performs much worse against prefill and GCG attacks.

C Jailbreaking Robustness Under Untargeted LAT

  • Untargeted LAT harms general performance less than targeted LAT (but not refusal training), yields comparable or slightly worse robustness in most cases, and fares much worse on prefill and GCG.During untargeted LAT, the adversary works only to make the model fail to output a refusal.Source: C Jailbreaking Robustness Under… · p. 23

D Jailbreaking Robustness Under an Alternate Autograder · E Backdoored Model MMLU Performance

Under HarmBench, attacks are labeled successful more often than under StrongReject, while qualitative method comparisons remain similar. On one backdoored model, DPO-LAT reduces MMLU performance by slightly less than one percentage point compared with DPO.

D Jailbreaking Robustness Under an Alternate Autograder

  • HarmBench is more likely than StrongReject to label attacks successful, but the overall trends remain similar.The alternate-autograder results are reported in Table 7.Source: D Jailbreaking Robustness Under an… · p. 23, Table 7 · p. 24
  • Table 7 reproduces the jailbreaking results with HarmBench in place of the StrongReject autograder used in Table 2.Source: Table 7 · p. 24
  • The qualitative comparisons between methods under HarmBench are similar to those in Table 2.Source: Table 7 · p. 24

E Backdoored Model MMLU Performance

  • Table 8 presents the MMLU performance comparison for backdoored models and describes LAT’s decrease relative to DPO as less than one percentage point.Source: Table 8 · p. 24
  • For backdoor removal, the study evaluates DPO-LAT and DPO destructiveness by measuring each model’s MMLU performance.The reported results cover a single model.Source: E Backdoored Model MMLU Performance · p. 24
  • LAT decreases MMLU performance by slightly less than one percentage point compared with DPO.Table 8 reports this comparison for a single model.Source: E Backdoored Model MMLU Performance · p. 24, Table 8 · p. 24

F Low Rank Adapters and Scaled Perturbation Constraints for WHP Unlearning

This section tests targeted LAT for WHP unlearning using PCA-whitened latent perturbations and rank-64 adapters. The approach yielded weaker unlearning but less tradeoff in general capabilities, and the authors leave further investigation of scaling for unlearning to future work.

F Low Rank Adapters and Scaled Perturbation Constraints for WHP Unlearning

  • Scaling and rank-64 adapters yielded weaker WHP unlearning alongside less tradeoff in general capabilities.Table 9 reports the scaling results and compares them with Table 4 in the main paper.Source: F Low Rank Adapters and Scaled… · p. 24, Table 9 · p. 25
  • For WHP, they PCA-whitened Harry Potter and coherent-genericized-text activations, constrained Lp perturbations there before de-whitening them onto latent states, and trained rank-64 adapters on all linear modules.The perturbations were added to the residual stream after de-whitening.Source: F Low Rank Adapters and Scaled… · p. 24
  • The authors speculate that unlearning may be especially suited to this scaling approach, but leave deeper investigation for future work.Source: F Low Rank Adapters and Scaled… · p. 25

G Tests for Robust and Competitive Unlearning in LLMs · H WMDP Unlearning Details

Section G evaluates Harry Potter unlearning with Familiarity under standard and prompt-extraction conditions; Section H specifies WMDP unlearning training and attack layers.

G Tests for Robust and Competitive Unlearning in LLMs

H WMDP Unlearning Details

  • For GA and GA-LAT, the authors use rank-64 LoRA; for RMU and RMU-LAT, they train full-rank MLP weights without LoRA, following Li et al. (2024a).LoRA is cited as Hu et al. (2021).Source: H WMDP Unlearning Details · p. 28
  • RMU trains layers ℓ, ℓ−1, and ℓ−2 and matches activations at ℓ; attacks on ℓ−2 require a smaller perturbation radius, and attacks work best one layer before ℓ.The layer choices follow Li et al. (2024a); the smaller radius avoids preventing successful unlearning.Source: H WMDP Unlearning Details · p. 28

Related papers

Loading 2407.15549v3…