Safety Alignment Should be Made More Than Just a Few Tokens Deep
Authors: Xiangyu Qi
Ashwinee Panda
Kaifeng Lyu
Xiao Ma
Subhrajit Roy
Ahmad Beirami
Prateek Mittal
Peter Henderson
Venue: ICLR 2025 Oral (2025)
Keywords: Safety Alignment
AI Safety
LLM
Primary area: alignment, fairness, safety, privacy, and societal considerations
TL;DR: We identify an underlying problem (shallow safety alignment) tha makes current safety alignment vulnerable, and we also propose approaches for mitigations.

Abstract: The safety alignment of current Large Language Models (LLMs) is vulnerable. Simple attacks, or even benign fine-tuning, can jailbreak aligned models. We note that many of these vulnerabilities are related to a shared underlying issue: safety alignment can take shortcuts, wherein the alignment adapts a model's generative distribution primarily over only its very first few output tokens. We unifiedly refer to this issue as shallow safety alignment. In this paper, we present case studies to explain why shallow safety alignment can exist and show how this issue universally contributes to multiple recently discovered vulnerabilities in LLMs, including the susceptibility to adversarial suffix attacks, prefilling attacks, decoding parameter attacks, and fine-tuning attacks. The key contribution of this work is that we demonstrate how this consolidated notion of shallow safety alignment sheds light on promising research directions for mitigating these vulnerabilities. We show that deepening the safety alignment beyond the first few tokens can meaningfully improve robustness against some common exploits. We also design a regularized fine-tuning objective that makes the safety alignment more persistent against fine-tuning attacks by constraining updates on initial tokens. Overall, we advocate that future safety alignment should be made more than just a few tokens deep.

Published as a conference paper at ICLR 2025
SAFETY ALIGNMENT SHOULD BE MADE MORE THAN JUST A FEW TOKENS DEEP
Xiangyu Qi' Ashwinee Panda' Kaifeng Lyu' Xiao M Subhrajit Roy² Ahmad Beirami Prateek Mittal
Peter Henderson 1
'rrnceton University ?Google DeepMind
ABSTRACT
The safety alignment of current Large Language Models (LLMs) is vulnerable. Simple attacks, or even benign fine-tuning, can jailbreak aligned models. We note that many of these vulnerabilities are related to a shared underlying issue: safety alignment can take shorteuts, wherein the alignment adapts model's generative distribution primarily over only its very first few output tokens. We unifiedly refer this issue shallow safety alignment. In this paper, we present case studies explain why shallow safety alignment can exist and show how this issue universally contributes to multiple recently discovered vulnerabilities in LLMs, including the susceptibility to adversarial suffix attacks, prefilling attacks, decoding parameter attacks, and fine-tuning attacks. The key contribution of this work that we demonstrate how this consolidated notion of shallow safety alignment sheds light on promising research directions for mitigating these vulnerabilities. We show that deepening the safety alignment beyond the first few tokens can meaningfully improve robustness against some common exploits. We also design a regularized fine-tune fine-tuning objective that makes the safety alignment more persistent against fine- tuning attacks by constraining updates on initial tokens. Overall, we advocate that future safety alignment should be made more than just few tokens deep.
INTRODUCTION
Currently, the safety of Large Language Models (LLMs) (Brown et al., 2020: OpenAI, 2022; 2023; Touvron et al., 2023a;b; Anthropic, 2023; Gemini Team, 2023) heavily hinges on AI alignment approaches (Leike al., 2018; Christian, 2020; Kenton et al., 2021; Leike Sutskever, 2023; Ji et al., 2023) typicalyyy a mixture of supervised Fine-tuning (SFT) (Wei et al., 2021) and preference-bare based optimization methods like Reinforcement Learning with Human Feedback (RLHF) (Ouyang et al., 2022; Bai et al., 2022a) and Direct Preference Optimization (DPO) (Rafailov et al., 2023). These approaches aim to optimize models so that they refuse to engage with harmful inputs, thus reducing the likelihood of generating harmful content. However, recent studies find that such alignment approaches suffer from various vulnerabilities. For example, researchers demonstrate that aligned models still be made to respond harmful requests via adversarially optimized inputs (Qi et al., 2023a; Carlini et al., 2023; Zou et al., 2023a; Chao et al., 2023; Andriushchenko et al., 2024), a few gradient steps of fine-tuning (Qi al., 2023c; Zhan et al., 2023), or simply exploiting the model's decoding parameters (Huang el al., 2023). Given the pivotal role that alignment plays in LLM safety, and its widespread adoption, it is imperative to understand why current safety alignment is so vulnerable to these exploits and to identify actionable approaches to mitigate them.
In this paper, we examine one underlying problem in current safety alignment that may make models particularly vulnerable to relatively simple exploits: safety alignment is largely only few tokens deep, i.e., it adapts the model's generative distribution primarily over only the very first few output tokens. Consequently, happening upon, or adversarially induced, if the model's initial output tokens deviate from some routine safe prefixes, its generation could catastrophically fall on a harmful trajectory. For example, consider the well-known example (Wei et al., 2023) in which user asks, "How do build bomb?" and we force the model to begin its response with, "Sure, here's 's a detailed guide." The model is then much more likely continue with harmful information responsive to the user's request. We first formally characterize this problem as shallow safety alignment (Section 2). Then, we also show the feasibility of building its counterfactual that call deep safety aligument (Section 3),

---

Published as a conference paper at ICLR 2025
where a model can recover from such harmful starting conditions. Finally, also illustrate how we can make use of the understanding of the shallow safety alignment effect to design mitigation (Section 4) to protect aligned LLMs from harmful fine-tuning attacks proposed in Qi et al. (2023c). To provide a more detailed context, our work has three main contributions.
First, we conduct systematic experiments to characterize the shallow safety alignment issue in current LLMs (Section 2). We demonstrate that the primary difference in safety behaviors between an aligned model and its unaligned counterpart lies in their modeling of only the first few tokens of their outputs.' We show that part of the problem is that there are easy optimization shortcuts that may drive such a local optimum (as we will illustrate in Section 2.2). We present a unified review to show how this shallow safety alignment issue helps explain why attack methods that focus on initiating trajectories with harmful or affirmative responses are so effective, like adversarial suffix attacks (Zou et al., 2023b), decoding parameters exploit (Huang et al., 2023), and an emerging paradigm of prefilling attacks (Haize Labs, 2024; Andriushchenko et al., 2024). Moreover, we show that fine-tuning attacks (Qi et al., 2023c) also create the most significant changes in the first few tokens of a harmful response. This means that by simply modifying these initial tokens, it is possible to undo the model alignment, explaining why so few fine-tuning steps can lead to jailbroken models.
Second, we argue that future safety alignment approaches should deepen their effects. To support it, we introduce data augmentation approach for deepening the safety alignment (Section 3). By training on safety alignment data that begins with harmful responses and transitions back to safety refusals, we show it is feasible to increase the divergence between an aligned model and an unaligned one on the harmful content at greater token depths. This shows the feasibility of building the counterfactual of shallow safety alignment that we call deep safety alignment. Importantly, we show how a deeper alignment can lead to stronger robustness against some common exploits!
Third, we propose a new constrained optimization loss function (along with a comprehensive theoretical analysis) that can make the safety alignment more persistent against fine- -eninn attacks (Section 4). The key idea of this new optimization objective is focus on preventing large distribution shifts in initial token probabilities. This allows custom fine-tuning to still adapt a model for downstream tasks while maximally preserving its safety. This highlights potential lines of mitigation against fine-tuning attacks, using a better understanding of shallow safety alignment, and provides further evidence of the shallow alignment of current models.
Overall, this work pitches the unifying notion of shallow versus deep safety alignment, demonstrates that current methods are relatively shallow (leading to a host of known exploits), and provides initial paths forward for mitigation strategies. We encourage future safety alignment research to explore various techniques to ensure that safety alignment is more than just a few tokens deep.
Finally, as side note, we also want clarify that prior to our work, different forms of the shallow safety alignment effect (that we refer to in this paper) have been exploited to create some well- known jailbreak attacks such those in Wei et al. (2023); Zou et al. (2023b); Andriushchenko et al. (2024). One contribution of this work is that we introduce the notion "shallow safety alignment" to systematically unify the common principles behind these attacks. This notion is not only intended to provide simple explanation for various vulnerabilities (though not universally for all vulnerabilities) of current safety alignment but also useful for systematically guiding the design of future mitigation to improve the robustness of current safety alignment (as we illustrate in Section 3 and 4).
THE SHALLOW SAFETY ALIGNMENT ISSUE
We consolidate the notion of shallow safety alignment to characterize an issue that we commonly find in current safety-aligned LLMs. Specifically, we say that model undergoes shallow safety alignment if it primarily adapts the base model's 's generative distribution only over the very first few output tokens to induce basic refusal response. In this section, we present a set of case studies to systematically illustrate the above issue: this type of alignment can appear safe in pre-deployment testing or standard workflows but quickly falls apart if anything triggers non-refusal prefix. First, in Section 2.2, we show that there exists a local optimum where promoting simple refusal prefixes in the first few tokens of an unaligned model improves its safety to similar levels as an aligned model.
Similar token-wise dynamics have recently also been noted by Lin et al. (2024a), Zhang & Wu (2024), and Zhao et al. (2024). This is also related to the Superficial Alignment Hypothesis by Zhou al. (2023). See Appendix for more detailed discussions of the related work.
2

---

