Robust LLM safeguarding via refusal feature adversarial training

Authors: Lei Yu, Virginie Do, Karen Hambardzumyan, Nicola Cancedda
Primary area: interpretability and explainable AI
Keywords: large language models, adversarial robustness, representation analysis, mechanistic interpretability

Abstract:
Large language models (LLMs) are vulnerable to adversarial attacks that can elicit harmful responses. Defending against such attacks remains challenging due to the opacity of jailbreaking mechanisms and the high computational cost of training LLMs robustly. We demonstrate that adversarial attacks share a universal mechanism for circumventing LLM safeguards that works by ablating a dimension in the residual stream embedding space called the refusal feature. We further show that the operation of refusal feature ablation (RFA) approximates the worst-case perturbation of offsetting model safety. Based on these findings, we propose Refusal Feature Adversarial Training (ReFAT), a novel algorithm that efficiently performs LLM adversarial training by simulating the effect of input-level attacks via RFA. Experiment results show that ReFAT significantly improves the robustness of three popular LLMs against a wide range of adversarial attacks, with considerably less computational overhead compared to existing adversarial training methods.

Decision: accept (poster)
Reviewer scores: 6, 6, 3, 8

Reviews and discussion:
<thread>
creation date:
1729848864448

thread_type:
official_review

<thread_object>
creator: reviewer_ocf9

<content>
summary:
The authors propose a novel robustification method in the context of LLMs that is built open, ablating the "refusal direction" of a given model in its latent space during training. The authors demonstrate how this approach relates to standard adversarial training / attacks and provide evidence for the effectiveness of their method.

strengths:
* The paper conducts various experiments to provide mechanistic insights onto how the proposed algorithm might be used as an alternative to adversarial training. 
* The authors also try to assess the limitations of the proposed approach (i.e., generalization to cases where the refusal direction might not be accurate anymore) 
* Comparisons to very recent baseline approaches are provided
* The proposed approach is considerably more efficient compared to other robustification methods

weaknesses:
* As far as I am aware the results from Figure 1 and Figure 2 are known and only provided to explain the concept. It would be nice if this is put into the right context. 
* The refusal results do not appear to be consistent with those of Xhonneux et al. (2024). A more comprehensive assessment of refusal behavior, e.g. with the OR-Bench, would further strengthen the paper 
* How stable are the different robustification methods with respect to hyperparameter choices? Would be interesting to assess how close each method is to its respective “optimum” specifically because the methods are directly compared to each other. Currently, I think it's not possible to give an accurate estimate of differences in accuracy robustness trade-offs between the different methods. However, the efficiency of the proposed approach is clear.

questions:
* Can the authors elaborate on hyperparameter tuning procedures? I understand that tuning hyperparameters of competitor approaches requires extensive resources. However, it would be interesting to know how much work was required to obtain the results in the paper for ReFAT.

I believe that this work offers interesting insights to the ICLR community and recommend accepting it. I am willing to improve my score after the rebuttal.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730409358304

thread_type:
official_review

<thread_object>
creator: reviewer_3mzw

<content>
summary:
Summary: In this paper, the authors demonstrate that adversarial attacks share a universal mechanism for circumventing LLM safeguards. Based on these findings, the authors proposed Refusal Feature Adversarial Training (ReFAT) to efficiently performs LLM adversarial training by simulating the effect of input-level attacks via RFA. Experiment results show that ReFAT significantly improves the robustness of three popular LLMs against a wide range of adversarial attacks

strengths:
+ proposed a new adversarial training method for safeguarding LLMs with improved efficiency
+ experiments show improvements over existing adversarial training solutions

weaknesses:
+ The idea of manipulating the refusal features in the activation space is also introduced in other works on steering vector or activation steering such as the following ([1][2][3]). The authors might also want to discuss and comment on the differences. Specifically, [1][3] applied this type of activation engineering to the jailbreaking tasks and it is also recommended to compare with

[1] Wang, Haoran, and Kai Shu. "Backdoor activation attack: Attack large language models using activation steering for safety-alignment." CIKM2024

