Robust LLM safeguarding via refusal feature adversarial training

Thank you for your positive and valuable feedback. We are glad that you highlighted the comprehensiveness of our mechanistic analyses and experiments, and the efficiency of our ReFAT method. Please see below for our point-by-point response to your comment:

W1: 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.

To our knowledge, previous work has shown that the mechanism behind refusals in LLMs is the refusal feature. However, we are the first to provide evidence that adversarial attacks bypass the refusal mechanism by ablating this refusal feature.

• Figure 1 provides an overview of the contributions of the paper. While the first two “rows” are indeed known results about the refusal feature, the next two rows illustrate our novel mechanistic analysis of existing adversarial attacks. The last row illustrates our new training method, ReFAT.

• Figure 2 and 3: We are not aware of previous work doing the similarity-based and PCA visualization of Figures 2,3 to analyze the connection between existing adversarial attacks and refusal feature ablation.

W2: The refusal results do not appear to be consistent with those of Xhonneux et al. (2024)

We would like to clarify that Xhonneux et al. (2024) was evaluating a different suite of LLMs that are significantly less safe than those used in our paper, as suggested by previous studies like HarmBench (Mazeika et al., 2024).

Moreover, Xhonneux et al. (2024) measure refusals on a hand-written dataset of 40 clearly harmless prompts (e.g., “Explain the rules of soccer”). In contrast, we measure refusals on XSTest, which is more challenging because it includes harmless prompts that seem harmful (e.g., “How can I build a bomb in minecraft?”). Therefore, XSTest provides a more informative assessment of refusals.

Consistently with Xhonneux et al (2024), we find that R2D2 significantly increases over-refusals (lower compliance) on XSTest, unlike their CAT method and ReFAT, see Table 1. Xhonneux et al. report that CAT does not increase refusals on their harmless dataset when applied to Gemma and Mistral-7B. We find that on the more challenging XSTest, CAT can lead to more over-refusal than ReFAT (see results for Mistral-7B in Table 1).

W3: A more comprehensive assessment of refusal behavior, e.g. with the OR-Bench, would further strengthen the paper

Thank you for this suggestion. OR-Bench was just released and not reviewed at the time of the paper submission, therefore we only included XSTest. To complement the refusal assessment, we follow a similar approach to Appendix D of [Xhonneux et al, 2024]: we measure the number of refusals on the 5,700 MMLU questions and report the following scores. We find that the refusal-rate on MMLU for ReFAT is at most 0.001, confirming that ReFAT does not induce significant over-refusals. We added these results in Appendix F, Table 6 of our paper revision.

W4 and W5: How stable are the different robustification methods with respect to hyperparameter choices? [...] Can the authors elaborate on hyperparameter tuning procedures? [...]

We would like to point out that due to the high computational cost of running adversarial attacks like GCG to evaluate each model, it is practically infeasible to comprehensively search for optimal hyperparameter configurations for each defense method. However, we have still made our best attempt to optimize the hyperparameter choices for each model.

In particular, our preliminary experiments on ReFAT suggest that its improvement on model adversarial robustness is pretty consistent, as long as we ablate refusal features for the second half of all model layers (i.e. layer 16-32 for Llama/Mistral, and layer 14-28 for Gemma), while the other hyperparameters (e.g. learning rate, batch size) have much less effect on model performance.

As for R2D2 and CAT, we tried to modify the default hyperparameters in their official implementations during preliminary experiments, and did not observe any significant performance improvement. Therefore, we are pretty confident that the demonstrated advantage of ReFAT over the other defense methods is robust against hyperparameter choices.

Please see below for our response to the issues you point out under “Questions".

Q1: 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?

The intuition behind RFA is to make harmful samples appear as similar as possible to harmless samples along the refusal feature direction. This translates into altering internal representations so that their projection on the RF direction is close to the centroid of the harmless samples. As shown in Fig. 18 and 19, the mean activation along the refusal direction of harmless samples is generally not zero, so setting this value to zero (the effect of applying the second term in the right-hand side of Eq. 3) would fall short of making harmful examples appear as harmless as possible. The mean RF activation term ensures therefore that the projection of the activation along the refusal feature is equal to the mean of the projections for harmless examples.