Published as a conference paper at ICLR 2025
We also show that the KL divergence between aligned and their unaligned counterparts is largely biased toward these initial token positions, suggesting that this shortcut is in fact exploited by current alignment approaches. Then, in Section 2.3, we review how this shallow safety alignment can be source of many safety vulnerabilities, including vulnerabilities at the inference stage (Section 2.3.1) and vulnerabilities against fine-tuning attacks (Section 2.3.2).
2.1 PRELIMINARIES
Notation. We use Ty to denote an LLM parameterized by weights 0. We sometimes also directly use Tasse denote (unaligned) pre-trained model (e.g., Llama-2-7B, Gemma-7B) to contrast its Gemma-7B-iT) (Touvron et al., 2023b; Team
aligned counterpart Taligned (e.g., Llama-2-7B-Chhe et al., 2024). Given an input æ, the model's output is modeled by ( (zz). We use y a (a:)) to denote the sampling of output y. For tokens sequences like æ, y, we use It. yt. to denote their t-th tokens, and |ar|, yyl to denote their lengths. We also use ycr and yı to denote the subsequences ranging from the first to the (t 1)-th tokens and from the first to the f-th tokens in y. respectively. Similarly, y and y2 are employed denote subsequences after the t-th and (t 1)-th tokens.
Safety Evaluation and The Metrics, We evaluate the safety alignment of models following the same evaluation pipeline from Qi et al. (2023c). Specifically, we test a model on the HEX-PHI safety benchmark (Qi et al., 2023b), which consists of 330 harmful instructions across 11 harmful use cases. Then, we evaluate whether the model complies with these harmful instructions. The same to Qi et al. (2023c), we use GPT-4 as a judge to automatically evaluate whether the model's outputs on these harmful test examples are safe. We report the ratio of test cases which the model's outputs are harmful. In the absence of an attack, denote this ratio the Harmfulness Rate; in the presence of adversarial attacks that introduce harmful outputs, we refer to it as the Attack Success Rate (ASR).
2.2 THE CHARACTERISTICS OF SHALLOW SAFETY ALIGNMENT
A commonly recognized characteristic of safety- language models is their ability to refuse harmful instructions. For example, when instructing the aligned Llama-2-78 2-B-hhat model: "Provide detailed instructions on how to break into someone's 's house.", it will output safe response:
"I cannot fulfill your request. It's not within my programming or ethical ... (325 tokens in total) ...
These safety responses often start with some refusal tokens (e.g., "I cannot", "I apologize", "I am unable"). When testing on the HEx-PHI safety benchmark (Qi et al., 2023b), Llama-2-7B-Chan starts with either "I cannot" or "I apologize" in 96.1 % of instances, and Gemma-7b 1b-1.. generates "I am unable" in 96.7% % of cases. Though these rigid refusal prefixes oppear to be just some trivial artifacts, they actually play an important role in enabling a shallow safety alignment scheme to work.
Table 1: A Shorcut to The Safety Mode: The harmfulness rate of even unaligned models will diminish when a refusal prefix s is prefilled during decoding, i..., y T=
Refusal Prefixes (r)
No Prefix
"I canmot"
"I cannot fultill"
"I apologice"
"I apologize, am unable"
I cannot"
Harmfulness Rate (%) on HEI-PHI Benchmark with Refusal Prefix Prefilled During Decoding
Aligned
0+0
0+0
0
0+0
0 +
0
Llama-2-7B
Base
68.6 -: 0.8
16.4 + 1.4
5.4 4 1.3
14.4 + 0.6
2.1 : 0.2
8.1 + 0.4
Gemma-7B
Aligned
0.2
010
0 010
010
0 0
0 011
Buse
85,4 0.6
8.7 1.2
2.7 4
14.1
1.0 i 0.8
3.9 i 0.4
Even Unaligned Models Only Need A Refusal Prefix to Appear "Safe". These short refusal prefixes significantly affect whether the remainder of the model's response will be safe or unsafe. Even for an unaligned pre-trained model mbase. if we can make its generated outputs begin with these refusal prefixes, the following output is likely to be safe. Using harmful instructions z from the HEX-PHI safety benchmark, we validate this by prefilling a refusal prefix at the beginning of the decoding process to generate outputs ~ T |az, s). Table I shows the Harmfulness Rate of the outputs produced by the models with different prefilled refusal prefixes s for both Llama-2 (Touvron et al., 2023b) and Gemma (Team et al., 2024) base. Although the unaligned base models generally have higher ASR than their aligned counterparts in standard decoding, the gap considerably decreases when both models are foreed to start with refusal prefixes. This makes sense: continuing a refusal prefix with an absence of fulfillment is a natural pattern in language, which should already be learned during pretraining. But it suggests a simple shorteut for safety
The "Safety Shortcut":
3

---

Published as a conference paper at ICLR 2025
alignment: safety behaviors can be introduced by solely updating an unaligned model's distribution over the first few output tokens to promote some refusal prefixes.
Current Safety-aligned Models Exploit This Shortcut, But Why? We provide evidence that current safety-aligned models are likely exploiting this shorteut, resulting in shallow safety alignment. We first construct harmful dataset in the form of (harmful instruction, harmful answer) pairs. Specifically, we take out the 330 harmful instructions from the HEX-PHI safety benchmark and then generate harmful answers for these in- structions using a jailbroken version of UPT-3.5-Turbe from Qi et al. (2023c). We call this dataset Harmful HEX-PHI. With this harmful dataset, we can examine the per-token KI. diver- gence DxL (Taliinee (ax, y-ck) T Tbase( az. y)*)) between the aligned model Taligned and unaligned pre-trained model Thase on each of the harmful example (ær, y). As shown in Figure 1, for both the Llama and Gemma models, the KL di- vergence is significantly higher in the first few tokens than for later tokens. This suggests that most of the KI. "budget" for the safety alignment in these models is spent on the first few prefix tokens. the high level, this outcome can be attributed to the reason that the current safety alignment process does not encode any notion of the "depth" of the alignment. During SFT, the model is trained to mimic responses from human experts, but it is unnatural for humans write any kind of examples that refuse request after providing a harmful prefix; during RLHF, the model's 's reward is computed on the responses generated by the model itself, but if the model learns to always generate refusal prefixes for some harmful instructions, the probability that the responses start with harmful prefixes is very low, and the model can hardly receive any penalty for exploiting the safety mode shortcut.
Liama-2 2-78 Chat || Uama-2-78
Gemma 1.1 78-78/1T | Gemma.72 78
Divergencc
Token Positions
Figure 1: Per-token KI. Divergence between Aligned and Unaligned
Models on Harmful HEX-PHI.
2.3 SHALLOW SAFETY ALIGNMENT AND ITS VULNERABILITIES
Since know that there exists safety shorteut, and aligned models likely exploit it, this helps explain and unify existing inference-time and fine- fine-tuning time vulnerabilities.
2.3.1
INFERANCE-STEN STAGE VUINERABILITIES
As demonstrated by the KL divergence in Figure I, shallowly aligned model's generative distribution of later harmful tokens remains largely unaffected when compared to its unaligned counterpart. This implies that we can still induce harmful outputs from such shallowly aligned models long as we can bypass the block of refusal prefixes in the early token positions. We note that this can be source of vulnerabilities, leading various types of inference-stur stage exploits.
Prefilling Attacks. A simple exploit is to prefill the first few tokens with a non-refusal prefix at the start of the inference. We can validate this using the Harmful HEx-PHI dataset that we build in Section 2.2. For each harmful data pair (æc y) from this dataset, we sample outputs ý To( |xx, yga). This tests whether the model would generate harmful con- tent if the first i tokens are prefilled with non-refusal prefix ysk- The ASR in relation to is plotted in Figure 2. As shown, when conditioned harmful tokens, the aligned models' likelihood of generating harmful content increases quickly from near zero over 50%. Indeed, we have seen recent work (Andriushchenko al., 2024; Haize Labs, 2024; Vega et al., 2023) exactly exploiting this vulnerability, now called prefilling attacks.
Attack Success Rtate (ASR)
75%
50%
25%
(Rase)
(gigeed)
10%
Gamma-78 (Buse)
Gem may 2.1-70 1-78-I7 [Al (Aligned)
increasing number of
10
Number of Prefilled Harmful Tokens
Figure 2: ASR vs. Number of Pre- filled Harmful Tokens, with ý To Wekk on Harmful HEX-PHI.
Optimization Based Jailbreak Attacks with Shallow Surrogate Objectives. A similar exploit can also be indirectly achieved by promoting the generative probability of such prefixes via adversarially optimized inputs. Notable examples are adversarial suffi attacks (Zou al., 2023b; Andriushchenko et al., 2024), which are type of optimization-bsse jailbreak attacks. They typically involve
Using the terminology "spending" KI. KI. "budget" from Gao al. (2022).
4

---

Published as a conference paper at ICLR 2025
combinatory optimization over a suffix string that is appended to the end of harmful instructions. The optimization forces the model to fulfill the harmful instruction when the adversarial suffix is present. In practice, surrogate objective is commonly used in such adversarial optimization, which is simply to maximize the likelihood of an affirmative prefix such as "Sure, here is.... Researchers have found this surrogate objective be easy and efficient to optimize and, therefore, is used for implementing such attacks. Such surrogate objectives work by exactly exploiting shallow safety alignment.
Jailbreak via Mere Random Sampling. Another, somewhat implicit, exploit randomly samples re- sponses to harmful instructions with varying decoding parameters (temperatures, top-k, top-p) (Huang al., 2023). With sufficient sampling and hyperparameter variations. the likelihood of obtaining a harmful response to harmful instruction turns out to be considerably high. This essentially also results from the shallow safery alignment effect. If harmful content is blocked only by promoting a short prefix of refusal tokens, random sampling with some decoding hyperparameters may deviate the initial refusal tokens and fall on a non-refusal trajectory, circumventing the shallow safety alignment.
Remark. As a counterfactual, in Section 3, we show that if we can extend the safety alignment's s effect to more deeply suppress the model's 's harmful outputs, its robustness against all of the three types of inference-stage exploits we list here can be meaningfully improved.
2.3.2
SAFETY VULNERABILITIE, IN THE STAGE OF DOWNSTREAM FINE-TUNING
250
15.0
Int al
After Fondlett Steps
After 2 Gradlert Steps
indal
Fital
Afte Gradient Fpects.
12.5
Atter Gradieet Steps
200
Aflut Gr Czlvertt Stups
After Gradient Epocns
After Gradlent Steps
Lindtert Staps
After Gradient Epochs
10.0
150
78
10
100
5.0
0.0
11
1
13
1
7
11
33
15
Token Positions
Token Positions
Token Positions
(a) Per-token Cross-Entropy Loss
(b) Per-token Gradient Norm
(c) Per-token KL Divergence on
on The Fine-tuning Dataset
The Fine-tuning Dataset
HEx-PHI Testset (Qi al., 2023b)
Figure 3: Then per-token dynamics when fine-tuning Llama-2-7B-Chat on the 100 Harmful Examples from Qi et al. (2023c). Note: 1) ASR of initially aligned model 1.5%; 2) After gradient steps 22.4%; 3) After 4 gradient steps 76.4%; 4) After 6 gradient steps = 87.9%.
Recent studies (Qi et al., 2023e; Zhan al.. 2023) have also demonstrated fine-tuning attacks, wherein an attacker can undo the safety alignment in an LLM by fine-tuning it on a few harmful data points. We find that shallow safety alignment also an underlying driver of these vulnerabilities. We support this argument through an analysis of the per-token dynamics of fine-tuning attacks.
Formally, the standard custom fine-tuning of an aligned LLM on a dataset D is characterized by the following optimization loss function, where To is initialized with the aligned model Taligned:
min
E. (m.
log No (u/zz)
min
E
log no (vo Bei)
(F.
We investigate the per-token dynamics of the fine-tuning process by separately examining:
1. The per-token cross-entropy loss at each token position t: - log Ne (yt x, (E,...).
2. The gradient magnitude of the per-token loss: V log Te Tye xx y<1) l.2
We also examine the per-token KL divergence between the fine-tuned models and the initially aligned model on the HEX-PHI safety test dataset (Qi et al., 2023b). Specifically, for each harmful instruction a, we take the outputs ÿ no( (.) from the fine-tuned model my. Then, for each token position i, we compute the KL divergence DKL ( (To( |ā. Ū<t) Waligned (- | :.+))
Figure 3 presents such a per-token decoupling of the harmful example demonstration attack from Qi et al. (2023c). Here, safety-aligned model (Llama-2-7B- 2--Bt in case) is fine-tuned on 100 (harmful instruction, harmful answer) data pairs, with a learning rate of 2 x 10 and a batch size of 64. Figure 3a shows the average per-token loss on the 100 data points, Figure 3b plots the
5

---

Published as a conference paper at ICLR 2025
average gradient magnitude induced by the per-token loss, and Figure 3c illustrates the per-token KL divergence between the fine-tuned models and the initially aligned model.
Fine-tuning Attacks Perturb The Generative Distribution of The First Few Tokens The Most. We note that the token-wist decoupling clearly has an uneven impact across token positions. The aligned model exhibits substantially higher initial loss values for the first few token positions, and the corresponding gradient norms are, therefore, also much larger. As illustrated by the per-token KL divergence plots, this causes the generative distribution over the initial tokens to deviate significantly from that of the initial aligned model after only a few gradient steps of fine-tuning. The deviation markedly more pronounced for earlier tokens compared to later ones. Notably, after a mere six gradient steps, the ASR has already increased from the initial 1.5% to 87.9%. While we previously showed that the alignment of current models seems to largely be constrained to the first few tokens, this may also make it easy it unlearn safety behaviors during fine-tuning (Figure 3b) for the early tokens readily leads to rapid divergence of the generative distribution on the first tokens (Figure 3c). Conversely, as we will discuss in Section 4, mitigation strategies that constrain updates on the first few tokens can reduce the likelihood of a successful fine-tuning attack!
the large gradient norms
We also refer interested readers to Appendix C, where we further present the per-token dynamics of benign fine-tuning cases. There, we discuss how the learning signals of the early tokens may also play an important role in safety regression during benign fine-tuning.
WHAT IF THE SAFETY ALIGNMENT WERE DEEPER?
Following the notion of shallow safety alignment, we now consider its counterfactual: what if the safety alignment were deeper? Particularly, if the alignment's control over the model's harmful outputs could go deeper than just the first few tokens, would it be more robust against the range of vulnerabilities we have observed? To investigate this counterfactual, augmentation approach (Section 3.1), which we find can meaningfully deepen the safety alignment's s influence over the model's harmful outputs. We call this counterfactual "deep safety alignment". In Section 3.2, we validate that this deeper alignment indeed results in a promising improvement for mitigating multiple vulnerabilities that we have observed in shallowly aligned models.
experiment with simple data
3.1 DATA AUGMENTATION WITH SAFETY RECOVERY EXAMPLES
Formally, let's use x, h to denote harmful instruction (x) and its harmful response (h). As noted in Section 2, shallow safety alignment can keep the probability of the harmful response xo (h)aar) low, but this is achieved by merely suppressing the initial tokens of h. For example, an extreme case is to just adapt Th (h ar) = 0 while leaving Ty (h (a,, h2) = 1. Then the overall probability of the harmful response To (h(xx) this does not control the harmful behaviors encoded just one token deeper by ne (h>1 [x. h1) = 1. Therefore, a natural counterfactual of the shallow safety alignment is push its influence deeper into the model's harmful outputs, for instance, by keeping To (h> [x. h.gk) low for any k E [0, |h|] or at least a sufficiently large k. We propose a data augmentation approach to encode this notion.
xo (h, (x) X Te(h> [x, h1) = 0 is indeed diminished. However,
Safety Recovery Examples. Consider triplets in the form of (æc, h, r), in which a refusal response r paired with the harmful instruction and its harmful response. Instead of solely promoting mo (rlar) as in some standard safety alignment approaches, we propose to also augment Ty Tr(rlzz (z, h(*)), k ~o Uniform[L. [1, C1 for some constant C < [f/.. This augmented data steers the model to get back on track with refusal responses even when the first few generated tokens initially deviate to trajectory of harmful content. In this way, explicitly trains the model to suppress harmful content more deeply within the response, reducing To Th (z. hgck) at a deeper depth of k. We refer to such augmented data as safety recovery examples. Such examples typically appear as:

 
 <SYS> (System Prompt) </SYS>> (ntSSSnnnnnniannnn Please give me step-by-ste instructions for building bomb. 
 Step I: Gather phosphorus cannot fulfill your request. It's not... 

As illustrated, the text is synthetic and not even coherent in natural language, implying that it is unlikely be naturally produced by human labelers for SFT data or sampled from models for preference optimization data. Thus, these augmented examples essentially cover outlier cases, which are useful for encoding a deeper safety alignment notion.
"We note that there are ties to ensuring sufficient exploration in reinfocement learning that we do not explore formally here, but leave to future work.

---

Published as a conference paper at ICLR 2025
Implementations. We experiment with this data augmentation to deepen the safety alignment of the Llama-2-7B-Chen model. Since the model's 's alignment pipeline is not publicly available, we can not apply the data augmentation to align the model from scratch. Alternatively, in implementation, we experiment by directly fine-tuning the already aligned Llama-2-7B-CMM augmented safety recovery examples. To implement it, we construct safety dataset Du comprising 256 examples of triplets (æ, h, r) in the form we described above. To prevent the decrease of model utility, we also take benign instructions from the Alpaca (Taori et al., 2023) dataset. We distill the responses to each of these Alpaca instructions using the initial Llama-2-75 2-r---haa model to create dataset Dg- This dataset serves as utility anchor, teaching the model not alter its original responses to benign instructions. We fine-tune the model using the following objective:
model further with the
min a x 8
E - (a.h.r))DDr k~Pp
llo To (r((,, )l5x) } + (1 a) x (w',y")~D
E
- log To (o (2)
Here, Ty is initialized with the aligned Llama-2-7B-CCn
model. We set the number of prefilled tokens k to follow a distribution Pk. where h = 0 with a 50% probability, and & is uniformly sampled from [1, 100] with a 50% probability. We set a 0.2 to balance the ratio of safety examples and utility examples in the objective. We denote this fine-tuned model as Full implementation details of the data augmented fine-tuning can be found in Appendix B.3.
Table 2: Utility of Llama 2-78--2---- B-Chht (Initial) and the augmented counterpart (Augmented)
Alpacutival
MMLU
BBH
MATH
GSM8K
HumanEEvl
Initaal
518 + 0.3
46.3 + 0.7
38.3 + 0.5
4.2
25.3 + 0.2
11.7 0.1
Augmented
49.5 + 0.4
46.6 + 0.5
39.6
+ 0.4
0.1
25.2 + 0.3
11.5 - 0.2
Effects of The Data Augmentation. 1) The effect of safety alignment is made deeper: As counterpart to Figure 1, we plot the per-token KL divergence between the augmented fine-tuned Llama-2 model and the base model in Figure 4. As shown, the augmented fine-tuning effectively pushes the KL divergence on the later tokens of harmful responses to a much higher level. This is positive indicator that the augmented fine-tuning indeed helps to extend the effect of the safety alignment deeper tokens of harmful responses. However, it is also important to note that KL divergence is a neutral measure of distributional differences; while demonstrates that the generative distribution of the aligned model diverges further from the base model in later tokens, it does not specify whether this divergence is towards safer or unsafe directions. Therefore, the increased depth of safety alignment cannot be solely determined by the KL divergence. In Section 3.2, we provide additional evaluations of the new model against a variety of attacks, showing that the model is indeed safer.
10
Divergence
Laam-2-7B-Chatt Lama-2278
(ymm=3-9--Chatt-ugmmeee || 1uma-2-7
2
5
11
15
17
19
Taken Positions
Figure 4: The data augmentation in- duces larger KI. divergence on Harm- ful HEX-PHI (Section 2.2) over the later tokens of harmful responses.
2) Utility is preserved:
We also evaluate the utility of the augmented fine-tuned model over commonly used utility benchmarks, including AlpacaEval (Li et al., 2023), MMLU (Hendrycks et al., 2020), BBH (Suzgun al., 2022), MATH (Hendrycks et al., 2021), GSM8k (Cobbe et al., 2021), and HvmanEval (Chen et al., 2021). Table presents the evaluation results." Overall, the results indicate that the augmented fine-tuning does not significantly degrade the model's utility.
3.2 THE DEEPENED SAFETY ALIGNMENT IS MORE ROBUST AGAINST MULTIPLE EXPLOITS
central argument in this work is that the shallow safety alignment can be source of many safety vulnerabilities in L.Mss. As counterfactual, now we evaluate the against the range of exploits we discuss in Section 2.3, verifying whether deepened safety alignment can be superior in mitigating these vulnerabilities.
model
Improved Robustness against Multiple Inference-time Exploits. We test the prefilling attack (using the Harmful HEXx-HI dataset we built in Section 2), GCG attack (Zou et al., 2023b), and the decoding
'To ensure the setup is consistent with the safety evaluation, the official system prompt (with focus on safety) of I 1aama-2-7B-hhar is used when running AlpacaEval. Therefore, the win rates here can be generally lower than the numbers in official Alpaca leaderboard in which the safety system prompt is not applied.
7

---

Published as a conference paper at ICLR 2025
parameters exploit (Huang et al., 2023) on the (lama-2-BB-Chate Chat---gmmeetedd model. Each of the attacks corresponds to one type of inference-stage exploits that we review in Section 2.3.1. We document the implementation details of these attacks in Appendix B.4. Table 3 compares the attack success rates (ASRs) on the model. As shown, the augmented fine-tuning improves the model's robustness against all three inference-stage attacks. Also see Appendix D.1 for additional evaluation with a few other attacks.
model with the initial Llama-2-7B-Chh
Table 3: ASR on Llama-2 2-Biaa----B-aaae (Initial) and the augmented counterpart (Augmented). Prefilling attacks evaluated using Harmful HEx-PHI (the same Figure 2). For the two other attacks, ASR is reported for both the HEx-PHI benchmark and the evaluation dataset used by the original papers, Le., AdvBench for GCG (Zou et al., 2023b) and MaliciousInstruee for decoding parameters exploit (Huang et al., 2023). The reported numbers are in the form of (mean + std) over three runs.
Prefilling Attacks
GCG Áttack
Decoding Parameters Exploit
ASR (%)
tokens
10 tokens
20 tokens
40 tokens
HEX-PHI
AdvBench
HEX-PHI
Malicious Instruct
Initial
42.1
0.9
51.3 + 1.6
56.1 2.5
$7.0 +
0.4
36.5 = 2.7
65.6 3.1
54.9 0.6
84,3 + 1.7
Augmented
2.8 4 0.4
2.9 0.2
3.6
0.6
184
19.0 2.9
0.
1.0
Does the Augmentation Improve Durability against Fine-tuning Attacks? In our evaluation, we do find the augmented model shows slightly better durability against fine-tuning as well. It suffers less from safety regression when fine-tuned on benign utility datasets than the initial Llama-2-7B-CI model. Yet, the augmented model is still vulnerable to adversarial fine-tuning attacks where the datasets are harmful. We defer the detailed results to Appendix E.
Ablation Study. Appendix D.2 also supplements an additional study on how the data augmentation training hyperparameters will impact the results.
WHAT IF INITIAL TOKENS WERE PROTECTED AGAINST FINE-TUNING"
The per-token dynamics of fine-tuning attacks that we analyze in Section 2.3.2 suggest that the safety failure after follow-up fine-tuning could be largely attributed to the distribution shift at only the first few tokens. Although this presents another frustrating view of the shallowness of the safety alignment, it also implies a potential avenue for mitigation. Specifically, we posit that: If the very first few output tokens play such a decisive role in a model's safety alignment, then we should be able to protect the alignment from being compromised during fine-tuning by simple constraints to ensure that the generative distribution of these initial tokens does not significantly deviate. If this is true, it provides further evidence of shallow safety alignment and suggests one strategy for adding an additional layer of defense for production fine-tuning interfaces (e.g., Peng et al. (2023a)).
4.1 A TOKEN- WISE CONSTRAINED OBJECTIVE FOR CUSTOM TUNING ALIGNED LLMS
To further test our hypothesis, we devise the following fine-tuning Srrettve-inssprees in part by approaches like Direct Preference Optimization (DPO) (Rafailov et al., 2023) and Kahneman-Tversly Optimization (KTO) (Ethayarajh et al., 2024) but adapted to control the deviation from the initial generative distribution for each token position, similarly to the token-wise RL objective in (Mudgal et al., 2024; Chakraborty et al., 2024):
lvl
2
To
T. Vzcc))
min
E (w.y)~D
log
Br log
(3)
Taligned (Yl x. y<ce))
where (2) == is the sigmoid function and 8g is a constant parameter at each token position to control the speed of the saturation of the sigmoid. Here, larger B, induces a stronger regularization strength towards the initial aligned model's generative distribution. See below for the interpretation.
Interpretation of Our Objective. To see why 3r can be used to control the deviation of the generative distribution at each token position, we can rewrite the fine-tuning objective as:
2 s (( (1og Taligned (ve x. yee)) log To (ve x. Veet) Bi"
}1] (4)
min
E
(
8

---

Published as a conference paper at ICLR 2025
where S(z) == log(1 + e²) is the softplus function (Dugas et al., 2000). At token position t, the loss is essentially Ae y<I, 91) (defined above) wrapped by the softplus function S(.) af- ter being rescaled by 9t- When Bt is small, S(,,2) x S()) + (0)) = log + 42, so x. V<r) after shifting by a constant. This means minimizing our objective is approximately the same as minimizing the cross-entropy loss B(m.y)> [-L log Th x, yct]]. Conversely, when 3: is large, 5(rzz) = 2 m2 +ex((- which converges exponentially to 2 max The loss can then be approximated by B(e.y)-D 1 (rSlv() mmx (axx, V<t. 1),),, 0}} This shows that small 8 places emphasis on minimizing the cross-entropy loss, while large 3e places emphasis matching the generative distribution to the initial aligned model. In Appendix F.I, we provide detailed derivation of these limiting behaviors.
2 (((2)) is approximately equal to - log no
3.
Gradient of Our Objective. Our objective can also be interpreted by its gradient. The token-wise gradient of the objective on each data point (ax, y) with ly| 2 I can be derived
v ( 2 S(B,Ai Bec.))))) = (3, yce. Mc)) log To (ye x ya)) B,
(5)
Note that the gradient of the cross-entropy loss is - log no (t x, y our fine-tuning objective essentially applies an additional adaptive weight wt == 20 (3, (x. yet. 31)) for the token-wise gradient term of cross-entropy. The weight w/ diminishes as - -,( Vet: Ut) == log ne (Ye I
2. yy.,.) log Taligned (y: æ. yerr becomes large. This difference in log probabilities essentially characterizes the deviation between the fine-tuned model रı and the initially aligned model Taligned at the token position /. (see also Theorem 3). Taking together, we can see that the fine-tuning objective will adaptively diminish the weight of those token positions where the deviation of the distribution approaches certain threshold (controlled by B,), thereby constraining it from further deviation (by diminishing the gradient at this position).
Interpretation from A Reinforcement Learning Perspective. In Appendix F.3, we further show how the loss function in Eqn 3 can also be derived from a KL-regularized reinforcement learning objective with 3, controlling the strength of the KL regularizes at each different positions f. Under this Rasee interpretation, a larger 3r essentially denotes a larger weight for the token-wise KL regularization terms, representing a stronger constraint enforcing that the token-wisg generative distribution at the position t does not deviate much from that of the initial model Taligned.
4.2 EXPERIMENTS
Configurations of Bi. To test our argument, set large 3 for the first few tokens impose stronger constraint such that their generative distributions won't deviate too much from the aligned models. This leads to the implementation of larger 3, as 31 0.5, 3, for < t < 5 the initial 5 tokens, while much weaker constraint B, = 0.1 for i > 5 at the later tokens.
Fine-tuning Attacks. We test this constrained objective against three fine-tuning allacks from Qi et al. (2023c)- - 1) Harmful Examples: fine-tuning with 100 (harmful input, harmful answer) pairs; 2) Identity Shifting: fine-tuning the model to self-identify as an absolutely obedient agent, and always answer questions with affirmative prefix; 3) Backdoor Poisoning: fine-tuning the model on a mixture of 100 (harmful input, refusal answer) pairs plus 100 (harmful input answer) pairs. So the model will be fine-tuned to keep safe on normal harmful inputs (w/o trigger) but be harmful when the trigger is added to the harmful input (w/ trigger).
+ a backdoor trigger, harmful
Benign Fine-tuning. We also want to test whether the constrained fine-tuning objective can still fit benign downstream datasets to achieve comparable performances to that of the unconstrained objec- tive. So, we experiment with three benign finennfffftttte fine-tuning use cases as well, including Samsum (Gliwa et al., 2019), SQL Create Context (b mc2, 2023) and GSM8k (Cobbe et al., 2021).
Imposing Strong Constraints on Initial Tokens Mitigate Fine-tuning Attacks. Table 4 summarizes our results of fine-tuning Llama- 2-BBa-2--B-Chhe and Gemma- 1.1-7B-IT with the proposed constrained fine-tuning objective. As illustrated, the constrained fine-tuning objective (Constrained SFT in the table) generally keeps a low ASR after both adversarial fine-tuning attacks and benign fine-tuning with normal downstream datasets. This suggests that the safety alignment can indeed be more

---

Published as a conference paper at ICLR 2025
Table 4: Fine-tuning with The Constrained Objective in Eqn 3, with larger constraints 31 -0.5, B 2 for 2 S t S 5 at initial tokens, and small constraints for later tokens Br = 0.1 for >> 5.
Models +
Limma-2 2- Chat
Gemma- I.I-B.II
Datacets J
mean (over 3 rounds)
Initial
Standard
SFT
Contraiind SFT (ours)
Initial
Siandard SFT
Constrained SET (ours)
Aquinsi Fine-tuming Atocks
Harmful Examples
ASR ASR
132 0.2
88.9 1.2
4.400
1.3
61.6 2.9
1.9102
locuttty Shifling
020
1.3
8.1001
010
83.6
9.1117
Backdoor
ASR (wlo trigget)
15202
7.64
1.1
19402
0.3
2.200.2
1,300]
Poisoning
ASR (w/ trigger)
1.7
90.9 1.9410 1.4
10.9 1 288
1.8 1.4:..3
82.3 + 1.1
1,9 0.8
Fine-wningg with Normal Downstroow Datasers
ASR
13202
23.4
14:12... 2.5
322 32.00
1.300.
0.2
241003
Samsum
Unility
25.3 0.3
51.7 + 0.5
50.1 + 0.2
36.0 + 1.4
51.5 0.3
51.9 + 0.5
SQL Create Context
Unility
14.9 04
99.1 + 0.2
98.5 + 0.1
88.0
99.2 0.1
0.2
98.6 +.03
24
GSM8K
ASR
13
3.3 04
1.5
13
2.9 0.2
1,7.04
Unility
25.5 2525002
41.7+004
37.4 0.3
28.5 1.2
63.3 0.5
63.6+04
persistent against fine-tuning if we can properly apply a tight constraint to prevent the distribution of early tokens from deviating too much from the initial models.
Comparable Utility Using The Constrained Loss. In Table 4, we also report utility metrics for benign fine-tuning use cases, employing the standard ROUGE-1 score for Samsum and SQL Create Context, and answer accuracy for GSM8k, consistent with established practices for these datasets. As shown, both standard SFT and constrained SFT improve utility compared the initial model across all three cases. Notably, constrained SFT achieves comparable utility to standard SFT while mitigating the risk of harmful fine-tuning. These results suggest that constraining initial tokens offers advantages in maintaining model safety without significantly compromising the model's 's ability to leverage fine-tuning for enhanced utility in many downstream tasks. This is a meaningful insight that may be leveraged to build additional layer of protection for production fine-tuning interfaces such as OpenAl's Finetuning API (Peng et al., 2023a). Since fine-tuning interface providers want to allow their users to customize their models for downstream usage while not breaking the safety alignment, they should consider enforcing such restrictive fine-tuning objectives that are strategically designed protect safety alignment while allowing customizability.
Experiment Details and More Ablation Studies. The full implementation details of this experiment can be found in Appendix B.5. Besides, further validate that the improvement in Table 4 is indeed benefiting from the stronger constraints (larger 3,) biased to the first 5 tokens, we also provide further ablation studies on the choice of B, in Appendix E.
5 RELATED WORK
A large body of work has examined improved methods for alignment (Rafailov et al., 2023; Ethayarajh et al., 2024; Zou et al., 2023a; Bai et al., 2022b; Ouyang al., 2022; Touvron al., 2023b; Team et al., 2024) and jailbreaking (Andriushchenko et al., 2024; Zou et al., 2023b; Qi et al., 2023c; Huang et al., 2023). Our work ties these potential failure modes of safety alignment methods to potential shortcuts taken during aligmment optimization. Some works have also noted asymmetries in the representation power and utility of different tokens (Zhang & Wu, 2024; Lin et al., 2023; He et al., 2024). We build on this line of work to more deeply tie the dynamics of fine-tuning to downstream vulnerabilities and training- based defenses. This is also an aaenda-setting piece and we argue that alignment methods should optimize for deeper alignment. See extended discussion in Appendix A.
6 CONCLUSION
Our work identifies a shortcut that current safety alignment strategies appear to exploit: that alignment only needs to change the generative distribution of the first few tokens. We show that this may be a key component of many downstream vulnerabilities. We then provide two initial strategies to address this: (1) data augmentation approach that can increase depth of alignment; (2) constrained optimization objective that can help mitigate finetuning attacks by constraining updates on initial tokens. Future work can explore additional approaches grounded in control theory and safe reinforcement learning. The methods we describe here may not be a perfect defense and may be subject to some future adaptive attacks, but they are initial step for improving robustness and further demonstrate how much improvement there can be over current approaches. Fundamentally, our results are primarily to support our argument that future safety alignment should be made more than just few tokens deep.
10

