Improving Alignment and Robustness with Circuit Breakers

Summary

This paper introduces a new technique called "short circuiting" that makes models more robust against outputting harmful texts without significant impacts to their helpfulness (a significant pareto improvement on the harmfulness-helpfulness tradeoff). Short circuiting works by training a model's hidden representation on harmful prompts to be perpendicular to the original representations, so that when it sees harmful prompts in the future its inputs will be off distribution. There is also an additional loss term that ensures that the model's representations on a set of normal prompts remains the same to preserve benign helpful behavior. This method achieves state of the art performance on adversarial attacks on text models, text+image models, and on an agent benchmark the authors create.

Strengths

- The method is very effective and general: short circuiting makes models significantly more robust to outputting harmful information while remaining almost the same level of helpful. - It is an incredibly positively impactful problem, and the results can be applied immediately to existing large scale multi-modal models. - This is an entirely new class of method for adversarial training that may now be built upon and improved. - The authors introduce a new harmful function calling dataset.

Weaknesses

- The presentation and paper organization could be improved in many places (see questions section) - There are many experiment details ommited - Not clear if the method will work on types of harmful output not seen in the fine tuning data - The harmful knowledge is still in the model, so it cannot be open sourced. Even if it was trained normally in Algorithm 1 (instead of LoRA) and released, bad actors may still be able to fine-tune the model to restore the harmful capabilities, as recent work shows. - In the checklist, the authors say they specify the compute resources used in the appendix, but I was not able to find this information.

Questions

Section 3 (the description of the methods) could be much improved: - It would help if there was more intuition on why short-circuiting makes sense as a method, for example by explaining that it seeks to put the hidden representations as much off-distribution as possible - The Loss subsection describes many potential loss methods that are brought up and then dismissed as explored. It is confusing which loss is used in the actual methods. - Line 149 states "We find that a practical approach is to remove harmful user requests from the Short Circuit Set." which does not make sense to me, as I thought the goal of the short circuit set was to precisely include harmful user requests. - This section is very long with prose and generally includes many unneeded details, e.g. much of paragraph 111. On the other hand, crucial details like which/how many layers are trained with LoRA are unstated. The second would benefit from more formalism of the method in the body of the text. Section C.1 should ideally also be merged into section 3 for similar reasons. - The paragraph on 212 discussing the results should focus more on the numbers and less on describing the virtues of short circuiting (which already occurs plenty in the introduction and the conclusion). To be clear, I think it's a great technique! But the results section should describe the results more objectively. - Possibly it would be helpful to show a non representation engineering baseline to show that the method is not overfitting to the training data. E.g. could do a LoRA finetuning run on the same dataset where the loss is on the output and towards refusals/maintaining original behavior. - Sections 4.1, 4.2, and 4.3 could use an intro paragraph at the top describing the problem setup. For 4.3 especially this is quite jarring. - It is mostly unclear what the short circuit and retain datasets are for each method. For example, how many samples do they contain? What are the "extra refusal data points" mentioned on line 182? - It would be helpful to explain all of the attack types in the list on line 188. - Typo: extra and on line 189

Rating

7

Confidence

4

Soundness

4

Presentation

2

Contribution

4

Limitations

Somewhat, the authors touch on only one limitation. The limitations section would be improved if it included a deeper discussion, e.g. on generalization to unseen harm classes or vulnerability to further fine-tuning.

Summary

This paper presents a novel short-circuiting method to ensure LLM safety alignment and its adversarial robustness against jailbreaking attacks. Particularly, the authors propose to conduct LoRA training with their short circuiting loss in the model representation space. The short-circuited models can effectively resist various jailbreaking attempts.

Strengths

1. The idea behind the proposed method is novel. Instead of trying to capture and patch against different jailbreaking attacks, the short circuiting method directly looks into the model representation space, and steers model representation at harmful state towards a short circuit direction. 2. The experiment results are solid, spanning across various jailbreaking attacks (whitebox & blackbox), 2 different models (Llama-3-8B-Instruct and Mistral-7B-Instruct-v2), and 3 different scenarios (LLM, MLLM, LLM Agent). 3. The authors also verify that the short-circuited models can retain similar-level model capabilities on LLM benchmarks like MT-Bench and MMLU.

Weaknesses