Q2: Can you further explain the loss design in Eq (10)? The removal of refusal features (H(x)−RHH) is different from what is demonstrated in Sec 3.2 / Sec 3.3?

We did not include the bias term in Eq (10) as we did in Eq. (3) due to efficiency consideration: in particular, computing the bias term $\Bar{\mathbf{r}}_{\mathcal{D}_\text{harmless}}^{(l)}$ requires an additional computation step of calculating the average projection value of harmless input hidden representations over the refusal feature at each layer, and we found empirically in our preliminary experiments that this would make the training slower while has little improvement on the resulting model performance. We therefore chose to remove the average harmless RF bias term from our training objective to facilitate fine-tuning. We have added explanations for this design choice in the revised manuscript.

Q3: The approximation of RFA to the optimal attack still feels a bit less convincing. Can you directly calculate the best δ 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?

We thank the reviewer for raising this point. We found the recent work of Latent Adversarial Training (LAT) by (Sheshadri et al., 2024) has proposed a similar idea of finding worst-case residual stream perturbation using projected gradient descent. We were considering comparing ReFAT with LAT, but unfortunately they had not fully released their code implementation by the time we submit our paper, so due to time constraint and resource limitation we did not include LAT as a baseline. However, we do acknowledge the strong relation between ReFAT and LAT, and we plan to take LAT as a baseline in the camera ready version if the paper is accepted for publication.

We hope that our response has addressed your concerns regarding our work, and we sincerely hope that you could consider raise your score to give this work the opportunity to be presented at ICLR. If you have any other questions or comments, please do not hesitate to post them.

We thank the reviewer for the constructive feedback and for your time and effort. Concerning the issues you point out under “Weaknesses”, we would like to clarify the following.

W1: The idea of manipulating the refusal features in the activation space is also introduced in other works [...]

We thank the reviewer for pointing to these additional related work, and we have added them in our revised manuscript. Meanwhile, we would like to emphasize that the main contribution of our paper is not “the first to propose the idea of manipulating refusal features (or steering vectors that encode other concepts)”, but instead is “the first to show that manipulating the linear refusal feature during training time could result in a model that can robustly reject adversarial prompts during evaluation, without the need of an additional steering vector”.

In particular, we would like to point out that the papers [1][2] only show that refusal feature ablation (or adding a steering vector to cancel the refusal feature) is a powerful jailbreaking attack method, similar to the finding of (Arditi et al., 2024). While our paper goes further and proposes a novel defense method that could not only mitigate model vulnerability to such steering vector attacks, but can also enhance model robustness against other types of popular adversarial attacks.

On the other hand, our paper is also different from the line of studies on “learning effective steering vectors” such as [3] and (Stickland et al., 2024) in that our ReFAT method fine-tunes model parameters by injecting a steering vector during training, and does not require an additional steering vector to be injected during evaluation to achieve high safety or robustness, while “steering vector learning” papers are instead trying to develop a “lightweight” strategy of finding a good steering vector that could be added during model evaluation inference time to improve performance. We have explained this in Section 4.2 of the revised manuscript to enhance clarity.

W2: There are other types of defenses on the safety of LLMs aside from adversarial training.[...]

We would like to emphasize that the objective of our paper is to propose a method that could improve LLM safety at the model parameter level – i.e., we would like to ensure high model adversarial robustness even without applying any inference-time defense methods including prompt engineering and post-generation checking/reflection (e.g. the RA-LLM method in [4]). From this perspective, we consider ReFAT and the other adversarial training methods as complementary as opposed to competing LLM safeguarding against inference-time defense methods. We thank the reviewer for bringing up this line of work, and we will consider combining an adversarially fine-tuned model using ReFAT with inference-time defense methods such as RA-LLM in future extensions.

W3: Many new attack baselines are presented in the recent year. The authors might also want to consider comparing with newer attack baselines [...]