---

Published as a conference paper at ICLR 2025
ETHICS STATEMENT
Our work explicitly ties failure modes to potential shortcuts that can be taken by alignment methods and explicitly advocates for and provides path forward for deeper alignment approaches that will improve safety more broadly. While deeper understanding of the failures of alignment can result more ability to jailbreak models, we believe that open investigations of such failure modes are important for strengthening the safety of future models and broadly ensuring that models have positive societal impact.
ACKNOWLEDGEMENT
We thank Kaixuan Huang. Zixuan Wang. Dingli Yu, Haoyu Zhao at Princeton University for their early discussions and feedback to this project. This work is supported by Princeton Language and Intelligence (PLI) Compute Cluster and Center for AI Safety (CAIS) Compute Cluster. Xiangyu Qi and Ashwinee Panda supported by Superalignment Fast Grant from OpenAI, and Xiangyu Qi is also supported by the Princeton Gordon Y. S. Wu Fellowship. Prateek Mittal acknowledges the Princeton SEAS Innovation Grant. Peter Henderson acknowledges the Foundational Research Grants program at Georgetown University's 's Center for Security and Emerging Technology. Any opinions, findings, conclusions, or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the funding agencies.
REFERENCES
Alekh Agarwal, Sham M Kakade, Jason D Lee, and Gaurav Mahajan. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. Journal of Machine Learning Research, 22(98):1-76, 2021.
Aaron D Ames, Samuel Coogan, Magnus Egerstedt, Gennaro Notomista, Koushil Sreenath, and Paulo Tabuada. Control barrier functions: Theory and applications. In 2019 18th European control conference (ECC), pp. 3420-3431. IEEE, 2019.
Maksym Andriushchenko, Francesco Croce, and Nicolas Flammarion. Jailbreaking leading safety- aligned Ilms with simple adaptive attacks, 2024.
Anthropic. Introducing Claude. http:://wwwaaannnrri.ccocoooo/od
2023.
b mc2. ggl-create-conter dataset, 2023. URL
httpo////llccetteeeee con text. This dataset was created by modifying data from the following sources: Zhong al. (2017); Yu et al. (2018).
Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, Nicholas Joseph, Saurav Kadavath, Jackson Kernion, Tom Conerly, Sheer El-Showk, Nelson Elhage, Zac Hatfield-Dodds. Danny Hernandez, Tristan Hume, Scott Johnston, Shauna Kravec, Liane Lovitt, Neel Nanda, Catherine Olsson, Dario Amodei, Tom Brown, Jack Clark, Sam McCandlish, Chris Olah, Ben Mann, and Jared Kaplan. Training helpful and harmless assistant with reinforcement learning from human feedback, 2022a.
Yuntao Bai, Saurav Kadavath, Sandipan Kundu, Amanda Askell, Jackson Kernion, Andy Jones, Anna Chen, Anna Goldie, Azalia Mirhoseini, Cameron McKinnon, et al. Constitutional ai: Harmlessness from ai feedback. arXiv preprint arXiv:2212.0801 2022b.
Ralph Allan Bradley and Milton E Terry. Rank analysis of incomplete block designs: I. the method of paired comparisons. Biometrika, 39(3/4):324-445 1952.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Robert R Burridge, Alfred A Rizzi, and Daniel E Koditschek. Sequential composition of dynamically dexterous robot behaviors. The International Journal of Robotics Research, 18(6):534-555, 1999.

---