Overall, I enjoy reading the paper quite much. But here are some of my major concerns: 1. **More insights into the method rationale.** The idea of short circuiting is impressive. Nevertheless, I find some the rationale of it not that intuitive. For example, in Line 44-45, the authors claim that "when a (short-circuited) model begins to generate such an output, its internal processes are interrupted, halting completion of the generation." However, the formulas in Algorithm 1 and Sec 3 only focus on steering SC model's representation $\text{rep}_{\mathcal{M}_\text{sc}}(x_s)$ away from the original model's representation $\text{rep}_{\mathcal M}(x_s)$, at the state where they are outputting a harmful output. This does not explain why a model would halt when it "begins to" provide a harmful output. I appreciate the visualization of representations in Fig 5, which partially help justify the claim in Line 44-45. Would be great to provide more insights why this would happen after training models via Algorithm 1. 2. **More details on Short Circuit Set.** I don't think the paper and apppendix provide enough information on how this dataset should be collected. 3. **Adaptive attacks?** The authors have claimed that short-circuiting can defend well against various types of adversarial jailbreaking attacks, including multiple white-box ones. Particularly, the authors highlight how short-circuiting is immune to PGD attack against multimodal models. I wonder whether these white-box attacks take into consideration of the LoRRA short-circuiting module, i.e., whether the attackers are (un)aware of the existence of the short-circuiting defense? If yes, why are they less effective than some black-box attacks (e.g., Input Embed is less effective than PAIR and TAP-T in Table 1)? If no, it'd be necessary to add an adaptive analysis on your method. I believe it'd be crucial to thoroughly discuss and highlight the potential adaptive vulnerabilities of short circuiting. 4. **More model capability evaluation.** It's important to validate the capability is well retained after short-circuiting. While I appreciate the authors' efforts on MT-Bench and MMLU, it'd be great to see more capability results (e.g., GSM8K, ARC-c, AlpacaEval, etc.). I'm also curious about the win rate of model responses before & after short-circuit training. 5. **Tailoring Short Circuiting.** In line 125-126, you mentioned that short-circuiting can be tailored to "prevent more narrowly tailored types of output." Can you provide any experimental results to justify this? For example, show how to apply short-circuiting to prevent models from generating copypright content, but not other types of potentially harmful content. I notice that in Fig 4, you have ablated across 6 types of harm -- the results there only show how training over one type of harm would generalize to other types.

Questions

1. In Table 1, why is PAIR the most successful attack (23.3 ASR), even better than white-box attacks like Input Embed (15.7 ASR)? The author claim these white-box attacks are "proxies for worst-case assumptions aboout attacker capabilities" (Line 54), but they don't seem to outperform other attacks against Short Circuiting. 2. Line 150: What do you mean by "removing harmful user requests" from the Short Circuit Set? Why? 3. Are there any ablation results on full parameter fine-tuning (showing it doesn't work stably), instead of LoRA? I'm asking this because sometimes LoRA may not be a preferred option (compared to full fine-tuning) to implement LLM safety -- malicious users can simply disable it to obtain harmful model outputs.

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

The authors have discussed some limitations in Sec 5. However, I don't think that discussion fully elaborates potential (adaptive) vulnerabilities of the proposed method. It would be necessary to highlight / discuss in what case the proposed short circuiting method may fail (in the context of jailbreaking).

Summary

This paper proposes a method that manipulates the representation of potentially harmful queries to reduce the harmful generation. Extensive experiments are conducted to verify the effectiveness.

Strengths

Experiments on different scenarios are provided to illustrate the effectiveness of the proposed method.

Weaknesses

1. The distinct difference between the proposed method and adversarial training is not clear. Generally speaking, adversarial training allows the defender to modify the model based on the adversarial examples to improve the robustness. The proposed method also fine-tune the LLM on potentially harmful queries, which can somehow be considered as a generative-model-version adversarial examples, to reduce the generation of harmful content. 2. The idea of manipulating the representation is not a new one. [1] has proposed to move the representation of harmful queries towards rejection. Then the contribution of this work is not very significant. 3. The details of representation used in this work are not clear. [1] On Prompt-Driven Safeguarding for Large Language Models. Chujie Zheng et al. ICML2024

Questions

See weakness.

Rating

4

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Yes

Summary

The authors propose and test a novel method called "short-circuiting" which aims to modify a generative model $M$ so that it is unable engage in a specified behavior $B$. The high-level idea of short-circuiting is as follows: 1. First, transcripts demonstrating the behavior $B$ are collected. The is called the "Short-Circuit Set". 2. Next, transcripts demonstrating behaviors that $M$ should retain are collected. This is called the "Retain Set". 3. A copy $M_\text{sc}$ of the model $M$ is made and $M_\text{sc}$ is fine-tuned with two objectives: - Short-circuit objective: for transcripts $x_\text{sc}$ from the short-circuit set, the internal activations of $M_\text{sc}$ when processing $x_\text{sc}$ should look different from the internal activations of $M$ when processing $x_\text{sc}$. - Retain objective: for transcripts $x_{r}$ from the retain set, the internal activations of $M_\text{sc}$ when processing $x_\text{r}$ should look very similar to the internal activations of $M$ when processing $x_\text{r}$. The authors try a few different implementations of this high-level scheme, and evaluate their best performing implementation against a suite of jailbreak attacks designed to elicit harmful completions, including attacks that get white-box access to the model. Evaluations are run for language models and vision-language models. The authors show that attacks have a much harder time eliciting harmful completions from the short-circuited models as compared to the original non-short-circuited models. Moreover, compared to the non-short-circuited models, the short-circuited models show essentially no degradation in performance on non-harmful tasks like MMLU and MTBench.