[2] Panickssery, Nina, et al. "Steering llama 2 via contrastive activation addition." ACL2024

[3] Cao, Yuanpu, et al. "Personalized Steering of Large Language Models: Versatile Steering Vectors Through Bi-directional Preference Optimization." NeurIPS 2024

+ There are other types of defenses on the safety of LLMs aside from adversarial training. It is also recommended to compare with other existing defense strategies, e.g., the following one [4].

[4] Cao, Bochuan, et al. "Defending against alignment-breaking attacks via robustly aligned llm." ACL 2024.

+ Many new attack baselines are presented in the recent year. The authors might also want to consider comparing with newer attack baselines (e.g. [5][6][7]) to demonstrate the effectiveness of the proposed work 

[5] Mehrotra, Anay, et al. "Tree of attacks: Jailbreaking black-box llms automatically." arXiv preprint arXiv:2312.02119 (2023).

[6] Li, Xirui, et al. "Drattack: Prompt decomposition and reconstruction makes powerful llm jailbreakers." arXiv preprint arXiv:2402.16914 (2024).

[7] Zhang, Tianrong, et al. "WordGame: Efficient & Effective LLM Jailbreak via Simultaneous Obfuscation in Query and Response." arXiv preprint arXiv:2405.14023 (2024).

+ The introduction of RFA and optimal attack and its relationship to the following adversarial training is not very clearly stated. See question below

questions:
- Can you give more details on why including the mean RF activation over harmless prompts term in Eq (3)? What does this term represent and why do we want to add this term to the activation?

- Can you further explain the loss design in Eq (10)? The removal of refusal features ($H(x) − R_{HH}$) is different from what is demonstrated in Sec 3.2 / Sec 3.3?

- The approximation of RFA to the optimal attack still feels a bit less convincing. Can you directly calculate the best $\delta$ direction in Eq (7) and its similarity to RFA? Also you are approximating AA with RFA, directly solving Eq (7) using adversarial training should give you even better performances. Can you show the performance gap here?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730618178153

thread_type:
official_review

<thread_object>
creator: reviewer_sytg

<content>
summary:
The paper addresses the vulnerability of LLMs to adversarial attacks, which can prompt these models to generate harmful, sensitive, or false information. These attacks exploit the models' inability to reliably detect and refuse harmful inputs. The authors highlight a mechanism called the Refusal Feature, which adversarial attacks exploit by ablating this feature to make harmful requests seem benign. Accordingly, the authors suggest a worst-case robust learning framework named ReFAT. During training, ReFAT dynamically computes the RF using two sets of inputs (harmful and harmless) and then ablates the RF for harmful inputs. This simulates the effect of adversarial attacks, training the model to make safety determinations without relying on the most salient features of input maliciousness.

strengths:
1. The suggested method seems to be novel and reasonable. ReFAT dynamically computes the RF using two sets of inputs (harmful and harmless) and then ablates the RF for harmful inputs. This simulates the effect of adversarial attacks, training the model to make safety determinations without relying on the most salient features of input maliciousness. 

2. The exploration of refusal features is interesting. The intervention mechanism seems to be reasonable while more discussion should be added to make it clearer. 

3. The paper provides compelling evidence that ReFAT effectively enhances model robustness across different types of LLMs and attack vectors.

weaknesses:
1. The authors rely on the previous findings of refusal features suggested by Arditi, and discuss the background in Sec 3.1 However, to me, I think the content is not self-contained enough, where I cannot understand the key heuristics behind Eq 3 as well as the following equation for their physical meanings. It seems that Eq 3 takes r_HH to somewhat measure the smoothness of the feature space, transforming h^l into a more informative space for intervention. I am quite curious about hope it is derived and why the third term in the rhs of Eq 3 exists (seems like a bias term?) Personally, I believe the authors should discuss more for this part, especially for their physical meanings, as the following sections of 3.2 - 3.3 critically rely on these formulations. 