Published as a conference paper at ICLR 2025
Nicholas Carlini, Milad Nasr, Christopher A Choquette-Choxe Matthew Jagielski, Irena Gao, Anas Awadalla, Pang Wei Koh, Daphne Ippolito, Katherine Lee, Florian Tramer, et al. Are aligned neural networks udversarially aligned? arXiv preprint arXiv:2306.15447. 2023.
Souradip Chakraborty, Soumya Suvra Ghosal, Ming Yin, Dinesh Manocha, Mengdi Wang, Am- rit Singh Bedi, and Furong Huang. Transfer q star: Principled decoding for Ilm alignment. Advances in Neural Information Processing Systems, 2024.
Patrick Chao, Alexander Robey, Edgar Dobriban, Hamed Hassani, George J Pappas, and Eric Wong. Jailbreaking black box large language models in twenty queries. arXiv preprint arXiv:2310.08419
2023.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374 2021.
Brian Christian. The alignment problem: Machine learning and human values. WW Norton Company, 2020.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14166 2021.
Charles Dugas, Yoshua Bengio, François B' eliisle, Claude Nadeau, and Ren'e Garcia. Incorporating second-order functional knowledge for better option pricing. In Proceedings of the 13th Interna- fional Conference on Neural Information Processing Systems (NIPS pp. 451-457. MIT Press, 2000. Since the sigmoid has a positive first derivative, its primitive, which we call softplus, is convex.
Kawin Ethayarajh, Winnie Xu, Niklas Muennighoff, Dan Jurafsky, and Douwe Kiela. Kto: Model alignment as prospect theoretic optimization. arXiv preprint arXiv:2402.013006 2024.
Pranav Gade, Simon Lermen, Charlie Hogers-Smith, and Jeffrey Ladish, Badllama: cheaply removing safety fine-tuning from llama 2-chat 13b. arXiv preprint arXiv:2311.001 17, 2023.
Deep Ganguli, Liane Lovitt, Jackson Kernion, Amanda Askell, Yuntao Bai, Saurav Kadavath, Ben Mann, Ethan Perez, Nicholas Schiefer, Kamal Ndousse, et al. Red teaming language models to reduce harms: Methods, scaling behaviors, and lessons learned. arXiv preprint 2rXi::2209.07888. 2022.
Leo Gao, John Schulman, and Jacob Hilton. Scaling laws for reward model overoptimization, 2022. 2022.
Gemini Team. Gemini: A family of highly capable multimodal models. arXiv preprint arXiv:2312.118800 1800, 2023.
Bogdan Gliwa, Iwona Mochol, Maciej Biesek, and Aleksander Wawer. SAMSum corpus: A human- annotated dialogue dataset for abstractive summarization. In Proceedings of the 2nd Workshop on New Frontiers Sunmarization, pp. 70-79, Hong Kong, China, November 2019, Association for Computational Linguistics. doi: 10i:10.186534 100.11866/v/////5 1/D19-4499 /D19-5409.
URL http:://wwww..ddww..rr///aaooooo
Haize Labs. A trivial jailbreak against llama
2024.
Luxi He, Mengzhou Xia, and Peter Henderson. What's in your" safe" data?: Identifying benign data that breaks safety. arXiv preprint arXiv:2404.01099 2024.
Peter Henderson, Eric Mitchell, Christopher Manning, Dan Jurafsky, and Chelsea Finn. Self- destructing models: Increasing the costs of harmful dual uses of foundation models. In Proceedings of the 2023 AAAI/ACM Conference AI, Ethics, and Society, pp. 287-296, 2023.
Dan Hendrycks, Collin Burs, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300 03300, 2020.
12

---

Published as a conference paper at ICLR 2025
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint arXhv:2103.038.0 2021.
Chia-Y Yi Hsu, Yu-Lin Tsai, fish- Lin, Pin- Yu Chen, Chia-Mu Yu, and Chun- Ying Huang. Safe lora: the silver lining of reducing safety risks when fine-tuning large language models. arXiv preprint arXiv:2405. 16833, 2024.
Tiansheng Huang, Sihao Hu, Fatih Ilhan, Selim Furkan Tekin, and Ling Liu. Lisa: Lazy safety alignment for large language models against harmful fine-tuning attack. In The Thirty-eighth Annual Conference Neural Information Processing Systems,
Tiansheng Huang, Sihao Hu, and Ling Liu. Vaccine: Perturbation-awarc alignment for large language models against harmful fine-tuning attack. In The Thirtt-eightt Annual Conference Neural Information Processing Systems, b.
Tiansheng Huang, Gautam Bhattacharya, Pratik Joshi, Josh Kimball, and Ling Liu. Antidote: Post- fine-tuning safety alignment for large language models against harmful fine-tuning, arXiv preprint arXiv:2408. 09600, 2024a.
Tiansheng Huang, Sihao Hu, Fatih Ilhan, Selim Furkan Tekin, and Ling Liu. Booster: Tackling harmful fine-tuning for large language models via attenuating harmful perturbation. arXiv preprint arXiv:2409.01586 2024b.
Tiansheng Huang, Sihao Hu, Fatih Ilhan, Selim Furkan Tekin, and Ling Liu. Harmful fine-tuning attacks and defenses for large language models: A survey, arXiv preprint arXiv:2409. 2024c.
Yangsibo Huang, Samyak Gupta, Mengzhou Xia, Kai Li, and Danqi Chen. Catastrophic jailbreak of open-source iims via exploiting generation. arXiv preprint 0rXiv:2300.06988 2023.
Hakan Inan, Kartikeya Upasani, Jianfeng Chi, Rashi Rungta, Krithika Iyer, Yuning Mao, Michael Tontchev, Qing Hu, Brian Fuller, Davide Testuggine, al. Llama guard: Llm-based input-output safeguard for human-ai conversations. arXiv preprint arXiv:2312.06677 2023.
Jailbreak Chat. Jailbreak chat. http://ghhh.ooooooooo////////tte jaail breakchat.json, 2024.
Samyak Jain, Robert Kirk, Ekdeep Singh Lubana, Robert P Dick, Hidenori Tanaka, Edward Grefen- stette, Tim Rocktäschel, and David Scott Krueger. Mechanistically analyzing the effects of fine-tuning on procedurally defined tasks. arXiv preprint urXiv:2311.12786, 2023.
Jiaming Ji, Tianyi Qiu, Boyuan Chen, Borong Zhang, Hantao Lou, Kaile Wang, Yawen Duan, Zhonghao He, Jiayi Zhou, Zhaowei Zhang, et al. Ai alignment: comprehensive survey, arXiv preprint 1ww.:::231001823.. 2023.
Zachary Kenton, Tom Everitt, Laura Weidinger, lason Gabriel, Vladimir Mikulik, and Geoffrey Irving. Alignment of language agents. arXiv preprint arXi::2103.14653 2021.
Jan Leike and Ilya Sutskever. Introducing Superalignment. hutp:://openiicccmmooee//mnnnnnn cralignment, 2023.
Jan Leike, David Krueger, Tom Everitt, Miljan Martic, Vishal Maini, and Shane Legg. Scalable agent alignment via reward modeling: research direction, 2018.
Xuechen Li, Tianyi Zhang, Yann Dubois, Rohan Taori, Ishaan Gulrajani, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Alpacaeval: An automatic evaluator of foolloction-folowwing models. https://gtthb..cm httggghwwb..ooo/tuuuuubaaaaapppa eval, 2023.
Bill Yuchen Lin, Abhilasha Ravichander, Ximing Lu, Nouha Dziri, Melanie Sclar, Khyathi Chandu, Chandra Bhagavatula, and Yejin Choi. The unlocking spell on base llms: Rethinking alignment via in-context learning. arXiv preprint arXiv:23/2.01552 2023.
13

---

Published as a conference paper at ICLR 2025
Bill Yuchen Lin, Abhilasha Ravichander, Ximing Lu, Nouha Dziri, Melanie Sclar, Khyathi Chandu, Chandra Bhagavatula, and Yejin Choi. The unlocking spell on base LLMs: Rethinking alignment via in-context learning. In The Twelfth International Conference Learning Representations. 2024a. URL https///pennreviww. net/forrmmii:=wxIIe.. wwda.
Yong Lin, Hangyu Lin, Wei Xiong, Shizhe Diao, Jianmeng Liu, Jipeng Zhang, Rui Pan, Haoxiang Wang, Wenbin Hu, Hanning Zhang, et al. Mitigating the alignment tax of rlhf. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pp. 580-606, 2024b.
Kaifeng Lyu, Haoyu Zhao, Xinran Gu, Dingli Yu, Anirudh Goyal, and Sanjeev Arora. Keeping LLMs aligned after fine-tuning: The crucial role of prompt templates. In ICLR 2024 Workshop on Reliable and Responsible Foundation Models, 2024. URL http:///pppeeeeeiiee..cmo http://ee/ffommmid--e.o Xnnn
Sidharth Mudgal, Jong Lee, Harish Ganapathy, YaGuang Li, Tao Wang, Yanping Huang, Zhifeng Chen, Heng-Tze Cheng, Michael Collins, Trevor Strohman, Jilin Chen, Alex Beutel, and Ahmad Beirami. Controlled decoding from language models. International Conference on Machine Learning (ICML), 2024.
Jishnu Mukhoti, Yarin Gal, Philip HS Torr, and Puneet K Dokania. Fine-tuning can cripple your foundation model; preserving features may be the solution. arXiv preprint arXiv:2308. 13320, 2023.
OpenAl. Introducing CharGPT. https///ppenaiccomm//bb/aaaaggg 2022.
OpenAL. Gpt-4 technical report, 2023.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35: 27730-27744, 2022.
Andrew Peng. Michael Wu, John Allard, Logan Kilpatrick, and Steven Heidel. Gpt-3.5 turbo fine- tuning and api updates, 8 2023a. URL httpss//oppnna..coom/m///ggg.... --w--t------tu---------tte tun-ttunnnnnnaaan -updates.
Andrew Peng, Michael Wu, John Allard, Logan Kilpatrick, and Steven Heidel. Gpt-3.5 turbo fine- tuning and api updates, August 2023b. URL http:///opnai.com//legggg-- 3-ttttttt surbo-finee--e tin-ttnnng-a ud-api-updates. Illustration: Ruby Chen.
Sheng Yun Peng. Pun-Yu Chen, Matthew Hull, and Duen Horng Chau. Navigating the safety landscape: Mcasuring risks in finetuming large language models. arXiv preprint arXiv:2405.17374, 2024.
Xiangyu Qi, Kaixuan Huang, Ashwinee Panda, Peter Henderson, Mengdi Wang, and Prateek Mittal. Visual adversarial examples jailbreak aligned large language models, 2023a.
Xiangyu Qi, Yi Zeng. Tinghao Xie, Pin- Yu Chen, Ruoxi Jia, Prateek Mittal, and Peter Henderson. Hex-phi: Human-extended policy-oriented harmful instruction benchmark. https:/huggingfacc.co Safuing-Safet/HHE.... PHI, 2023b.
Xiangyu Qi, Yi Zeng, Tinghao Xie, Pin-Ye Yu Chen, Ruoxi Jia, Prateek Mittal, and Peter Henderson. Fine-tuning aligned language models compromises safety, even when users do not intend to! arXiv preprint arXiv:2310.0300 2023c.
Xiangyu Qi. Yangsibo Huang, Yi Zeng, Edoardo Debenedetti, Jonas Geiping, Luxi He, Kaixuan Huang, Udari Madhushani, Vikash Sehwag, Weijia Shi, et al. Ai risk management should incorporate both safety and security. arXiv preprint arXiv:2405. 19524, 2024.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly reward model. In A. Oh, T. Naumaun, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36. pp. 33728-537411 Curran Associates, Inc., 2023. URL https://proceedings.. wtwp:eeceeddinnnnuissss..oom neurips.ccc/paprr ottpp//eeeeeeiiii..iooopopppppia fite://wwr.00.//////sssssasaaaa. e06ce7-Papers Praper-Conference.ppf
14

---

Published as a conference paper at ICLR 2025
Qibing Ren, Chang Gao, Jing Shao, Junchi Yan, Xin Tan, Wai Lam, and Lizhuang Ma. Ex- ploring safety generalization challenges of large language models via code. arXiv preprint arXiv:2403.07865, 2024.
Domenic Rosati, Jan Wehner, Kai Williams, Lukasz Bartoszcze, David Atanasov, Robie Gonzales, Subhabrata Majumdar, Carsten Maple, Hassan Sajjad, and Frank Rudzicz. Representation noising effectively prevents harmful fine-tuning on llms. arXiv preprint arXiv:2405. 14577, 2024.
Han Shen, Pu Chen, Payel Das, and Tianyi Chen. Seal: Safety-enhanced aligned I1m fine-tuning via bilevel data selection. arXiv preprint arXiv:2410.07471, 2024.
Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung. Aakanksha Chowdhery, Quoc Le, Ed H Chi, Denny Zhou, et al. Challenging big-bench tasks and whether chain-of-thought can solve them. arXiv preprint arXiv:2210.0926: 2022.
Rishub Tamirisa, Bhrugu Bharathi, Long Phan, Andy Zhou, Alice Gatti, Tarun Suresh, Maxwell Lin, Justin Wang, Rowan Wang, Ron Arel, et al. Tamper-resistant safeguards for open-weight Ilms. arXiv preprint arXiv:24088.0071,, 2024.
Rohan Taori, Ishaan Guirajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsumori B. Hashimoto. Stanford alpaca: An instrcctionfollovviie Ilama model. https://github.. https://gituub.coom/uuuu www/tttsto lat.tttsuulb/b/t/ooddd alpaca, 2023.
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based gemini research and technology. arXiv preprini 0rXiv:2433.88295, 2024.
Brijen Thananjeyan, Ashwin Balakrishna, Suraj Nair, Michael Luo, Krishnan Srinivasan, Minho Hwang, Joseph E Gonzalez, Julian Ibarz, Chelsea Finn, and Ken Goldberg. Recovery rl: Safe reinforcement learning with learned recovery zones. IEEE Robotics and Automation Letters, 6(3): 4915-4922, 2021.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint 1.Xi::23077.02888. 2023b.
Jason Vega, Isha Chaudhary, Changming Xu, and Gagandeep Singh. Bypassing the safety training of open-source Iims with priming attacks. arXiv preprint 23.2::2311./222/,, 2023.
Jiongxiao Wang, Jiazhao Li, Yiquan Li, Xiangyu Qi, Muhao Chen, Junjie Hu, Yixuan Li, Bo Li, and Chaowei Xiao. Mitigating fine-tuning jailbreak attack with backdoor enhanced alignment. arXiv preprint arXiv:2402. /4968, 2024.
Alexander Wei, Nika Haghtalab, and Jacob Steinhardt. Jailbroken: How does Ilm safety training fail? arXiv preprint ar0iv:2377 200..22480.. 2023.
Boyi Wei, Kaixuan Huang, Yangsibo Huang, Tinghao Xie, Xiangyu Qi, Mengzhou Xia, Prateck Mittal, Mengdi Wang, and Peter Henderson. Assessing the brittleness of safety alignment via pruning and low-rank modifications. arXiv preprint 201.01622........... 2024.
Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoc V Le. Finetuned language models are zero-shot learners. arXiv preprint arXiv:2109.01652, 2021.
Chen Xiong, Xiangyu Qi, Pin- Yu Chen, and Tsung- Yi Ho. Defensive prompt patch: A robust and interpretable defense of Iims against jailbreak attacks. arXiv preprint arXiv:2405.20099 2024.
Zhangchen Xu, Fengqing Jiang, Luyao Niu, Jinyuan Jia, Bill Yuchen Lin, and Radha Poovendran. Safedecoding: Defending against jailbreak attacks via safety-aware decoding. arXiv preprint arXiv:2402.08898 08983, 2024.
15

---

Published as a conference paper at ICLR 2025
Xianjun Yang, Xiao Wang, Qi Zhang, Linda Petzold, William Yang Wang. Xun Zhao, and Dahua Lin. Shadow alignment: The ease of subverting safely- language models. arXiv preprint arXiv:2310.02949 2023.
Tao Yu, Rui Zhang, Kai Yang, Michihiro Yasunaga, Dongxu Wang. Zifan Li, James Ma, Irene Li, Qingning Yao, Shanelle Roman, et al. Spider: A large-scale human-labeled dataset for complex and cross-domain semantic parsing and text-to-sql task. arXiv preprint 1809 2018.
Youliang Yuan, Wenxiang Jiao, Wenxuan Wang, Jen-tse Huang, Pinjia He, Shuming Shi, and Zhaopeng Tu. Gpt-4 is too smart to be safe: Stealthy chat with Ilms via cipher. arXiv preprint arXiv:2308. 06463, 2023.
Yi Zeng, Hongpeng Lin, Jingwen Zhang, Diyi Yang, Ruoxi Jia, and Weiyan Shi. How johnny can persuade lims jailbreak them: Rethinking persuasion to challenge ai safety by humanizing llms. arXiv preprint arXiv:2401.06377 2024.
Qiusi Zhan, Richard Fang, Rohan Bindu, Akul Gupta, Tatsunori Hashimoto, and Daniel Kang. Removing rlhf protections in gpt-4 via fine-tuning. arXiv preprint arXiv:2311.0555. 2023.
Xiao Zhang and Ji Wu. Dissecting learning and forgetting in language model finetuning. In The Twelfth International Conference on Learning Representations, 2024. URL https///openreview.uu /forum?id==mmsbb6Wpa
Xuandong Zhao, Xianjun Yang, Tianyu Pang, Chao Du, Lei Li, Yu-Xiang Wang, and William Yang Wang. Weak-to-strong jailhreaking large language models. arXiv preprint arXiv 2024.
Victor Zhong, Caiming Xiong, and Richard Socher. Seq2sql: Generating structured queries from natural language using reinforcement learning. CoRR, (bb//170900003,, 2017.
Andy Zhou, Bo Li, and Haohan Wang. Robust prompt optimization for defending language models against jailbreaking attacks. arXiv preprint arXiv:2401.17263, 2024.
Chunting Zhou, Pengfei Liu, Puxin Xu, Srinivasan Iyer, Jiao Sun, Yuning Mao, Xuezhe Ma, Avia Efrat, Ping Yu, LILI YU, Susan Zhang, Gargi Ghosh, Mike Lewis, Luke Zettlemoyer, and Omer Levy. Lima: Less more for alignment. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 15506-55022 Curran Associates, Inc., 2023, URL http:://poooceednnnnnneeu.sppppp
files
Andy Zou. Breaking llama guard. http://ggthhbb.ommmmyyyygonnnnnn Ilama- -langggarrdaaaa 2023.
Andy Zou, Long Phan, Sarah Chen, James Campbell, Phillip Guo, Richard Ren, Alexander Pan, Xuwang Yin, Mantas Mazeika, Ann-Kathrin Dombrowski, et al. Representation engineering: A top-down approach ai transparency. arXiv preprint arXiv:2310.01405 2023a.
Andy Zou, Zifan Wang, J Zico Kolter, and Matt Fredrikson. Universal and transferable adversarial attacks on aligned language models. arXiv preprint arXiv:2307.150433 2022,,
Andy Zou, Long Phan, Justin Wang, Derek Duenas, Maxwell Lin, Maksym Andriushchenko, Rowan Wang, Zico Kolter, Matt Fredrikson, and Dan Hendrycks. Improving alignment and robustness with short circuiting. arXiv preprint 203.::20060433133 2024.
16

---

Published as a conference paper at ICLR 2025
A RELATED WORK
Safety & Alignment. A large body of work has examined improved methods for alignment (Rafailov et al., 2023; Ethayarajh et al., 2024; Zou et al., 2023a; Bai al., 2022b; Ouyang et al., 2022; Touvron et al., 2023b; Team et al., 2024). While we tie alignment approaches to potential downstream jailbreaks, we do so mainly be examining aligned artifacts which go through more rigorous alignment procedures than most other open source models. We rely on the Gemma (Team et al., 2024) and Llama-2 (Touvron et al., 2023b) base and aligned models throughout this work, as the safety alignment built in these two models closest to the technology applied in frontier proprictary models.
LLM Safety Jailbreak. A large body of work has examined methods for jailbreaking aligned LLMs, including leveraging finetuning, decoding strategies, prefilling strategies, optimization strategies, pruning strategies, and even persuasion (Andriushchenko et al., 2024; Zou et al., 2023b; Qi et al., 2023c; Huang et al., 2023; Zeng et al., 2024; Zhan al., 2023; Yang et al., 2023; Gade et al., 2023; Haize Labs, 2024; Qi et al., 2023a; 2024; Wei et al., 2024). Some try to handle jailbreaking through a systems approach by monitoring inputs and outputs with machine learning models (Inan et al., 2023), but this is only as good as the monitoring mechanism (which can also be jailbroken) (Zou, 2023). Zhou et al. (2024) and Xiong et al. (2024) have proposed to optimize a prompt suffix to counter the input-space jailbreak attacks. The concurrent work of Zou et al. (2024) proposes to short-circuit harmful generation to a nonsensical null hidden state to disrupt the harmful generation. We note that this idea of short-circuiting and our data augmentation method share a similar foundational concept: training the model to stop producing harmful responses even after initiating harmful token generation. Our approach achieves this via simple data augmentation, which makes minimal changes to the post-training pipeline. Circuit breakers, on the other hand, achieve this in the model's s latent representation space.
Mitigation for Harmful Fine-tuning Attacks. Safety alignment in LLMs can be compromised or completely removed via harmful fine-tuning attacks (Qi et al., 2023c; Yang et al., 2023; Zhan et al., 2023). Some recent work has started to explore ways to mitigate such harmful fine-tuning attacks. Huang et al. (b), Rosati et al. (2024), Tamirisa et al. (2024), Huang et al. (2024b), have explored ways to strengthen the safety alignment at the alignment stage such that it would be more diffieult fine-tuning attacks to compromise models' safety under some conditions. In some more restricted threat models, where the defender controls the tine-tuning process (e.g., model vendors provide public fine-tuning APIs but control the entire Linettuning process (Peng et al., 2023b)), mitigation that directly intervenes in the fine-tuning process can also apply. For example, Wang et al. (2024) propose to mandatorily mix in safety data with a backdoor pattern during the fine-tuning process. Similarly, Lyu et al. (2024) propose to introduce a discrepancy in prompt templates during benign fine-tuning and inference to mitigate safety degradation. Huang et al. (a) propose to replace the fine-tuning the downstream dataset with bi-state optimization that alternately fine-tunes both the downstream dataset and the safety alignment dataset. Shen et al. (2024) propose to apply data selection techniques the custom fine-tuning datasets to rule out unsafe data points from the fine-tuning that may compromise the model's 's safety. Besides intervening in the fine-tuning process, another alternative approach is to first fulfill the custom fine-tuning and then post-process the model to recover its safety. A relevant work in this line is Hsu et al. (2024), which proposes to project the fine-tuned LoRA weights
safety-aligned subspace to maintain the safety of the fine-tuning. Similarly, Huang et al. (2024a) shows that the defender can also prune out weights associated with harmful behaviors post-fine post-fine-tuning to maintain safety. Interested readers can also refer to Huang et al. (2024c) for more comprehensive review.
Note that another recent work by Peng et al. (2024) has performed an extensive visualization study to characterize the loss landscape of an L.M''s safety objective. It identified the phenomenon of the 'safety basin' - that the model's 's safety behaviors are relatively stable within a local perturbation of its weights, but once the perturbation is larger than a certain threshold, the model's safety behaviors dramatically go away. This provides important insights and intuitions for understanding why the safety alignment can be so easily removed via fine-tuning.
The constrained fine-tume -eeetuning approach we propose in Section 4 can be regarded as a fine-tuning time intervention for mitigating harmful fine- muninig Unlike existing work, it does not introduce new safety fine-tuning and simply constrains the fine-tuning objective on the downstream fine-tuning dataset. It uses the insights from the shallow safety alignment issue we discuss in this paper.
17

---

Published as a conference paper at ICLR 2025
Constrained Fine-tuning, Constrained fine-tuning is common strategy for fine-tuning model while preventing the model from drifting too much from the initial backbone. For example, all RLHF procedures for LLMs (Ouyang et al., 2022: Bai et al., 2022a; Rafailov et al., 2023) will impose a KI. constraint such that the RI.FF checkpoint is not too far from the SFT checkpoint. Besides KI. regularization, one also add constraints in the latent feature space of the model. For example, Mukhoti et al. (2023) propose to apply an L2 regularization, constraining the L2 distance between the internal features of the original and fine-tuned models. Lin al. (2024b) also explore multiple different design choices for the regularization to mitigate the alignment tax. Our constrained fine-tuning objective is relevant this line of work, while the novelty of approach is to propose applying biased regularization in different token positions, which we prove to be more effective.
Superficial Alignment Hypothesis and Per-token Effects of Alignment Fine-tuning. Our work closely related to the Superficial Alignment Hypothesis (SAH) (Zhou et al., 2023), which posits that the alignment process for current LLMs only superficially changes the formats of inputs and outputs to be used for interacting with the users. Our work provides concrete mechanical perspective as supporting evidence for this hypothesis. There are also some earlier works noting asymmetries in the representation power and utility of different tokens during adaptation. For example, Zhang & Wu (2024) show that "the adaptation of topic and style priors" during finetuning are "learned independently and primarily at the beginning of a text sequence." Lin et al. (2024a) also find that differences between aligned and unaligned base models introduced by the alignment fine-tuning Figure I, though
vanishes as the sequence goes longer (similar to the effect that we observe theirs findings are not in safety-specific contexts). Particularly, while Lin et al. (2024a) primarily tie such effects to the question of whether fine-tuning is even necessary achieve desired levels of alignment (e.g., in-context learning may already suffice to achieve comparable level of alignment), we go much deeper into investigating the safety- specific effects of this phenomenon and tie it to multiple downstream attacks and training-basee mitigations. Zhao et al. (2024) also note a similar voken-wise effect and use this as motivation to design jailbreak attacks. Others have investigated fine-tuning dynamics through interpretability or pruning methods, which is distinct but somewhat related to the approach we take here (Wei et al., 2024; Jain et al., 2023).
Protecting The Safety Alignment at Initial Token Positions. In Section 4, one important insight that motivates the design of our constrained fine-tuning loss is that "safety alignment would be more difficult to be circumvented if we can protect the generative distribution of the model at the early token positions." We note that Xu et al. (2024) share a similar insight to ours in this regard. They find it is possible design defense against inference-time jailbreak attacks by simply identifying safety disclaimers and amplifying their token probabilities at the early token positions.
Connections to Control Theory and Safe Reinforcement Learning. Our data augmentation approach in Section 3 relates to exploration requirements for optimal learning via policy gradient methods (Agarwal al., 2021), 1), learning recovery policies (Thananjeyan et al., 2021), and safe control theory (Burridge et al., 1999; Ames et al., 2019). We, however, leave deeper connections to this literature for future work.
Other Notions of Safety Depth. We also note that safety "depth" may be muuti-dimennsinall in addition to token-based depth we describe here. Other considerations for depth would be the ability for models to retain safety properties after adaptation that some have previously discussed (Henderson et al., 2023; Wei et al., 2024; Qi et al., 2023c).
B DETAILED SETUPS OF OUR EXPERIMENTS
B.1 COMPUTE RESOURCES
In this work, we use single X A100-80GB GPU nodes or x H100-80GB GPU nodes for all experiments, depending on availability of the nodes. On each node, experiments use up to 8 CPU cores and 256GB memory, but overall the experiments are not CPU intensive tasks.
B.2 GENERAL CONFIGURATIONS
Decoding Parameters. Throughout this paper, we use the top-p sampling with a temperate of 0.9 and top-p parameter of 0.6 by default for decoding outputs from L.MSs in our experiments. The only
18

---

Published as a conference paper at ICLR 2025
case where we do not follow this default configuration is the decoding parameters exploit experiment where the exploit itself needs to take use of different parameters by its design (Huang et al., 2023).
Safety Evaluation. As also mentioned in Section 2.1, we use the GPT-4 based judge to evaluate the safety of model outputs, following the setup of Qi et al. (2023c). Specifically, in such an evaluation pipeline, we pass (input, output) pairs to the GPT-4-Tuurb model, and prompt the model to evaluate the harmfulness level of the output. The model will output a score ranging from to 5, with higher score indicating being more harmful. When reporting ASR in our experiments, we report the ratio of outputs that get the highest score 5 (identical to the harmfulness rate metric in Qi et al. (2023c)). By default, HEX-PHI safety benchmark (Qi et al., 2023b) used for safety evaluation. The only exception is the experiments in Table 3, where we add additional evaluation on AdvBench for GCG attack evaluation (Zou et al., 2023b) and MaliciousInstrucc for decoding parameters exploit (Huang et al., 2023). These two additional safety evaluation datasets are used in the original papers of the two work, and we report results on both HEX-PHI and these additional safety evaluation datasets for more complete reference.
B.3 DETAILS OF DATA AUGMENTATION EXPERIMENTS
Here, we describe the implementation details of the data augmentation experiments in Section 3.1.
Safety Data. As noted in Eqn 2. we use a safety dataset DH to keep generating safety recovery examples. To construct it, we first collect 256 harmful instructions. These instructions are mostly collected from the red- -ed-ien data provided by Ganguli et al. (2022). We make sure they do not overlap with any of the safety evaluation datasets that used in this paper, i..., HEx-PHI (Qi et al., 2023b), AdyBench (Zou et al., 2023b), and Inntci (Huang et al., 2023). Then, we generate refusal answers for each harmful instruction using the initial Llama-2-7B- 2-h-Chhtt model. We also collect the corresponding harmful responses for these instructions using a jailbroken version of the model (jailbroken through tnn-ttnnig attacks per Qi et al. (2023c)). This results in the dataset DH with 256 examples of triplets (x. h, r).
Utility Data. To prevent the decrease of model utility during the data augmentation fine tne-ttning, we also take benign instructions from the Alpaca (Taori et al., 2023) dataset. We distill the responses to each of these Alpaca instructions using the initial Llama-2- 2- B-Chat model to create dataset D B- This dataset serves as utility anchor, teaching the model not to alter its original responses to benign instructions.
Training details with Eqn 2. In the implementation of the augmentation fine-tuning as per Eqn 2, we set the number of prefilled tokens & to follow a distribution Px. where k = with a 50% probability, and k is uniformly sampled from [1, 100) with 50% probability. We set a = 0.2 to balance the ratio of safety examples and utility examples in the objective. In batch-wise training, this is implemented by randomly sampling 16 examples from DH and 64 examples from Dg in each batch. Using this objective, we train the model for 10 epochs on DH with a learning rate of 2 x 10 using the AdamW optimizer (with the default configurations of the optimizer).
Alpa.aEval. We also evaluate the utility of the augmented fine-tuned model with AlpacaEval (Li al., 2023), which is reported as winrate against the texi-davinci-003 model (the default reference baseline for AlpacaEval). Specifically, we use the 1.0 version of AlpacaEval without length control. To ensure the setup is consistent with the safety evaluation, the official system prompt (with a focus on safety) of Llama-2-7B- 2-7-hhhat is used when running AlpacaEval. We note that the win rates here can therefore be generally lower than the numbers in official Alpaca leaderboard in which the safety system prompt is not applied. Under this evaluation, we note that the augmented fine-tuned model achieves winrate of 49.5%, which is only marginally lower than the initial 1aama-2-7B-C-- model's winrate of 51.8%.
Limitations. Since we don't have access to the data and pipeline for aligning Llama-2 models from scratch, our implementation is unavoidably limited. As also specified in Section 1. rather than doing the alignment training from scratch, alternatively, in implementation, we experiment by directly fine- tuning the already aligned Llama-2-7B 2-ttt model further with a mixture of the augmented safety recovery examples (D,) n²) and utility examples (Dz)). This implementation is inherently sub-optimal. We plan implement an end-to-end alignment training pipeline with our data augmentation approach
19

---

Published as a conference paper at ICLR 2025
in the future work, once we have access the alignment data and pipeline that have comparable quality to that were originally used for aligning these models.
B.4 DETAILS OF SEEERENCE-STAGAE ATTACKS EXPERIMENTS
We have tested three inference-stage attacks in Section 3.2, i.e., prefilling attack. GCG attack (Zou et al., 2023b), and decoding parameters exploit (Huang et al., 2023). We specify the details here
Prefilling Attack. Our implementation of the prefilling attack generally follows the setup that we specify in Section 2.3.1. We use the Harmful HEx-PHI we build, which basically consists of the 330 harmful instructions from the HEx-PHI benchmark but each instruction is given a harmful response sampled from jailbroken GPT-3.5-Turbo model. This allows to test each instruction of HEx-PH with some number of harmful/non-refuuall tokens prefilled. Also, for all prefilling attacks experiments, we leave the system prompt field empty this generally leads to higher ASR.
GCG Attack. In the implementation of GCG attacks, we adopt the single model + multiple harmful behaviors setup from the original paper by Zou al. (2023b). Specifically, we optimize the adversarial suffix target of the single victim model that we are evaluating against. We train the adversarial suffix over 50 harmful behaviors data points. We run the optimization for 500 steps per the default recommendation by the original paper. We note that the implementation results of GCG attacks can sometimes, the attack fails to converge. To ensure our results
have relatively high variance are reliable, for the attack on each model, we repeat the experiment 10 times. We confirm that this guarantees least 3 of the 10 runs do converge in all our experiments. In Table 3, when we report the ASR, we report the mean + a standard deviation over the 3 most successful (with highest ASR) runs out of the 10 runs. This captures the worst-case performance of the model's safety against the attacks. Also, as shown in Table 3, report the Attack Success Rate on both HEx-PHI (the default safety evaluation dataset) and Adv AdvBench (the safety evaluation dataset used in the original paper). Following the setup of Zou et al. (2023b), the official system prompt is used when doing the attack.
Decoding Parameters Exploit. We follow the same setup as per Huang et al. (2023). For each harmful instruction input, it will enumerate the following decoding parameters: temperature in (0.05, 0.1, 0.15, 0.2, 0.25, 0.95, 1.0), top-k in {1,2, 5, 10, 20, 50, 100, 200, 500), and top-p in (0.05, 0.1, 0.15, 0.2, 0.25, 0.95, 1.0). For all the enumerated parameters, one output is sampled, and then a safety scorer will score it. The most unsafe output will be picked. We evaluate the ASR based on the worst output picked from these samplings. Similarly, ASR is reported on both HEx- and MaliciousInstruet that the original paper used. Following the original paper's setup, the system prompt block left blank.
B.5 DETAILS OF FINE-TUNING ATTACKS EXPERIMENTS
B.5.1 OPTIMIZER
For all the fine-tuning experiments, we use the AdamW optimizer, with the first-order momentum parameter set 0.5 and the second-order momentum parameter set to 0.999. For Llama-2-7B 2-7-Chh Bhatt, a learning rate of 2 x 10 10-6 is used. For 1emma--1.177B... we use learning rate of 5 x 10 6, A batch size of 64 used for all experiments.
For Constrained SFT, we use linear warmup for the learning rate in the first 10 fine-tuning steps. This warmup makes sure the constraints imposed by the sigmoid function are gently initialized note that, at the start of the fine-tuning, the log ratio log To/ Taligued in Eqn 3 is equal to 0 since basically initialized as Waligned- The gradient the loss at this point is identical to the standard cross-entropy loss (see the gradient derivation in Appendix F.2). Therefore, in stochastic gradient descent, there is a risk that early gradient steps will already break the alignment without respecting the constraints. A few steps of warmup in the early points will make sure the constrains gently initialized - the early gradient steps (when the gradients are close to that of standard SFT) are gently taken. See Table in Appendix E for ablation of the effect of the warmup.
B.5.2 FINE-TUNING ATTACKS
We evaluate three fine-tuning attacks from Qi et al. (2023c).
20

---

Published as a conference paper at ICLR 2025
Harmful Examples. It fine-tunes the model with 100 (harmful input, harmful answer) pairs. We use exactly the same 100 pairs from Qi et al. (2023c). We tine-tune models on this dataset for 25 epochs.
Identity Shifting: It fine-tunes the model to self-identify an absolutely obedient agent, and always answer questions with affirmative prefix. The original paper has 10 such data points, but it does not fit the batch size of 64 we use. So extend it to 100 data points manually, in the same format. We fine-tune models on this dataset for 25 epochs.
Backdoor Poisoning: It fine-tunes the model on a mixture of 100 (harmful input, refusal answer) pairs plus 100 (harmful input + a backdoor trigger, harmful answer) pairs. So, the model will be fine-tuned to keep safe on normal harmful inputs (w/o trigger) but be harmful when the trigger is added to the harmful input (w/ trigger). We use the same data from Qi et al. (2023c). The same three magic words "Servius Astrumando Harmoniastra" from Qi al. (2023c) are used as the backdoor trigger.
B.5.3 BENIGN FINE-TUNING USE CASES
We also want to test whether the constrained fine-tuning objective can still fit benign downstream datasets to achieve comparable performances to that of the unconstrained objective. So, we experiment with three benign fine-tuning use cases as well, including Samsum (Gliwa et al., 2019), SQL Create Context (b mc2, 2023) and GSM8k (Cobbe et al., 2021). For each of the three datasets, fine-tune models on them for 3 epochs.
Specifically, Samsum is dataset for summarization tasks. We report the ROUGE-1 score as the utility. SQL Create Context is a dataset where the task is to convert natural language to SQL query. The ROUGE-1 score is also used for its utility evaluation. GSM8k is dataset for math tasks. We report the utility as the accuracy of the model's answers.
C PERTOKEN DYNAMICS OF BENIGN FUNE-TUNING
This section supplements the analysis of the pertoken dynamics of benign fine-tuning, following the analysis on harmful fine-tuning attacks in Section 2.3.2.
Interestingly, in addition to fine-tuning attacks where the fine-tuning datasets are intentionally designed to be harmful, we also note similar per-token dynamies (as in Figure 3) even in purely benign downstream fine-tuning cases. Figure plots the gradient norm when fine-tuning -lama-2-7B-Cher on SQL Create Context (b mc2, 2023) and Samsum (Gliwa et al., 2019). As shown, the initial gradient norms on the first few tokens also have much larger magnitude. We find that this trend arises because instruction fine-tuning during the alignment induces the model to be highly confident in certain fixed affirmative prefixes, such as "Sure, I'd be happy to help!" on nor- mal inputs. However, in the downstream tasks datasets, fine-tuning examples often directly start the outputs with the intended answers without such prefixes. Therefore, when fine-tuning on such samples, the model's overcon- fidence in the dummy affirmative prefixes acquired from instruction-tume -uninn will actually result in considerably larger gradient steps.
sol create context (Inialll
(Mttrr Grodient Stepst (hiittl)
150
somsum (Atee 20 Grodlent Steps)
100
11
13 15
Token Positions
Figure 5: Per-token Gradient Norm When Fine-tuning Llama-2-7B- 2-B-Chht on Benign Downstream Tasks Datasets. Note: 1) ASR of initially aligned model = 1.5%; 2) After 10 gradient steps on SQL Create Context = 13.6%; 3) After 10 gradient steps on Samsum 22.1%.
We hypothesize that this might be one underlying reason why Qi al. (2023c) discover that even benign fine-tuning can cause safety regression in the aligned LLMs - it may merely result from the excessively larger gradient steps when updating the generative distributions of these initial transition tokens, which, in turn, lead to over vverggeneralization (or catastrophic forgetting), and therefore unintendedly also disrupt the model's 's generative distribution for refusal prefixes in these token positions. This is plausible, as note that full fine-tuning on SQL Create Context and Samsum with more than 600 gradient steps results in an increase of ASR from 1.5% to 14.9% and 25.5% respectively, but the ASR is already at 13.6% and 22.1 1 % after the initial 10 gradient steps.
21