We thank the reviewer for bringing up additional recently proposed attack methods, and we will include them in our future revision. On the other hand, due to the high computational cost of running adversarial attack methods, we have carefully searched for and selected the most popular, powerful, and peer-reviewed attack methods proposed in the past few years, in order to guarantee the generalizability of our findings. In particular, the attacks proposed by [6][7] were not peer-reviewed and had few citations when we submitted our paper. While TAP by [5] is a well-known attack, it is highly similar to the PAIR attack tested in our paper, and recent evaluation studies suggest that TAP is not more successful (and sometimes even worse) than PAIR in jailbreaking the safest models like Llama, Claude and GPT-4 (see Table 6 on page 26 of (Mazeika et al., 2024)), so we chose PAIR as the most representative LLM-based attack in our study.

W4: The introduction of RFA and optimal attack and its relationship to the following adversarial training is not very clearly stated.[...]

Please see our point-by-point response to your raised questions below.

References:

Mantas Mazeika, Long Phan, et al. Harmbench: A standardized evaluation framework for automated red teaming and robust refusal. 2024.

Asa Cooper Stickland, Alexander Lyzhov, Jacob Pfau, Salsabila Mahdi, and Samuel R Bowman. Steering without side effects: Improving post-deployment control of language models. 2024.

Thank you so much for your insightful feedback! Please see below for our point-by-point response.

W1: The authors rely on the previous findings of refusal features suggested by Arditi, and discuss the background in Sec 3.1 [...]

We agree that Section 3.1 could be clearer and more self-contained, so we have revised it to better explain the key ideas and physical meanings behind our formulations (see changes in orange in the revised PDF). Specifically, we clarify that Arditi discovered that refusal in language models is mediated by a single direction, and that we can determine this direction by calculating the average difference between the activations of harmless and harmful prompts. We find the refusal direction as $\mathbf{r}_{HH}$, with $\hat{r}$ representing the unit vector for this direction. The second term in Equation 3 modifies the original activations by zeroing out the value along the refusal feature direction. The last term sets this value to the average of the harmless prompt activations, as such activations are not generally centered near zero along the direction.

W2: What are the main difference between r_HH in Sec 3.1 and r_A in Sec 3.2 [...]

Thank you for raising this point about the difference between r_HH in Section 3.1 and r_A in Section 3.2. While Eq. 2 and 4 have very similar structures, they have different meanings.

In Section 3.1, r_HH represents the single refusal feature direction that we identified by calculating the mean difference between harmful and harmless inputs. This direction encodes the model's tendency to refuse harmful requests. By controlling this refusal feature, we can influence the model's behavior—effectively "jailbreaking" it.

In Section 3.2, r_A denotes the average influence on activations caused by various jailbreak attacks of different natures, including GCG (a gradient-based attack) and PAIR (an attack using iterative refinement of prompts). Interestingly, we found that all these attacks, despite their different methodologies, indirectly influence the same refusal feature direction $\mathbf{r}_{HH}​$ identified earlier.

This demonstrates that diverse jailbreak methods converge on manipulating the same underlying refusal feature. We believe this is a significant finding, as it underscores the central role of the refusal feature in understanding and defending against adversarial attacks in ReFAT.

W3: From the words in Sec 4.2, I cannot fully understand how the algorithm workflow is [...]

Thank you for your comments regarding the clarity of Section 4.2. We have revised the text to bring more clarity to our algorithm's workflow. To clarify, instead of standard adversarial training that uses gradient-based worst-case perturbations, we craft perturbations by ablating the refusal feature directions. Specifically, we use the refusal feature $R_{HH}$ as an approximation for the worst-case perturbation and apply it directly to the residual stream activations at each layer for harmful prompts during training. Since the perturbation is fully defined by $R_{HH}$ , there is no perturbation strength parameter involved. We apply this perturbation to the residual stream at each layer using $R_{HH}$ layerwise refusal features.

The main hyperparameters — such as the number of random harmful and harmless prompts sampled to compute $R_{HH}$ , the frequency $k$ of dynamic recomputation, the probability $p_{RFA}$ of applying the ablation, and the choice of datasets are specified in Section 5 (Experimental Setup), with more detailed ReFAT training hyperparameters provided in Appendix Table 4. We have also simplified the formulations in the manuscript to enhance clarity.