2. What are the main difference between r_HH in Sec 3.1 and r_A in Sec 3.2, it seems that they are exactly the same thing. Therefore, the observation of their similar cosine similarity seems to straightforward. 

3. From the words in Sec 4.2, I cannot fully understand how the algorithm workflow is, could the authors provide more details. Moreover, it seems that the suggested worst-case learning method involves a lot of hyper parameters, such as the adopted layer of perturbations, strength of perturbations, I cannot find the specific configuration of these parameters. 

4. What's the connection between the suggested method and many previous alignment methods, especially for those methods that consider the binary feedback setup. Could the authors consider to add more baseline methods of alignments in their experiments.

questions:
Please refer to the Weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730722328873

thread_type:
official_review

<thread_object>
creator: reviewer_xyhb

<content>
summary:
This paper addresses the vulnerability of large language models (LLMs) to adversarial attacks, which can elicit harmful responses by bypassing model safeguards. The paper reveals a shared mechanism among various adversarial attacks: they function by ablating a dimension in the residual stream embedding space, termed the "refusal feature" (RF), which acts as a linear predictor of input harmfulness. This refusal feature, discovered by Arditi et al. (2024), is the mass mean difference in hidden representations between harmful and harmless instructions, and is essential for the model to generate safe responses. Building on this insight, the paper proposes Refusal Feature Adversarial Training (ReFAT), which enables the model to recognize harmful instructions and maintain robustness. Experimental results show that ReFAT enhances the robustness of multiple LLMs against a wide array of adversarial attacks, while reducing computational overhead compared to traditional adversarial training methods.

strengths:
1. The method is innovative, extending Arditi's Refusal Feature and further proving the similarity of adversarial attack and RFA mechanisms through causal theory, and proposing a more efficient adversarial training method based on this.
2. The method proposed in this paper can be objectively evaluated, including success rate of attack, generation performance, efficiency evaluation method, and the limitations are also analyzed.
3. The paper is well organized, first analyzing the general mechanism of adversarial attacks in the semantic space, and then naturally proposing an adversarial training method based on this feature, and proving the feasibility of the method through a large number of experiments and supplementary materials.

weaknesses:
1. Using Llama-3-8B-Instruct-generated XSTest responses as the gold standard answers for supervised fine-tuning might pose potential issues in the design of scientific experiments, especially when Llama-3-8B-Instruct itself is also a subject of the subsequent experiments. This design could violate basic principles concerning independence in experimental standards.
2. The conclusion that rejecting direction leads to a performance degradation cannot be well demonstrated by experiments. The end of Section 3.1 states that experimental results show that performance degrades when the rejection direction is simply set to zero. However, the experimental figures in Appendix D only show that "harmless example features are not centered near zero," which leads to a speculation and intuitive inference that model performance degradation may be a possible consequence of this phenomenon?

questions:
1. In Figure 3 - PCA visualization, why does the AutoDAN section have noticeably fewer experimental samples from HarmBench compared to the other three attack algorithms?
2. Section 3.1 mentions that Harmful samples were sampled from AdvBench, which to my knowledge is a dataset for adversarial suffix-type attacks. GCG is also a suffix attack. Intuitively, HarmBench (adversarial) should be quite similar to Harmful since they are both suffix attack samples. However, the actual results do not reflect this. Please explain or analyze this phenomenon.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732132786201

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Our response to all reviewers

comment:
We thank the reviewers for their comments and constructive criticism. We are happy to read that all of them find the proposed method novel and efficient, and concur that our experiments clearly demonstrate improvements over popular recent adversarial training methods. We are also delighted to see that most reviewers took notice that the ReFAT method stems from an analysis of the effect that known adversarial attacks have on model representations, rather than simply from trying out an idea.

The reviewers’ comments helped us realize the shortcomings of the current manuscript. We uploaded a revised version where we highlighted in **orange** the differences from the initial submission. We would appreciate if the reviewers could check our revised paper.