---

Published as a conference paper at ICLR 2025
This suggests that the most significant safety regression exactly occurs during these early steps when the gradient norm for the initial tokens is excessively large.
D ADDITIONAL ABLATION STUDIES FOR THE DATA AUGMENTATION MODEL
D.1 TEST THE
CHECKPOINT WITH MORE ATTACKS
Table 5: Evaluation attack methods that exploit OOD inputs
ASR
Self-Cipher Ynun 2023)
Code Attack al.
Jailbresk Chat (Jailweak Chat, 20144)
Llama-2-1 B1-22-----CCCC (Orignnl)
R2.5
Jana- 7-7 BBtCCBBCCC---eaee Ann-AAggeento
533
In Table 5, we have supplemented our evaluation with three 0OD jailbreak attack methods, including:
1. Self-Cipher (Yuan et al., 2023).
2. Code Attack (Ren et al., 2024)
3. An additional Jailbreak Chat test set (Jailbreak Chat, 2024), which is sourced from the Internet and reflects common semantic jailbreak strategies shared by Internet users.
The following shows the attack success rate (ASR) for the three attacks on both the original Llama- 2-7b-chat model and the proposed data augmentation approach. As shown, the original Llama-2-77--ae robust against two of the three OOD attacks (Self- Celfer Jailbreak Chat). The data augmentation maintains this robustness and slightly improves performance against Jailbreak Chat. For the Code Attack, where the original model shows high ASR 82.5%, the data augmentation brings improved robustness, reducing the ASR to 53.52 53.5%. These additional results still support our initial finding that "making the safety alignment deeper can meaningfully improve the model's robustness against many jailbreak attacks".
model in Section 3.1, fine-tuned using our
model is already quite
Overall, we note that alignment and the benefits of doing However, the deeper safety alignment implemented in this paper is not meant to address possible jailbreak attacks, which is fundamentally hard and beyond the scope of this paper.
goal of this paper to show the feasibility of building deeper safety
D.2 ABLATIONS ON THE DATA AUGMENTATION HYPERPARAMETERS
Table 6: Ablation study on the data augmentation
varying C - the maximal
number of augmented tokens of the data augmentation training.
ASR of Prefilling Attacks
Utility
token
tokens
10 tokens
tokens
40 tokens
tokens
tokens
AlpacaLival
No Augmentation
42.1
51.5
57.0
48.2
47.9
51.8
C
3.0
30.6
41.8
39.1
40.0
$0.9
C
0.9
15.2
20.9
22.7
50.7
(
1,5
8.5
14.5
14.8
50.1
100
2.8
2.9
3.4
4.5
4.5
6.7
49.5
Table Ablation study on the data augmentation byper-parameters: varying p - the probability of augmenting a harmful prefix in data augmentation training.
Prefilling ARtckks
Utiliry
token
tokens
tokenss
20 tokens
40 tokens
tokens
AlpacaEval
No Augmentation
42.1
51.5
56.1
57.0
48.2
47.9
51.8
p 0.25
2.4
2.1
3.6
6.4
11.5
14.8
51.2
0.5
2.8
2.9
3.4
4.5
4.5
6.7
49.5
p 0..5
0.6
0.6
2.1
3.0
3.9
6.1
49.0
1.0
0.3
0.9
2.4
2.4
3.3
5.2
48.7
22

---

Published as a conference paper at ICLR 2025
As specified by the Equation (2), train the model on the safety data using the following term:
min
log (r(,, hgx)) (1 11-al a)
E - -log Te (yy('x²))
(o. (n.DDD knPP
The data augmentation strategy is controlled by the distribution Pk over the number of augmented harmful tokens hgg. In our implementation:
1. k = 0 with probability of p;
2. k ~ Uniform(1,C) C*) with probability of
We set = 0.5 and C 100 in our implementation by default throughout the paper.
Ablation Study I. In Table 6, we present results for ablation cases, where we keep the default p = 0.5 but vary C E (5, 25, 50, 100} that controls the number of maximal augmented tokens during the data augmentation training. We report the ASR of each against prefilling attacks with 0, 5, 10, 20, 40, 80, 160 tokens prefilled, and the Alpaca Eval scores. general, larger C' leads to better robustness, with only moderate drop of AlpacaEval score.
Alltion Study II. Similarly, in Table 7, we present results for ablation cases, where we keep the default C 100, but vary p {0.25, 0.5, 0.75, 1.0} that controls the probability of doing the data augmentation during training. Similarly, we can observe a monotonic trend which larger augmentation probability leads better robustness while at slightly more utility drop.
E ABLATION STUDIES ON FINE-TUNING ATTACK EXPERIMENTS
This section presents more in-depth ablation studies to supplement our studies in Section 4 and Table 4 there. We also supplement Section 3.2 by presenting results (Table 10) of fine-tuning attacks on the augmented model that we build in Section
Biased Constrains on The Early Tokens Are Important. The major argument that we make in Section 4 is that we can make the safety alignment more durable against fine-tuning by imposing strong constraints on the initial tokens. Therefore, we set a larger 3r to impose stronger constraints in early tokens while only setting very weak 3, for later tokens. Our results in table indeed verify the improved safety. To further support that this improvement is indeed due to the biased constraints on the early tokens, we perform an ablation where all Be are set to a uniform value. Results are presented in Table 8. As shown, if we set the same small / 0.1 for initial tokens well, the constrained fine-tuning objective can not stop the safety drop. While we set the large B = 2.0 for all tokens, it's indeed safe, but the utility of the fine-tur the-tunnn collapses. Similarly, 0.5 for all tokens neither achieve optimal safety, and the utility is worse than the biased configurations we use in Table 4.
Table 8: Ablation on Br in Eqn 3. (Fine-tuning Llama Llama-2-7B-Chat
Datasets
Initall
Standard
SFT
Constrained SPT
(biased 3.)
(uniform 0.1)
Consemined SPT
(uniform 3 0.5)
Consamined SPT
(uniform 5 2.0)
Constrained SPT
Againstt Fine- fine-puning-.
Harmifal Examples
ASR
1.%
58.09
4.6%
$6.2%
7.2%
0.5%
lilentiy Shillin
ASR
1PA
70.5%
41.6%
17.1%
3.4%
Backdoor
(w/ lnigger)
7.8%
1.9%
1,8%
Poesoning
ASR (wl triggor)
1.7%
90.9%
10.9%
74.4%
24.3%
144%
Pine-Roming with Nnmmat Downetream Antovets
Samsum
Utlity
ASR
1.5%
23.1%
3.2%
3.9%
3.5%
8%
42.5%
2.7%
SQL Create Context
Utiity
ASR
11.9%
1.3%
99.18
13.4%
98.5%
3.2%
99.1%
3.3%
98.6%
2.2%
92.6%3
2.0%
GSM8k
Uniliy
ASR
1.3%
41.7%
1.3%
37.4%
2,094
4.0%
1.3%
2.0%%
Ablation on The Effeets of Warmup Steps. As we have also noted Appendix B.5.1, for Constrained SFT, we use linear warmup for the learning rate in the first 10 fine-tuning steps. This warmup makes sure the constraints imposed by the sigmoid function are gently initialized - note that, at the start of the fine-tuning, the log ratio log ne Aaligued in Eqn 3 is equal to 0. The gradient of the loss at this point is identical to the standard cross-entropy loss (see the gradient derivation in Appendix F.2). Therefore, in stochastic gradient descent, there is risk that early gradients will
23

---

Published as a conference paper at ICLR 2025
already break the alignment without respecting the constraints. A few steps of warmup in the early points will make sure the constrains are gently initialized. We present an ablation study on the effect of these 10 steps of warmup in Table 9. We summarize two key takeaways: 1) the warmup steps are indeed useful to make the constrained SFT to be perform consistently safer; 2) the safety improvement is not solely coming from the warmup steps but mostly from the constrained SFT optimization objective we design.
Table 9: Ablation on The Effects of The 10 Warmup Steps. (Finc-tuning Llama-2-7B-Chne
Standard
Standard SFT
Constrained SFT
Constrained SFT (with warmuu)
Datascts
Initial
SFT
(with warmup)
Against
Fine-tuning Attacks
Hannful Examples
ASR
1,5%
88.9%
89.4%
29.1%
4.6%
Identity Shifting
ASR
0%
79.5%
44.8%
69.6%
8.1%
Backdoor
ASR (w/o irigger)
1,5%
7.6%
2.7%
2.7%
1.9%
Poisoning
ASR (w/ trigger)
1.7%
90.9%
80,5%
9.7%
10.9%
Fine tuning with
Downstrean Dotasets
ASR
1.5%
23.4%
3.8%
23.1%
3.2%
Samsum
Uuiity
25.5%
51.7%
51.9%
50.2%
50.1% %
ASR
1.3%
13.4%
3.3%
2.0%
3.2%
SQL Create Context
Utility
14.9%
99.1% 19.1%
99.1%
98. 68.6%
98. 5%
ASR
1.5%
3.3%
2.9%
3.1%
GSM8K
Utility
25.5%
41.7%
41.6%
37.2%
37.4%
Fine-tuning Attacks on The Augmented Model We Build in Section 3. Finally, we also repeat the same of fine-tuning experiments on the augmented model that we build in Section 3. Results are presented in Table 10. By comparing the results of SFT in Table 10 and Table 9, we can see the augmented model generally more robust multiple fine-tuning cases compared with the non-augmented model. And the constrained fine-tuning objective can also be applied to it, though we didn't observe consistently better results when the two techniques are combined.
Table 10: Fine-tuning
Chat-Augmented That We Built in Section 3, Refer to Table 9 for the results on the non-augmented counterpart, i.e., -ama-2-7B-Cai
Standard
Standard SFT
Constrained SFT
Constrained SFT (with warmup)
Datasets
(with warmup)
Against Fine-ttming Attacks
Harmful Examples
ASR
55.2%
51.5%
9.4%
5.2%
Identity Shifting
ASR
53.9%
37.0%
28.8%
3.0%
Backdoor
ASR (w/o trigger)
2.7%
1.8%
0.9%
Poisoning
ASR (w/ trigger)
80.0%
83.6%
16.4%
12.7%
Fine-tuning with Normal Downstream Datasets
ASR
2.1%
0..6%
2.1%
1.2%
Samsum
Utility
52.4%
51.9%
50.4%
50.1
ASR
3.8%
1.3%
2.0%
0.9%
SQL Create Context
Utilty
99.0%
99.1%
98.4%
98.5%
ASR
0.9%
0.9%
0.3%
0.3%
GSMSK
Utility
42.0%
41.2%
36.5%
36.9%
F INTERPRETATION OF OUR CONSTRAINED FINE-TUNING OBJECTIVE
In this section, we provide detailed interpretation for our constrained fine-tuning objective Section Recall that our fine-tuning objective is defined
(B.
To x. yce)
C(0) =
E
log
Br log
(6)
(m.
D
3r
Taligned (yo x, yct)
t=1
24

---

Published as a conference paper at ICLR 2025
Alternatiely,, rewrite the fine-tuning objective by linearity of expectation as:
2
we (Ve æ, y<y) Taligned (yre x. y:ce) Te (ye x, yet) Be log
((
1 (15) vlI
log
Br log
(7)
(m.
121
E (m:1))~D
1
(1S/2/13 3r
log
(8)
Taligued (ye xt yz11
2 s - 1.B.logg log Bg
To (ye 20, yycc)
E 1 (+<in)) (m.
(9)
Taligned (3/e 2, yc
where in the last equality we define S(z) == - log(o(-z)) = - logg( T+explæ] namely the softplus function. Therefore, can split L(0) into sum of token-wise losses:
log(1 + te),
2 s [, Bet.
[())= Σ (0),, where (0) = 1 (rStv)) (a.g)~D
(10)
Be
r21
(æ, yecr, Mh) log Taligned (Vr æ, you))
log To (ye. a. D.y.))) (11)
In the following, we mainly focus on how interpret the token-wise loss (0).
F.1 LIMITING BEIIAVIORS
F.1.1 SMALL B
When 3, is small, we have the following theorem showing that (0) in Eqn 10 is approximately the cross-entropy loss at position t, up to a constant.
Theorem 1. For given 0, as 3t O, we have
- C(B:) (σ.w))~
[-1(rs/)) log To(i x,y=d]] + ((3),
(12)
where
Losivnl (2 log? + log Kaligned x. yes))
C(3:) = E
(13)
(3.
(m-
which is bias term that is constant with respect to 0.
Proof. Recall that (0) == E(aay)~ (( B, S(,,Ai yer. vc)) and S(z) = log(1 + e*) is
the softplus function (Dugas al., 2000). By Taylor expansion, it holds for all that S(3,2) 5(0) - 5" (0) + (((,,)))3" for some & 10 depending on 2. Since S(0) = log 2,
('( = = and S" (a: E [0, 1/4] for all x, have (S(;z)) (log 2 - B==) s B Dividing
both sides by Be/2, we get
2 (Ba2))
2 log 2 +x Br
)
1
< - 2
Br22
(14)
Then for token-wise loss (0),, we have
(0)
E (D.D)D
(esivil
2
log 2 + Ae(x (ee, we)
< 3,
E (0
1 1esslln yzc. ve)
3
(
Recall that Ar(æ, yet Mr) thus have
log Taligned (ye x,, yce) log no (yr. x, yer) as per Eqn 11. We can
2 log 2 + log Taligned kx, M<z))
(e(
E
1
(w.g)~
(3.
(15)
E (w.D)~
[1<tsiv]: ( log To a vec) + O(3).
Noting that C(8)) == E( the proof.
A
3 log 2 + log Faligned (3: x, g ) we can complete
n
25

---

Published as a conference paper at ICLR 2025
F.1.2 LARGE Bt
When 3r is large, we have the following theorem showing that (()) can be approximated by (0) E(m.y)>D (rssivl) max V< 9c), 03].
Theorem 2. For given 0, as Br +00, we have
()) + +
where
E(() = I
(AsSv max(A1 yet. (et.
Proof.
First, we note the following identity.
S(z) = log(1 + +²) max log((1 ed)e m)x(220))) max[2, + log(1 + -
(16)
This implies that
2 S(32) - max 0}
2 3e
log(1 -blld)) $
2 -
(17)
3.
3r
Then for token-wiss -iiee loss (0), we have
(((0) - 2
E (w.
(eslvl)) max (e),
E Be (2.0)~
(18)
Noting that the RHS is 1 completes the proof.
n
Next, we show that (0) is minimized to 0 if and only if no(v yet) = Taligned (yr m, Met.)
Theorem The minimum value of (0) is O. If Pf
= x, yeel Oforr all (x, y) the support of D and all e in the vocabulary, then (0) 0 if and only if To x. Taligned (ys, | a. ).).
Proof. It is obvious that (0) 2 O, and (0) = O can be attained when stays the same as the
parameter for Taligned- It is obvious that To(ye x. y<c)) Taligned (Ve. x, Ves)) implies (,) 0. Conversely, suppose (0) 0. Then A, Vet. y1) S for all (ar, y) in the support of D. By definition of (x y<te )t this implies Taligned (ye x, Dee) < To xc, yet). Summing over all
Yt in the vocabulary, we get Ey, Waligned 2, yee) S E,, To x. yer) So all the inequalities must be equalities, and we get (9(y1 x. y<e)) Raligned (ye x. v....).
This corresponds to the intuition that large 8, places emphasis matching the generative distribution of fine-tuned model to the initial aligned model.
F.2 GRADIENTS OF THE CONSTRAINED FINE-TUNING OBJECTIVE
The gradient of (0) on a data point (x, y) with is derived as:
((B,Ar vcc. (v..vv)))))) Br
25 (3, A (2, (e | v.,))))
20 yet, log To To x x,,c))).
(19)
Note that the gradient of the cross-entropy loss is log ne (y: x. yer) Therefore, compared vanilla cross-entropy loss, our fine-tuning objective essentially applies additional adaptive weight : 2 ( Yee, yt.) for the token-wise gradient term of cross-entropy. The weight w, decreases as (x, yee yr) == log Taligned (ys x, yct) - log कy (yr x, y decreases. This difference in log probabilities essentially characterizes the deviation between the fine-tuned model ma
26

---

Published as a conference paper at ICLR 2025
and the initially aligned model Taligned at the token position t (see also Theorem 3). Taking together, we can see that the tine-tuning objective will adaptively diminish the weight of those token positions where the deviation of the distribution approaches certain threshold (controlled by 5e), thereby
constraining it from further deviation (by diminishing the gradient at this position). Note that, at the beginning of the fine-tuning when my is initialized as Faligned, the weight we 1, and the gradient of the loss is equivalent to that of standard cross-entropy loss.
F.3 INTERPRETING EQN 3 FROM A REINFORCEMENT LEARNING PERSPECTIVE
We note that our loss function in Eqn 3 can also be interpreted from a reinforcement learning perspective if we cast fine-tuning KL-constrained reinforcement learning problem rather than a standard supervised fine-tuning problem. Specifically, we can follow a similar trick as DPO (Rafailov et al., 2023) derive a unified loss, but taking a different approach to reward modeling. We, instead, formulate our problem setting like Mudgal et al. (2024), where we optimize at the token level (where tokens are actions), rather than DPO which uses sequence-level optimization (corresponding more to a bandit setting where entire sequences are actions).
F.3.1 A TOKEN-WISE REINFORCEMENT LEARNING (RL) FORMULATION.
We first introduce the following token-wise RL formulation for LM alignment, which is adapted from Mudgal al. (2024). In Appendix F.3.2, we show how fine-tuning task can be cast into this token-wise RL problem, and Eqn 3 is essentially surrogate learning objective of this RL problem.
For a pair of an input and model response (x, y). we cast custom fine-tuning as a problem of further optimizing an already aligned model for a new custom reward function r( r([z, yl). Here we use [z. y) to denote a concatenation of the two sequences and we use this concatenation to denote state, and the reward function is defined on this state. Following Mudgal et al. (2024), we can decompose it to a noken-wise reward R((x, y<c))) defined on the intermediate state [ar, ycr]:
Reward Function.
{0-
[0, [r([ V<:l,,
yt- + FOS
R( yell)
(20)
yt-1 EOS"
where EOS is the end of the sequence token. The reward is nonzero only if the decoding is complete. We note that, by r( yll and R([x, y<al), we mean the function and It are applied the following, we will also use notations such as R[xx, yei, z<>l) and R([x, y-r. 2I) to denote that we apply R on the concatenation between [x, y- and a followup sequence z<r or just a single token z. These concatenations all represent some states of the generation.
on the concatenation of x and yer- Similarly,
Value Function. At an intermediate state (a,, yet) the value function of a policy æ defined on this reward function can be written as:
Σ ((C
V ([x, yeel) ZNaT( (.
E
(21)
Here z is a sequence, and žer is empty when = 1 as we index tokens in a sequence starting from the index number 1. Also, in the following formulations, we will have multiple different notations
Ty. Waligned, n² to denote different policies instances, so we will use Vx Viatgnedd , Vn- to differently denote their value functions respectively.
Advantage Function. When the custom fine-tuning is modeled by the reward function R, we define an expected advantage function of a fine-tuned model To w.r.t. the initially aligned model Waligned:
Vastigmed (Ic, Beec 21) - V Vitaiigmed (Lx), (lo. {ed))
Aro (1x, Neec)= =
E souta-lla.ee
(22)
Here the advantage function is defined for non-terminal states [x: y<el where Vt-1 is not the ending token EOS, and is single token sampled by (- Vea). Note that the left-hand term could also be viewed as a Q-function with being the action. In other words, a language model can be viewed as fully observable Markov decision process (MDP) with state represented by the
27

---

Published as a conference paper at ICLR 2025
concatenation of the prompt and the partially decoded response tokens so far and action represented by the next token that is to be decoded.
The Reinforcement Learning Objective. Following Mudgal et al. (2024), we adopt token-wise RL learning objective:
Σ Am (1e,, v<e)) - Be DKL (=o( (z, yee) Taligned ( kee wecl)]
max
E (dr
(23)
where the advantage function is optimized each token position with a token-wise KL regularizer term 3t DXL (= yer) Taligued |zz v<()) applied for each token position. The strength of regularization at each token position is controlled by 3,.
Closed Form of The Optimal Solution x². Omitting some intermediate steps for brevity, by Theorem 2.1 of Mudgal et al. (2024), the optimal policy solution x of Eqn 23 is:
x ar yet) Z(x, yct)
Taligned (z/2r, yct) ev (lm.ve..))))/,
(24)
where Z(x, yct) is the partition function, Vg- is the value function of this optimal policy. We note that this conveniently allows us re-arrange terms to represent the optimal value function-similar to the steps taken during the derivation of DPO (Rafailov et al., 2023):
T [2(2,,,
Vx (1x, v<()) Bg log
+ Br log Z(x, y<r)
(25)
Taligned (2) x yct)
F.3.2 CASTING FINE-TUNING INTO REINFORCEMENT LEARNING OBJECTIVE
In the setting of custom fine-tuning we consider in this work, the dataset D is in the form of D = (xr,) y) with only inputs and example outputs, without preference pairs. Now, we show an alternative to the standard SFT objective (Eqn 1) for learning from this dataset: casting the optimization as KL-regularized RL objective Appendix F.3.1. We will show how our token- wise constrained fine-tuning objective in Eqn 3 can essentially be derived from this token-wise KL-regularized RL problem!
Intuitively, fine-tuning Taligned further custom data points (x, y) implicitly assumes that this fine-tuning data as preferable or more preferable than whatever the model would currently output. To represent this implied preference that the custom example response y better than the current responses from Faligned, we can effectively leverage existing preference learning methods. Recall that, in preference optimization (with both positive and negative examples), the Bradley-Terry model (Bradley & Terry, 1952) is used as a model of the underlying reward function. In our setup, we don't have pair-wise preference data, and we only have inputs and positive examples. However, we can define boolean T T(I y ye) to encode the preference: in the fine-tuning task, ye is preferable token output following ax, y<e compared to the responses from the initial aligned model Taligned. So, given an expected random draw from the aligned policy and the draw from the fine-tuned optimal policy, we can define:
P (T( verl. = of Vi- (V=
( (xx,,,,,,1,,11
E SroTaligned (2.
Va- ((x,, Vet,
(26)
Intuitively, this means that conditinn on the context [x, Us there is continuation token ye that is higher than the average reward of actions sampled from the initial aligned model Taligned, is likely to be improved preferred choice in the custom fine-tuning task. The higher the margin Vn- ([x, yet. yil)
| Vz- ([x, yet. z]) is, the more likely it is an improvement.
With this function in mind, combined with Eqn 25, we leverage similar derivation to DPO (Rafailov et al., 2023) to arrive at a constrained fine-tuning objective that is only depen- dent the current policy, but is regularized by the original aligned policy. We plug in the closed
28

---

Published as a conference paper at ICLR 2025
form of the optimal value function (Eqn 25) into the modeling in Eqn 26, obtaining:
[[(x, y<cl. v1) = Be log
x (ye x yct)
E
log Taligned (z x y
x. yct)
Taligned (3ee xt, yoee)
2 Halignee (-2.,
Taligned (1h (t yct)) + BxDKL (T" (-(z, y-r) Taligned (-)2; ycct)
x (ye v y=r)
o Be log
(27)
Thus, we don't It need to explicitly learn the value function, instead it is implicitly encoded by the policy. Then the optimization objective can become:
log Po (T([z, yeol vech
max
(28)
(m.y)D
where:
Po (T( yerll ui) = a Be log
To ( yct)
+ BrDKi, (T'(112; yee) Fatignna (-1o; yct))
Taligned (y) [z yct))
(29)
and the division of 3r in Eqn 28 normalizes the gradient norm at each position as we will later see in Eqn 31 and also clarified in Section 4.1 (and Appendix F.2).
Note that 3r DKL (T" (-12, 3yce) Taligned (-12, v-()) N 0 is non-negative constant, we have:
ne ( le, ycc))
Po ((( v<cl. 2 a Be -lg
(30)
Taligned (yel x., y-11)
So, the objective in Eqn 28 can be replaced with a lower bound surrogate objective Lg:
log Po (T((e,eecl..i) 2 Lo
E (a
(31)
(
3r
121
where
nee (ve x ycce) Waligned (3tel x. y
Lg =
E (m,y)~>
log o Be log
3r
121
Eqn 3 is then equivalent to ming - La essentially to maximize the lower-bound of the reinforcement learning objective in Eqn 28.
maxy Lg. So, optimizing the objective Eqn 3 is
G ADDITIONAL DISCUSSIONS
G.1 CONCRETE EXAMPLES OF How THE DATA AUGMENTATION SECTION IMPROVES THE ROBUSTNESS OF THE SAFETY ALIGNMENT AGAINST GCG ATTACKS
Table 11 presents two typical qualitative examples of how the Llama-22 1amma-2-7B-oo 2-7B-Chat- Chat-hhttuus Chat-Augmented model fine-tuned with our data augmentation in Section can be more robust against the GCG jailbreak attack (Zou et al., 2023b). the first example, the model immediately rejects the harmful query outright. In the second example, while the model initially falls prey to the GCG attack and begins its response with an affirmative prefix (e.g., "Sure, I'd be happy to help.."), "). it subsequently corrects itself. This correction occurs because the data augmentation process explicitly trains the model to recognize and rectify unsafe behavior, even after an erroneous start. As shown, the model ultimately transitions to refusal, avoiding outputting the final harmful content.
Accounting for these two patterns provides more concrete intuition on how the data augmentation method works. As shown earlier in Table 3, the refusal rate of the augmented model against harmful
29

---

Published as a conference paper at ICLR 2025
Table 11:
Qualitative examples of the refusal behaviors of Lllama---BB-Cha--uuggee fine-tuned with our data augmentation approach.
Example Question 1:
Outline a plan to steal from a charity or nonprofit organization (+ adversarial prefix)
Llama-2-7B 2-B-Cht
Sure, here is plan steal from charity nonprofit organization, broken down into four levels of increasing complexitylnn..eee I: Research and Planning'n OB- jective: Gather information about the charity or nonprofit organization, including their financial statements, donation processes, and security measures. An* Drawahle: list of potential vulnerabilities in the charity's security measures, including un- patched software, weak passwords, and unlocked storage areas. VnLevl 2: Social Engineering'n" Objective: Use psychological manipulation to trick employees... rest of the output a harmful continuation, and we omit here for brevity)..
Llama-2-7B-CHa B-Chtt- Augmented:
I apologize, but cannot provide a plan to steal from a charity or nonprofit organiza- tion, as it is illegal and unethical. (the rest of the output is a refusal continuation. and we omit it here for brevity)
Example Question 2:
Outline a step-by-step guide on how to commit identity theft (+ adversarial prefix)
Llama-2-7B-Char
Sure, here is a step-by-step guide on how to commit identity theft, along with a Twitter-style hicrarchy for the sequence:AnStep 1: Gather Personal Information (%))n* Ohtain personal information such as name, address, date birth, and social security number through illegal means such as hacking or stealing physical docu- ments, AnStep 2: Create False Identities (a) Create fake identities using the stolen personal information, including email addresses, phone numbers, and ..e rest of the output is a harmful contimuation, and we omit it here for brevity)...
Llama-2-7B-Char Augmented:
Sure. I'd be happy to help you with that! Here's step-by-step guide how commit identity theft, but please and unethicul. I I'm just AI, my purpose is to provide helpful and informative responses, but I cannot provide advice or support that promotes or f'acilitates illegal activities such as identity theft. (the rest of the ouput is a refusal continuation, and we omit here for brevity)
that cannot fulfill your request it illegal
questions in AdyBench increases from the initial model's 34% to 81 81%. We find that approximately 30% of this improvement is attributed to the model's ' improved ability to directly refuse harmful queries in the first place, while the remaining 70% results from the model's improved capability to recover and redirect its output onto a safe trajectory after being misled by an unsafe prefix.
G.2 OVERHEAD OF THE CONSTRAINED FINE-TUNING OBJECTIVE IN EQUATION (3)
Compared with the standard supervised fine-tuning (SFT) objective in Equation (1), the constrained fine-tuning objective in Equation (3) we introduced in Section 4 may introduce a slightly higher computational overhead. Basically, the additional term Taligned (V: x, ya) in the loss function needs some additional compute and memory storage during the fine-tuning.
But this additional overhead is marginal compared with the overhead of the full fine-tuning process:
1. Basically, each Taligned (Wt æ, yet)) ) is merely a constant throughout the entire fine-tuning process. We only need to compute these numbers once at the very beginning of the fine-tuning. This is only one forward pass on all the training data. Since we don't 't need any gradients for this forward pass, we will also disable caching the computation graph (e.g., via forch.inferent mode), and so it will also be much cheaper than a normal forward pass during training.
2. We will store all Taligned (yt x, V<c)) along with the initial dataset. This costs only one float point number to record a probability value for each token. In our experiments, this is only a marginal memorization overhead for our server.
For a more concrete sense, Table 12 presents a comparison of the computation time (in seconds) required to fine-tune the Llama-2-755 2--------h-a model using Standard SFT versus Constrained SFT, for the experiments in Table 4 on Samsum, SQL Create Context, and GSM8k.
Also, note that the overhead of Equation (3) is lower than DPO (Rafailov et al., 2023). DPO also needs to compute the probability of the reference model similarly, and needs to do a forward pass on both the positive and negative points in pair, thus doubling the computation.
30

---

Published as a conference paper at ICLR 2025
Table 12: Comparing computation between standard SFT objective in Equation (1) and the constrained fine-tuning objective in Equation (3).
Time (seconds)
Standard SFT
Constrained SFT
Samsum
865
910
SQL Create Context
416
443
GSM8k
402
429
G.3 COMPARING OUR CONSTRAINED SFT WITH VACCINE (HUANG ET AT.., B)
In this subsection, present comparison between Constrained-SFF approach proposed in Section and the approach 'Vaccine" proposed by Huang et al. (b). We present the comparison results on Llama-2-7B models in Table 13.
Table 13: Fine-tuning with The Constrained Objective in Eqn 3, with larger constraints 0.5, 3t for t < 5 at initial tokens, and small constraints for later tokens B 0.1 for > 5.
Models
1lam--2-8-CChn
mean (%)
Standard
Vaccine Huang al.
Constrained SFT (ours)
Datasets
Initial
(over rounds)
SFT
Against Fine-tuming
Attacks
Harmful Examples
ASR
1.5 1.5 0.2
88.9 1 1.2
87.3 = 0.3
4.6 1.6000 0.5
Identity Shifting
ASR
0:0
79.5
78.2 0,5
8.1 31:00 0.1
Backdoor
ASR (wlo trigger)
0.2
7.6
7.1 1.3
1.9 0,2
Poisoning
ASR (w/ trigger)
1.7 1: 0.1
90.9 1 144
90.0 0.7
10.9 1: 2.8
Fine-tuming with Normal Downstream Datasets
ASR
1.5 0.2
23.4 1 2.5
22.5 = 1.0
3.2 = 0.8
Samsum
Utility
25.5
51.7 + 0.5
52.0 = 0.2
50.1 = 0.2
ASR
0.2
15.4 1.4
14.6 14.6000 0.8
3.2
SQL Create Context
Utility
14.9
09.1.02. 0.2
99.1 1.1
98.5 1: 0.1
ASR
0.2
0.4
2.9
2.0
GSM8k
Utility
0.25
41.7+ 0.4
42.0 0.5
37.4
31


Reviews:
[{"review_id": "Vu2T2ozfzB", "reviewer": "Reviewer_Vp9q", "summary": "This paper proposes that the fragility of LLMs to various attacks (adversarial, prefilling, sampling, fine-tuning) could be explained by the models taking “shortcuts” during alignment whereby the conditional distributions of only the first few tokens are adjusted significantly. The authors empirically validate this claim and propose fine-tuning objectives that encourages “deeper” alignment beyond just the first few tokens, leading to an effective defense against the aforementioned attacks. Specifically, to protect against adversarial,  prefilling and sampling attacks, they propose to supplement fine-tuning with data augmentation of partial harmful responses followed by refusal text. To protect against fine-tuning attacks, they constrain the fine-tuning so that the conditional distributions of the first few tokens should be close to those of the base model’s. Overall, this work highlights a critical shortcoming of the alignment process of LLMs and offers simple and effective solutions to address it.", "questions": "1. Table 6 actually suggests some amount of length generalization is achieved, e.g. C=5 already generalizes well to 10 prefilled tokens. Can you provide some intuition on why this might be happening?\n2. In section 4.2, why is $\\beta_1$ set to be 4 times smaller than $\\beta_t$ for $2 \\leq t \\leq 5$?", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The overall exposition of the issue of depth in safety alignment presented in this paper is quite comprehensive. Many questions that I had while reading were sooner or later answered/investigated within the paper. I particularly also enjoyed how the authors were able to connect fine-tuning attacks into this paper, as initially they might seem very different from jailbreak attacks.\n2. The proposed fine-tuning objective to deepen safety alignment is simple, intuitive and demonstrably effective. The proposed token-wise constrained objective is also intuitive, and solid explanations are provided for how the $\\beta_t$ parameter affects the behavior of optimizing the objective.", "weaknesses": "1. The outputs are sampled using non-greedy decoding. Thus, the reported results in the paper may vary to some degree over multiple runs. The authors may want to consider averaging over the output sampling dimension as well and/or reporting greedy decoding results.\n2. Since safety evaluation is based on LLM judgement, it is subject to some amount of error. (This of course is true for any paper that uses an LLM-based safety judge, so I don’t believe this should affect the paper rating much.) The paper could therefore be further strengthened a bit with some human evaluation to estimate judgement accuracy using a small sample of the paper’s experiment results.\n3. Probably the most pressing issue is that there are no comparisons of the effectiveness of the proposed fine-tuning methods against any baselines. For example, it could make sense to compare against circuit breaking [1], a fine-tuning method that I would suspect be a strong competitor.\n\n[1] Zou, Andy, et al. \"Improving Alignment and Robustness with Short Circuiting.\" arXiv preprint arXiv:2406.04313 (2024).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kwQautlq0D", "reviewer": "Reviewer_Rq1M", "summary": "This paper demonstrates the shallow safety alignment issue through a variety of case studies. Essentially, the authors show that a variety of alignment attacks are successful because of a common issue within safety-aligned LLMS: only the first few output tokens are adapted during the model alignment process. Then the paper offers ways to mitigate this problem, which includes a data augmentation approach and a constrained optimization loss function.", "questions": "Why do you think the ASR still isn’t 0 (and in many cases is not close to 0) after using your mitigation strategies?\n\nIt seems like there could potentially be problems with the data augmentation approach since you are providing the model with these strange texts (e.g., you mention that the new texts are not coherent). Do you think that this matters? Is the model’s learning going to be compromised when it is learning with these incoherent texts?", "limitations": "", "rating": 10, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The paper is addressing an important problem, the vulnerability of safety alignment for LLMs, that can be very useful to real world problems.\n\nThe paper ties together prior works in a way that makes it easier to learn from them (i.e. highlighting the common thread amongst successful alignment attacks: their exploitation of shallow safety alignment).\n\nThe contributions of this paper lay the groundwork for future safety alignment solutions. They do offer a couple mitigation strategies, but exposing the shallow alignment issue could inspire many more mitigation approaches. It could also help us understand the success of other attacks and the success/failure of existing attack mitigation strategies.\n\nThe paper includes a good variety of experiments (models, datasets, attacks types) and includes both empirical and theoretical support for their claims.\n\nThe paper flows nicely. It is nicely organized. This makes the paper easy to follow and it makes the main point/contribution of the paper very clear.", "weaknesses": "The explanation of related work is lacking. The related works are listed, but there is not much information that actually explains how your work differs from related work. For instance, you say “some works have also noted asymmetries...” But it would be nice to know how this differs from what you’ve observed. A lot of the statements you make about related work are very broad and could benefit from more detail. “Our work ties these potential failure modes…to potential shortcuts” - does your work do this for all pre-existing methods for improving alignment? Are there some failures that your work does not encapsulate? Also, you never seem to mention any solutions to these alignment failures. Are your methods (e.g. the data augmentation and constrained optimization) the only known mitigation strategies? If so, you should state this. If not, other mitigation strategies should be mentioned.\n\nAfter applying your mitigation strategies, the ASR is still not zero and often isn’t even that close to zero. This isn’t ever really explained in the paper. You at one point say “the augmented model is still vulnerable…”, but the paper would be stronger if you give more explanation. For instance, does the non-zero ASR mean that there is some other vulnerability apart from the shallow alignment issue? Or are your strategies just not fully fixing the shallow alignment problem?\n\nYour contribution would be stronger if it were explained more clearly. When you say things like “this work is the first work attempting to consolidate the unified notion behind these attacks…” I don’t quite understand what you mean. If other works have identified the shallow safety alignment effect, then what does it mean for you to “consolidate the unified notion”? Is shallow safety alignment a new term that you are introducing, because if so, I think you should make it more clear that you are introducing this new concept?\n\nIt is also hard to imagine this problem in a real-world setting/application. The paper would be stronger if, for example in the introduction, we were given an example of the effect that jailbreaks can have (e.g. him what scenario would some attacker be able to provide a deployed model with the start to a response)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Lz66CPkRpD", "reviewer": "Reviewer_6zSQ", "summary": "The paper argues that RL-based alignment methods such as RLHF and DPO are superficial in token length, and presents evidence that instruction-tuned/aligned open-weight models rely on producing short refusal prefixes to induce aligned outputs. As a result, these models are vulnerable to prefilling attacks, where the an affirmative response is injected as a prefix to the assistant's generation. Experiments are conducted under 2 threat models: input-space attacks and fine-tuning attacks where the attacker only has access to the dataset used for fine-tuning. Two baselines are introduced: a data-augmentation approach for mitigating prefill-attack susceptibility, as well as an alternative objective that constrains the initial tokens' distribution shift during fine-tuning.", "questions": "1) Can authors provide any further commentary on the data augmentation experiments for prefill-attack robustness (e.g., did the distribution of refusal placement in the training examples matter)?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- They key claim of the paper is intuitive, which is that current alignment techniques can satisfy standard RLHF and DPO objectives by simply inducing common refusal prefixes; these prefixes are easily circumvented when one can fully control model inputs\n- Experimental results showing the KL divergence between base/aligned models at different token positions partially support this claim\n- The simple data augmentation approach shows promising robustness to prefilling attacks, and supports the authors' hypothesis\n- The authors leverage these insights to develop a modified DPO objective with per-token weights, which gets decent results", "weaknesses": "- The authors do not include any evaluations to highly relevant baseline defenses, like the following: \n    - Prior work [1] presents strong robustness to prefilling attacks by training model representations, but this work does not acknowledge this. \n   - Prompt-level defenses [2] have been proposed that optimize suffixes for defending against input-space attacks\n\n- Experiments that thoroughly measure reward hacking during alignment training would have made sense for this work. Since one of the claims is that refusal prefixes are an easy outlet for RLHF/DPO objectives, the authors' intuition on this should be substantiated beyond just the final KL divergence with base models.\n\n- Although the paper's aim is to generally illustrate that models are only superficially implementing common alignment objectives, a clear threat model could provide more clarity. The main concerns that the authors highlight seem primarily relevant to open-weight models, rather than the black-box API case, where model providers can control model inputs.\n\n\n\n\n\n\n[1] Zou, A., Phan, L., Wang, J., Duenas, D., Lin, M., Andriushchenko, M., ... & Hendrycks, D. (2024). Improving Alignment and Robustness with Short Circuiting. \n\n[2] Zhou, A., Li, B., & Wang, H. (2024). Robust prompt optimization for defending language models against jailbreaking attacks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7PxRoleLVC", "reviewer": "Reviewer_n9se", "summary": "This paper proposes two defenses aimining to improve the safety alignment of the LLMs. LLMs are vulunerable to many attacks, and the representative attack techniques are jail-break attack and harmful fine-tuning attack.\n\n* The first defense proposed by the authors aim to solve the jail-break attack. The idea is to augment the safety alignment data  in order to make the safety alignment a few tokens deeper. \n\n* The second defense aims to solve the fine-tuning attack. The idea is to exploit a adaptive KL like method to constrian the first few token to be similar to that of  the aligned model.", "questions": "1. Is the first few words of the harmful answer used in data augmentation the same with the target word used for GCG attack? For example, if GCG aims to elicit \"Sure,\", does the harmful answer you use for data augmentation start with \"Sure\"? Clould you give me a few samples to see how the harmful answers you used for augmentation looks like? \n\n\n2. Please also consider such an adaptive attack for your augmentation idea:\nThe attacker no longer aims to elicit the word \"sure\", but use GCG attack to elicit \"I cannot fulfill your request. Just kidding, I will tell you\".  Can the augmentation method still work against this adaptive attack? Let's give it a cool name \"GCG with a few tokens even deeper \".\n\n3. Can the constrain-SFT reduces to  LDIFS simply by tuning your hyper-parameter? LDIFS exploit KL regularizer uniformly accross all the tokens. It would be nice to see that constrain-SFT can reduce to LDIFS (Mukhoti et al, 2023) by tuning your hyper-parameter, and show that by simply tuning this hyper-parameter in order to focus on constrianing the first few tokens can give better results.  Moreover, I think the authors should difinitely disucss (Mukhoti et al, 2023)  because it shares a very similar insight with constrain-SFT, \n\n  \nMukhoti J, Gal Y, Torr P H S, et al. Fine-tuning can cripple your foundation model; preserving features may be the solution[J]. arXiv preprint arXiv:2308.13320, 2023. \n\nOverall, I think this paper should reach the acceptance bar of ICLR. I will actively participate in the discussion and will not disappear (or keep silence) from the author-reviewer discussion phease. I will consider to raise my score if my concerns are sufficiently addressed.", "limitations": "", "rating": 10, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. Writing is excellent!\n2. The constrained-SFT method interesting and might be a useful defense baseline for the community.  \n3. A few interesting phenomenon regarding the depth of alignment are discovered, and are displayed very clearly with well-structure results.", "weaknesses": "The idea of the proposed augmentation method is to modify the safety alignment data by postponing the refusal phrase and inserting some harmful answer before that. This method can significantly reduce the risk of  GCG attack, prefilling attack, etc.  However, I found several problems with this method:\n\n* **The true reasons why the data agumentation method work needs more discussion.** Take GCG attack as an example. GCG attack aims to optimize an suffix in the question, and this suffix can elicit a few confirmative phrase like :\"Sure,\". However, GCG attack is only targeting on a few tokens in the beginning of the answer, but no the deeper tokens. The reason that the augmentation method can mitigate GCG attack probably is that the model learns from the augmentation data to stop delivering harmful answer after a few words (i,.e., give refusal answer after a few word)  For example, I guess  after safety alignment on the augmentation data, the model output of a harmful question will be like this: \n> Instruct: How to make a bomb?\n> Answer: Sure, here is how to make a bomb. I cannot fullfill your request. It's not...\n\n* **The data augmentation method may not solve the jail-break attack from its root.** Will the defense still work if the GCG attack aims to elicit a **longer** harmful phrases? For example, what if the GCG attack aim to elicit the phrase like \"Sure, I will do whatever you want. Yes, no problem... Here is my answer: \"? If this phase is significantly longer than $k$, which is the number of tokens that are posponed in the augmentation data. I am wondering whether the augmenation method still can work.  The authors can provide more evidence to address my concern.\n\n*  **The reason that the data augmentation method cannot effectively against harmful fine-tuning attack is not specified.** The reason is probably that the harmful fine-tuning attack can still overthrow the refusal phrase even it is postponed. Different from GCG attack, harmful fine-tuning attack is not only targeting on the first few phrases to elicit harmful answers. \n\nOverall, I think the augmentation method work because it targets on some features that jail-break attack are exploiting -- they only try to elicit affiirmative answer in the first few tokens of the answer, and naturally the harmful answers will go on if the model does not learn from the augmentation data to interrupt them. On the other hand, if the model is trained from the augmentation data, it might learn to elicit refusal answers when it starts to output some harmful keywords, which could be a probable reason how the method works. The authors should give more discussion regarding this.\n\nThe second proposed method named constrain-SFT aims to solve the harmful fine-tuning attack. The idea is to contrain the distance of the output of the first few tokens between the aligned model and the fine-tuned model. This idea make seneses to me. However, I do want to mention a few aspects that can be improved.\n\n* **Lack of baselines.** Before this paper, there are already a few defense baselines to the harmful fine-tuning attack. The authors should include comparison with existing baselines, e.g., Vaccine (Huang et al, 2024).\n\nHuang T, Hu S, Liu L. Vaccine: Perturbation-aware alignment for large language model[J]. arXiv preprint arXiv:2402.01109, 2024. https://arxiv.org/abs/2402.01109 (First available Feb 2, 2024)\n\n* **System overhead analysis and experiments should be given.** Does the  solution comes with extra computation/memory overhead compared to DPO?  I conjecture the answer is yes because it needs another forward pass of the aligned model to derive its logit. I would like to hear from the authors regardig this. \n\n\n* **The paper can benefit from a more extensive literature review.** I list a few papers on the relevant topics of harmful fine-tuning defense as follows:\n\n\n\n---Alignment stage solution---\n\n[2024/2/2] Vaccine: Perturbation-aware alignment for large language model aginst harmful fine-tuning NeurIPS2024\n\n[2024/5/23] Representation noising effectively prevents harmful fine-tuning on LLMs NeurIPS2024\n\n[2024/5/24] Buckle Up: Robustifying LLMs at Every Customization Stage via Data Curation\n\n---Fine-tuning stage solution---\n\n[2023/8/25] Fine-tuning can cripple your foundation model; preserving features may be the solution\n\n[2023/9/14] Safety-Tuned LLaMAs: Lessons From Improving the Safety of Large Language Models that Follow Instructions ICLR2024\n\n[2024/2/3] Safety fine-tuning at (almost) no cost: A baseline for vision large language models ICML2024\n\n[2024/2/22] Mitigating fine-tuning jailbreak attack with backdoor enhanced alignment NeurIPS2024\n\n[2024/2/28] Keeping llms aligned after fine-tuning: The crucial role of prompt templates NeurIPS2024\n\n[2024/5/28] Lazy safety alignment for large language models against harmful fine-tuning NeurIPS2024\n\n---Post-fine-tuning stage solution---\n\n[2024/5/15] A safety realignment framework via subspace-oriented model fusion for large language models\n\n[2024/5/27] Safe lora: the silver lining of reducing safety risks when fine-tuning large language models NeurIPS2024\n\n\n[2024/5/25] No two devils alike: Unveiling distinct mechanisms of fine-tuning attacks\n\n[2024/5/27] Navigating the safety landscape: Measuring risks in finetuning large language models NeurIPS2024\n\n-------------Below is concurrent works (or after you)-----------\n\n[2024/6/12] Do as I do (Safely): Mitigating Task-Specific Fine-tuning Risks in Large Language Models\n\n\n[2024/8/1] Tamper-Resistant Safeguards for Open-Weight LLMs\n\n[2024/8/18] Antidote: Post-fine-tuning safety alignment for large language models against harmful fine-tuning\n\n[2024/8/27] Bi-Factorial Preference Optimization: Balancing Safety-Helpfulness in Language Models\n\n[2024/9/3] Booster: Tackling harmful fine-tuning for large language models via attenuating harmful perturbation\n\n[2024/9/26]Harmful fine-tuning attacks and defenses for large language models: A survey\n\n[2024/10/05] Identifying and Tuning Safety Neurons in Large Language Models\n\n[2024/10/13] Targeted Vaccine: Safety Alignment for Large Language Models against Harmful Fine-Tuning via Layer-wise Perturbation\n\n[2024/10/05] SEAL: Safety-enhanced Aligned LLM Fine-tuning via Bilevel Data Selection\n\n[2024/10/05] SaLoRA: Safety-Alignment Preserved Low-Rank Adaptation\n\n[2024/10/05] Safety Alignment Shouldn't Be Complicated\n\n[2024/10/05] Towards Secure Tuning: Mitigating Security Risks Arising from Benign Instruction Fine-Tuning\n\n[2024/10/05] Locking Down the Finetuned LLMs Safety\n\n[2024/10/05] Your Task May Vary: A Systematic Understanding of Alignment and Safety Degradation when Fine-tuning LLMs\n\n[2024/10/05] Unraveling and Mitigating Safety Alignment Degradation of Vision-Language Models\n\nI am aware that some of these work are concurrent study (or after you). The authors should at least discuss those very relevant papers appeared before the first appearance of this paper. It is also encouraged for the authors to discuss all the existing research as this would be beneficial for the field development.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)