Strengths

1. [Quality] Short-circuiting is a simple, cheap, and general method for preventing models from engaging in specified behaviors. 2. [Originality] Short-circuiting seems novel based on my understanding of the literature. 3. [Quality] Short-circuiting shows good empirical performance at defending against attacks, greatly reducing attack success rate compared to non short-circuited models without damaging performance on MMLU and MT Bench. 4. [Significance] Short-circuiting can be quite naturally applied to prevent LLMs from exhibiting clearly bad behavior. This is an open and important problem in the field of AI safety, thus the method has actual real-world consequences.

Weaknesses

The main thing the paper lacks is an analysis of why short-circuiting fails to reduce attack success rates to *zero*. Achieving zero ASR is one of the end-goals of adversarial robustness research. The paper would be significantly stronger if it provided an explanation of why short-circuiting does not meet this objective and suggested possible improvements to short-circuiting.

Questions

**Major suggestion** I feel that you should include transcripts of successful and unsuccessful attacks for both baseline models and short-circuited models. This would help the reader develop better intuitions around what short-circuiting does, the situations where it works, and the situations where it doesn't work. This might also give some insight into why short-circuiting is unable to achieve 0 ASR. This would also address concerns around the quality of jailbreak evaluation as raised by papers like "A StrongREJECT for Empty Jailbreaks" (https://arxiv.org/abs/2402.10260). **Minor questions** - What is the “System Level” attack in Figure 2? It would be good to match the axis labels here with the different attacks from Table 1?

Rating

8

Confidence

4

Soundness

4

Presentation

3

Contribution

4

Limitations

See weaknesses and question sections.

Summary

This work presented a method to short-circuit model when they attempt to output harmful content. This is conducted through representation engineering the internal representation of short-circuit model (LoRRA adaptor + forzen original model) corresponding to harmful content towards an objective orthogonal to it's original representation of original model through a short-circuit dataset and short-circuit loss, but keep model performance by keep those two representation close on retain dataset using retain loss. The method is model invariant and simple to apply. The work demonstrate three cases: LLM, MLLM, and AI Agents, using this method, they reduce ASR by a significant amount for different attack types including unseen ones comparing to refusal training and adversarial training, while retain the original model performance better compare two the other defense baselines. The author presented comprehensive experiments for the three cases also provided ablation and analysis discussion.

Strengths

1. The defense techniques proposed here by design could be invariant to input but focus on representation that generates harmful output, therefore have nice generalization potential. 2. Strong generalization across different attacks and even unseen attacks comparing to refusal training and adversarial training. 3. Strong performance for mitigating attacks targeting LLM, MLLM, and AI Agents. 4. High performance retain compare to other defense techniques. 5. The method is simple and model invariant.

Weaknesses

Writing: 1. Figure 4 is not used anywhere, I suppose you refer it around section 4.4 2. Table 1 is not used, it should be used around line 212? For table 1, the baseline: refusal trained and adv trained are not described in section 3. Also the scores reported in the table is not clear: I assume for above half is ASR, what is used for bottom half? (typo of MT-Bench result?) Experiment setting and results are not well explained, check questions below. 4. Figure 1 is not referred anywhere.

Questions

1. I am confused about the explanation at line 140 - 143: what does training data referring to here, and for "performance is generally improved when the training data better aligns with the domains we aim to short-circuit and retain.": Do you mean the opposite that short-circuit and retain data should be aligned to the training data of the model? 2. From line 145 to line 150, I could understand your may want to remove harmful user requests from short circuit set to address when an attack by pass the refusal mechanism, but description here is not clear to me, especially sentence from 145 to 148 where you describe the necessity of adding refusal data to retain set. It is helpful to explain that whole paragraph more clearly so that later ablation result in Figure 4 make more sense. 3. From Table 1: when No Attack, the refusal trained and adv trained model still have high ASR? better explanation is needed. I am willing to increase score if my questions are addressed.

Rating

7

Confidence

3

Soundness

2

Presentation

1

Contribution

3

Limitations

Author stated their method is designed to preventing harmful content generation when harmful is defined by short-circuit data, but the method is not suitable for attacks that changing class label since class label can not be defined as harmful or not in general.