In Summary:
* We added additional related work as suggested by the reviewers, especially those pointed out by reviewer 3MZw. 
* We rewrote large parts of Section 3.1 (background for the ReFAT method) to clarify it and make it more self-contained.
* We added pseudo-code for ReFAT as Algorithm 1.
* We conducted additional experiments on over-refusal on MMLU as Table 6.

Moreover, reviewers also suggested potential additional experiments. While it is not possible to complete such experiments in the timeframe of the discussion, we do intend to add at least some of them to the CR version, if the paper is accepted for publication. In particular, we should be able to add gradient-based continuous latent attacks such as LAT (Sheshadri et al., 2024) as an additional baseline defense, as suggested by reviewer 3MZw.

Finally, we thank you again for your thoughtful reviews. We hope we have addressed your concerns, and hope that our response would make you feel comfortable to raise your scores and give this work the opportunity to be presented at ICLR. We remain of course open to more discussion in the coming days.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732560939128

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Our response to all reviewers

comment:
We thank you all again for your thoughtful reviews, which have helped us significantly improve the paper. We are encouraged that now three reviewers out of four are supportive of accepting our work!
 
The fourth reviewer is of different advice, on the ground that the combination of refusal feature ablation and adversarial training would be trivial, and comparison against some recent baselines missing. We acknowledge their objections, but would like to emphasize a few points:
 
* Our work is not simply a combination of Refusal Feature Ablation (RFA) and adversarial training. Besides the sheer leaderboard-climbing aspect, a key contribution is the finding that adversarial attacks of very different natures all work by ablating the refusal feature. We do think that this is a non-trivial result.
* Building on this insight, we designed ReFAT to target this shared vulnerability, enabling the model to robustly counteract diverse attacks. Importantly, our results demonstrate that ReFAT is effective against attacks that are fundamentally different in nature, underscoring the broad applicability of our approach.
* Compared to inference-time representation-editing mitigations, the step to adversarial training extends the scope of the defense to cases when the attacker has access to model weights (and can therefore just choose not to intervene at inference time), but does not have the resources or the competences to alter them. With powerful open-weight models becoming commonplace, we believe this is important.
* The efficiency of ReFAT compared to all alternative adversarial training methods, which is not being questioned, can make a substantial difference to practitioners with limited access to GPUs.
* We had finite computational resources and had to arbitrate between the number of models, adversarial attacks, general performance benchmarks, over-refusal benchmarks, and baseline mitigation strategies. We may be able to add a limited number of experimental conditions in time for an eventual CR version, but we believe we struck a reasonable balance, and we explained in our replies to this reviewer the factors driving our selection of baselines.

At a higher level, the field of mechanistic interpretability is sometimes criticized for not yielding actionable insights in improving LLMs: we believe that this work to be evidence of the opposite, and we expect our work could inspire others to develop more interpretable and efficient methods to improve LLMs in other domains.

We hope that the reviewing panel will strengthen their support, and the program committee will give this work the opportunity to be presented at ICLR. Thank you.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces a defense method for protecting LLMs against jailbreaking, based on the "refusal feature" observation of Arditi et al (2024). It first shows that several adversarial attacks ablate the same refusal feature direction and based on this observation design an efficient approximation to adversarial training, ReFAT. In experiments on 3 different LLM models, ReFAT is shown to be competitive with baselines while being significantly more efficient.

Strengths of the paper include the novel and well-motivated design of the ReFAT algorithm, including the insight that refusal feature ablation is a somewhat universal adversarial attack mechanism. ReFAT is relatively simple to implement and computationally efficient. Experimental results are promising.

Weaknesses are a lack of discussion and comparison to other recent works (which can also affect the generalizability of the results) and a more thorough evaluation of the approximation quality of the refusal feature attack vs true adversarial training. 

Overall the work presents an interesting insight along with a well-motivated, efficient yet effective defense mechanism for LLMs. This is an active and important area of research and the AC believes this paper will add to the conversation; the AC thus recommends acceptance. The authors should include the additional comparisons in the final version of the paper as promised in their responses.