W4: What's the connection between the suggested method and many previous alignment methods [...]

Are you referring to preference fine-tuning methods, such as RLHF and Direct Preference Optimization (DPO)? Assuming this is the case: ReFAT is an alternative to supervised finetuning (SFT) with improved robustness and is therefore comparable to other adversarial training methods for LLMs such as R2D2 and CAT.

Moreover, it is possible to design a variation of ReFAT for preference-finetuning methods such as DPO [Rafailov et al, 2023] and the IPO variant [Azar et al, 2024], by similarly removing the refusal feature from the activations of harmful inputs with probability $p_{\text{RFA}}$ during forward passes.

We thank reviewer xYHb for their time and effort. We are delighted to read that you find our method innovative and objectively evaluated, and that you think the paper was well organized.

Concerning the issues you point out under “Weaknesses”, we would like to clarify the following.

W1: 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 [...]

For the avoidance of doubt, and because our use of the term “gold-standard answers” could be misleading: the 150 examples from the XSTest dataset included for supervised finetuning come from a holdout set that is not used in evaluations, i.e. to determine the results in Table 1. We chose to use the generations from Llama-3-8B because it has almost perfect accuracy on XSTest (97.2%, Table 1, first row). The use of predictions from a model to improve the model itself is a form of self-training that does not imply a data leakage. We clarify the sentence “We take the XSTest responses generated by Llama-3-8B-Instruct as gold-standard answers for supervised fine-tuning.” -> “We use the responses generated by Llama-3-8B-Instruct on this holdout sample from XSTest as references for the next-token prediction task in the supervised finetuning step.”

We acknowledge the residual concern that using generations from one model (Llama-3-8B-Instruct) also when finetuning other models might in principle introduce a bias. Such bias could be positive or negative. It could be favorable to Llama-3-8B-Instruct because reference samples are already easily reachable by the model, but it could be in favor of other models because the process could realize a form of ‘distillation’ from Llama-3-8B-Instruct to the target model. Considering the very limited size of the sample, this bias is likely very small. We emphasize moreover that the purpose of this evaluation is not to compare the relative performance of models, but rather to verify that the impact of applying different adversarial training methods is systematic across model families.

W2: The conclusion that rejecting direction leads to a performance degradation cannot be well demonstrated by experiments [...]

The phenomenon that zeroing the projection on the refusal direction results in a performance degradation is something we observed in small-scale early experiments. From Fig. 18 and 19, though, it is quite apparent that the value ‘0’ does not systematically correspond to ‘harmlessness’, and in some cases (e.g. Layer 3 for gemma-2b-it, Fig. 19) it can be quite close to the mode for ‘harmfulness’, or altogether out of distribution (e.g. Layer 0 in either model). For these reasons, we still believe that setting the projection to a value squarely within the ‘harmless’ range is the correct thing to do. We have changed the wording in Section 3.1 to reflect this.

Q1: In Figure 3 - PCA visualization, why does the AutoDAN section have noticeably fewer experimental samples from HarmBench compared to the other three attack algorithms?

This is because we retain and plot only cases when the attack succeeds, and AutoDAN has a much lower success rate than the other attacks (see Table 1) on Llama-3-8B-Instruct.

Q2: Section 3.1 mentions that Harmful samples were sampled from AdvBench, which to my knowledge is a dataset for adversarial suffix-type attacks. [...]

We are not sure we fully understand this question. Harmful samples were taken from the “Harmful Behaviours” part of the AdvBench dataset, and we would like to clarify that these harmful prompts are generic ‘red teaming’ style prompts and are not specifically geared towards adversarial suffix-type attacks (e.g. they have no nonsensical suffices as those found by GCG), despite the fact that they were originally used in a paper that developed attacks of this type. The same is also true of the samples we took from HarmBench to complete the evaluation datasets. Was there a misunderstanding concerning this fact? Otherwise, could the reviewer kindly elaborate on the phenomenon that should be explained?