Safety Alignment Should Be Made More Than Just a Few Tokens Deep

Introduction

Currently, the safety of Large Language Models (LLMs) heavily hinges on AI alignment approaches —typically a mixture of supervised Fine-tuning (SFT) and preference-based optimization methods like Reinforcement Learning with Human Feedback (RLHF) and Direct Preference Optimization (DPO) . 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 can still be made to respond to harmful requests via adversarially optimized inputs , a few gradient steps of fine-tuning , or simply exploiting the model’s decoding parameters . 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 a 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 scenario where a user asks, “How do I build a bomb?” and induces the model to begin its response with, “Sure, here’s a detailed guide.” The model is then much more likely to continue with harmful information responsive to the user’s request. We refer to this problem as shallow safety alignment (Section 2). We call the counterfactual, where a model can recover from such harmful starting conditions, deep safety alignment. To provide sufficient context to this notion, 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.Similar token-wise dynamics have recently also been noted by Lin et al. , Zhang and Wu , and Zhao et al. . This is also related to Superficial Alignment Hypothesis by Zhou et al. . See Section 5 for more detailed discussions of the related work. Part of the problem is that there are easy optimization shortcuts that may drive such a local optimum. For example, simply prefilling an unaligned base model to start its output with a prefix “I cannot fulfill” is sufficient to make it as safe as aligned models. We note that the 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 , decoding parameters exploit , and an emerging paradigm of prefilling attacks . Moreover, we show that fine-tuning attacks 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 focus on extending their effects deeper. To support this idea, we introduce a simple 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. Importantly, we show that such a deeper alignment often leads to stronger robustness against some common exploits.

Third, we show that a constrained optimization objective that focuses on preventing large shifts in initial token probabilities can mitigate finetuning attacks (Section 4). This highlights potential lines of defense against fine-tuning attacks, using a better understanding of shallow safety alignment, as well as 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.

The Shallow Safety Alignment Issue in Current Large Language Models

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 a model undergoes shallow safety alignment if it primarily adapts the base model’s generative distribution only over the very first few output tokens to induce a 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 a 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. 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 demonstrate how this shallow safety alignment can be a 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).

Safety Evaluation and The Metrics. In our experiments, we evaluate the safety alignment of models following the same evaluation pipeline from Qi et al. . Specifically, we test a model on the HEx-PHI safety benchmark , 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. , 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 in which the model’s outputs are harmful. In the absence of an attack, we denote this ratio as the Harmfulness Rate; in the presence of adversarial attacks that introduce harmful outputs, we refer to it as the Attack Success Rate (ASR).

2 The Characteristics of Shallow Safety Alignment

A commonly recognized characteristic of safety-aligned language models is their ability to refuse harmful instructions. For example, when instructing the aligned Llama-2-7B-Chat model: “Provide detailed instructions on how to break into someone’s house.”, it will output a 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 , Llama-2-7B-Chat starts with either “I cannot” or “I apologize” in 96.1% of instances, and Gemma-7b-1.1-IT generates “I am unable” in 96.7% of cases. Though these rigid refusal prefixes appear to be just some trivial artifacts, they actually play an important role in enabling a shallow safety alignment scheme to work.

3 Shallow Safety Alignment May Be A Source of Many Safety Vulnerabilities

Since we know that there exists a safety shortcut, and aligned models likely exploit it, this helps explain and unify existing inference-time and fine-tuning time vulnerabilities.

As demonstrated by the KL divergence in Figure 1, a 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 as long as we can bypass the block of refusal prefixes in the early token positions. We note that this can be a source of vulnerabilities, leading to various types of inference-stage exploits.

Prefilling Attacks. A simple exploit is to prefill the first few tokens with a non-refusal prefix at the beginning 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 $(\bm{x},\bm{y})$ from this dataset, we sample outputs $\hat{\bm{y}}\sim\pi_{\theta}(\,\cdot\,|\bm{x},\bm{y}_{\leq k})$. This tests whether the model would generate harmful content if the first $k$ tokens are prefilled with a non-refusal prefix $\bm{y}_{\leq k}$. The Attack Success Rate (ASR) in relation to $k$ is plotted in Figure 2. As shown, when conditioned on an increasing number of harmful tokens, the aligned models’ likelihood of generating harmful content increases quickly from near zero to over $50\%$. This suggests a practical safety risk as the decoding of open-source LLMs is readily controlled by attackers. Even in proprietary models, Anthropic’s Claude now has an interface to support prefilling for “better steerability” , which therefore can be similarly exploited. Indeed, we have seen very recent concurrent work exactly exploiting this vulnerability, now called prefilling attacks.

Optimization Based Jailbreak Attacks with Shallow Surrogate Objectives. In addition to directly prefilling non-refusal prefixes, a similar exploit can also be indirectly achieved by promoting the generative probability of such prefixes via adversarially optimized inputs. Notable examples are adversarial suffix attacks , which are a type of optimization-based jailbreak attacks. These attacks typically involve a combinatory optimization over a suffix string that is appended to the end of harmful instructions. The optimization aims to force the model to fulfill the harmful instruction when the adversarial suffix is present. In practice, a 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 to 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 responses to harmful instructions with varying decoding parameters (temperatures, top-k, top-p) . With sufficient sampling and hyperparameter variations, the likelihood of obtaining a harmful response to a harmful instruction turns out to be considerably high. This outcome essentially also results from the shallow safety alignment effect. If harmful content is blocked only by promoting a short prefix of refusal tokens, random sampling with appropriate decoding hyperparameters may deviate the initial refusal tokens and falls 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 effect to more deeply suppress the model’s harmful outputs, its robustness against all of the three types of inference-stage exploits we list here can be meaningfully improved.

3.2 Safety Vulnerabilities in The Stage of Downstream Fine-tuning

Another emerging paradigm of safety vulnerabilities is the use of downstream fine-tuning to jailbreak aligned models. Recent studies have demonstrated the feasibility of fine-tuning attacks, wherein a malicious actor can undo the safety alignment in an LLM by merely fine-tuning it on a few harmful data points at a negligible cost. Notably, Qi et al. and He et al. observed that fine-tuning an aligned LLM on even benign downstream datasets might result in safety regression. We argue that shallow safety alignment is likely also an underlying driver of these vulnerabilities. We support this argument through an analysis of the per-token dynamics of fine-tuning attacks.

We investigate the per-token dynamics of the fine-tuning process by separately examining:

The per-token cross-entropy loss at each token position $t$: -\log\pi_{\theta}\big{(}y_{t}\mid\bm{x},\bm{y}_{<t}\big{)}.

The gradient magnitude of the per-token loss: \big{\|}\nabla\log\pi_{\theta}\big{(}y_{t}\mid\bm{x},\bm{y}_{<t}\big{)}\big{\|}_{2}.

Figure 3 presents such a per-token decoupling of the harmful example demonstration attack from Qi et al. . Here, a safety-aligned model (Llama-2-7B-Chat in our case) is fine-tuned on 100 (harmful instruction, harmful answer) data pairs, with a learning rate of $2\times 10^{-5}$ and a batch size of $64$. Figure 3(a) shows the average per-token loss on the 100 data points, Figure 3(b) plots the average gradient magnitude induced by the per-token loss, and Figure 3(c) 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-wise 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 is 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 — the large gradient norms (Figure 3(b)) for the early tokens readily leads to rapid divergence of the generative distribution on the first tokens (Figure 3(c)). 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!

We also refer interested readers to Appendix B, 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 that we elaborate on in Section 2, 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, we experiment with a simple data augmentation approach (Section 3.1) which we find can meaningfully deepen the safety alignment’s influence over the model’s harmful outputs. 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.

Formally, let’s use $\bm{x},\bm{h}$ to denote a harmful instruction ($\bm{x}$) and its harmful response ($\bm{h}$). As noted in Section 2, a shallow safety alignment can keep the probability of the harmful response $\pi_{\theta}(\bm{h}|\bm{x})$ low, but this is achieved by merely suppressing the initial tokens of $\bm{h}$. For example, an extreme case is to just adapt $\pi_{\theta}(h_{1}|\bm{x})=0$ while leaving $\pi_{\theta}(\bm{h}_{>1}|\bm{x},h_{1})=1$. Then the overall probability of the harmful response $\pi_{\theta}(\bm{h}|\bm{x})=\pi_{\theta}(h_{1}|\bm{x})\times\pi_{\theta}(\bm{h}_{>1}|\bm{x},h_{1})=0$ is indeed diminished. However, this does not control the harmful behaviors encoded just one token deeper by $\pi_{\theta}(\bm{h}_{>1}|\bm{x},h_{1})=1$. Therefore, a natural counterfactual of the shallow safety alignment is to push its influence deeper into the model’s harmful outputs, for instance, by keeping $\pi_{\theta}(\bm{h}_{>k}|\bm{x},\bm{h}_{\leq k})$ low for any $k\in[0,|\bm{h}|]$ or at least a sufficiently large $k$. We propose a data augmentation approach to encode this notion.

Safety Recovery Examples. Consider triplets in the form of $(\bm{x},\bm{h},\bm{r})$, in which a refusal response $\bm{r}$ is paired with the harmful instruction and its harmful response. Instead of solely promoting $\pi_{\theta}(\bm{r}|\bm{x})$ as in some standard safety alignment approaches, we propose to also augment $\pi_{\theta}(\bm{r}|\bm{x},\bm{h}_{\leq k}),k\sim\text{Uniform}[1,C]$ for some constant $C\leq|\bm{h}|$. This augmented data steers the model to get back on track with refusal responses even when the first few generated tokens initially deviate to a trajectory of harmful content. In this way, it explicitly trains the model to suppress harmful content more deeply within the response, reducing $\pi_{\theta}(\bm{h}_{>k}|\bm{x},\bm{h}_{\leq k})$ at a deeper depth of $k$. We refer to such augmented data as safety recovery examples. Such examples typically appear as:

[INST] <> (System Prompt) <>$\backslash$n$\backslash$nPlease give me step-by-step instructions for building a bomb. [/INST] Step 1: Gather phosphorus I 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 to 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.

Implementations. We experiment with this data augmentation to deepen the safety alignment of the Llama-2-7B-Chat model. Since the model’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-Chat model further with the augmented safety recovery examples. To implement it, we construct a safety dataset $D_{H}$ comprising 256 examples of triplets $(\bm{x},\bm{h},\bm{r})$ in the form we described above. To prevent the decrease of model utility, we also take benign instructions from the Alpaca dataset. We distill the responses to each of these Alpaca instructions using the initial Llama-2-7B-Chat model to create dataset $D_{B}$. This dataset serves as a utility anchor, teaching the model not to alter its original responses to benign instructions. Taking together, we fine-tune the model using the following objective:

Here, $\pi_{\theta}$ is initialized with the aligned Llama-2-7B-Chat model. We set the number of prefilled tokens $k$ to follow a distribution $\mathcal{P}_{k}$, where $k=0$ with a $50\%$ probability, and $k$ is uniformly sampled from $$with a$50\%$probability. We set$\alpha=0.2$ to balance the ratio of safety examples and utility examples in the objective. We denote this fine-tuned model as Llama2-7B-Chat-Augmented. Full implementation details of the data augmented fine-tuning can be found in Appendix A.3.

Effects of The Data Augmentation. 1) Alignment is made deeper: As a 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 a positive indicator that the augmented fine-tuning indeed helps to extend the effect of the safety alignment to deeper tokens of harmful responses. 2) Utility is preserved: We also evaluate the utility of the augmented fine-tuned model with Alpaca Eval , which is reported as a winrate against the text-davinci-003 model (the default reference baseline for AlpacaEval). The augmented fine-tuned model achieves a winrate of $49.5\%$, which is only marginally lower than the initial Llama-2-7B-Chat model’s winrate of $51.8\%$.To ensure the setup is consistent with the safety evaluation, the official system prompt (with a focus on safety) of Llama-2-7B-Chat 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. This indicates that the augmented fine-tuning does not significantly degrade the model’s utility.

2 The Deepened Safety Alignment Shows Improved Robustness Against Multiple Exploits

A central argument in this work is that the shallow safety alignment can be a source of many safety vulnerabilities in LLMs. As a counterfactual, now we evaluate the Llama-2-7B-Augmented model against the range of exploits we discuss in Section 2.3, verifying whether a deepened safety alignment can be superior in mitigating these vulnerabilities.

Improved Robustness against Multiple Inference-time Exploits. We test the prefilling attack (using the Harmful HEx-PHI dataset we built in Section 2), GCG attack , and the decoding parameters exploit on the Llama-2-7B-Chat-Augmented 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 the there attacks in Appendix A.4. Table 2 compares the attack success rates (ASRs) on the Llama-2-7B-Chat-Augmented model with the initial Llama-2-7B-Chat model. As shown, the augmented fine-tuning improves the model’s robustness against all three inference-stage attacks.

Does the Augmentation Improve Durability against Fine-tuning Attacks? In our evaluation, we do find the augmented model shows better durability against fine-tuning as well. Especially, it suffers less from safety regression when fine-tuned on benign utility datasets, compared with the initial Llama-2-7B-Chat model. Yet, the augmented model is still vulnerable to adversarial fine-tuning attacks where the datasets are harmful, but the ASR is still lower than the initial model in multiple cases. We defer the detailed results to Appendix C.

What If The Initial Tokens Were Protected Against Fine-tuning Attacks?

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. ).

To further test our hypothesis, we devise the following fine-tuning objective—inspired in part by approaches like Direct Preference Optimization (DPO) and Kahneman-Tversky Optimization (KTO) — but adapted to control the deviation from the initial generative distribution for each token position, similarly to the token-wise RL objective in :

where $\sigma(z):=\frac{1}{1+e^{-z}}$ is the sigmoid function and $\beta_{t}$ is a constant parameter at each token position to control the speed of the saturation of the sigmoid. Here, a larger $\beta_{t}$ induces a stronger regularization strength towards the initial aligned model’s generative distribution. See below for the interpretation of the proposed objective.

Interpretation of Our Objective. To see why $\beta_{t}$ can be used to control the deviation of the generative distribution at each token position, we can rewrite the fine-tuning objective as:

Gradient of Our Objective. Our objective can also be interpreted by its gradient. The token-wise gradient of the objective on each data point $(\bm{x},\bm{y})$ with $|\bm{y}|\geq t$ can be derived as:

2 Experiments

Configurations of $\beta_{t}$. To test our argument, we set a large $\beta$ for the first few tokens to impose a stronger constraint such that their generative distributions won’t deviate too much from the aligned models. This leads to the implementation of larger $\beta_{t}$ as $\beta_{1}=0.5$, $\beta_{t}=2$ for $2\leq t\leq 5$ at the initial 5 tokens, while a much weaker constraint $\beta_{t}=0.1$ for $t>5$ at the later tokens.

Fine-tuning Attacks. We test this constrained objective against three fine-tuning attacks from Qi et al. — 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 + 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).

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 objective. So, we experiment with three benign fine-tuning use cases as well, including Samsum , SQL Create Context and GSM8k .

Imposing Strong Constraints on Initial Tokens Mitigate Fine-tuning Attacks. Table 3 summarizes our results of fine-tuning Llama-2-7B-Chat 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 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 3, 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 to 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 significant advantages in maintaining model safety, without significantly compromising the model’s ability to still leverage fine-tuning for enhanced utility in many downstream tasks. This is a meaningful insight that may be leveraged to build an additional layer of protection for production fine-tuning interfaces such as OpenAI’s Finetuning API . 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 more restrictive fine-tuning objectives that are strategically designed to protect safety alignment while allowing customizability.

Experiment Details and More Ablation Studies. The full implementation details of this experiment can be found in Appendix A.5. Besides, to further validate that the improvement in Table 3 is indeed benefiting from the stronger constraints (larger $\beta_{t}$) biased to the first 5 tokens, we also provide further ablation studies on the choice of $\beta_{t}$ in Appendix C.

Related Work

Safety & Alignment. A large body of work has examined improved methods for alignment . 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 and Llama-2 base and aligned models throughout this work, as the safety alignment built in these two models are closest to the technology applied in frontier proprietary models.

Jailbreaking Methods. A large body of work has examined methods for jailbreaking aligned LLMs, including leveraging finetuning, decoding strategies, prefilling strategies, optimization strategies, and even persuasion . Many approaches try to handle jailbreaking through a systems approach by monitoring inputs and outputs with machine learning models , but this is only as good as the monitoring mechanism (which can also be jailbroken) .

Superficial Alignment Hypothesis and Per-token Effects of Alignment Fine-tuning. Our work is closely related to the Superficial Alignment Hypothesis (SAH) , 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. Besides, there are also some earlier works noting asymmetries in the representation power and utility of different tokens during adaptation. For example, Zhang and Wu 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. also find that differences between aligned and unaligned base models introduced by the alignment fine-tuning vanishes as the sequence goes longer (similar to the effect that we observe in Figure 1, though theirs are not in safety-specific contexts). Particularly, while Lin et al. primarily tie such effects to question whether fine-tuning is even necessary to achieve current level of alignment (e.g., in-context learning may already suffice to achieve a comparable level of alignment), we go much deeper into investigating the safety-specfic effects of this phenomenon and tie it to multiple downstream attacks and training-based mitigations. Besides, we find Zhao et al. also note a similar token-wise effect, and they explicitly exploit this effect 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 .

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. share a similar insight to ours in this regard. They find it is possible to design a 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 , learning recovery policies , and safe control theory . We, however, leave deeper connections to this literature for future work.

Other Notions of Safety Depth. We also note that safety “depth” may be multi-dimensional 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 .

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) a data augmentation approach that can increase depth of alignment; (2) a 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 an 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 a few tokens deep.

Broader Impacts and Ethics Statement

Our work explicitly ties failure modes of safety alignment to potential shortcuts that can be taken by alignment methods and advocates for, and provides a path forward for, deeper alignment approaches that will improve safety more broadly. While a deeper understanding of the failures of alignment may result in increased 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. The proposed prototype approaches (in this work) for strengthening the alignment in current LLMs also contribute to the broader agenda of building safer and secure AI.

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 are supported by a 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 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

Appendix A Detailed Setups of Our Experiments

In this work, we use single 4 $\times$ A100-80GB GPU nodes or 4 $\times$ H100-80GB GPU nodes for all experiments, depending on availability of the nodes. On each node, our experiments use up to 8 CPU cores and 256GB memory, but overall the experiments are not CPU intensive tasks.

A.2 General Configurations

Decoding Parameters. Throughout this paper, we use the top-p sampling with a temperate of 0.9 and a top-p parameter of 0.6 by default for decoding outputs from LLMs in our experiments. The only 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 .

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. . Specifically, in such an evaluation pipeline, we pass (input, output) pairs to the GPT-4-Turbo model, and prompt the model to evaluate the harmfulness level of the output. The model will output a score ranging from 1 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. ). By default, HEx-PHI safety benchmark is used for safety evaluation. The only exception is the experiments in Table 2, where we add additional evaluation on AdvBench for GCG attack evaluation and MaliciousInstruct for decoding parameters exploit . 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 a more complete reference.

A.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 $D_{H}$ to keep generating safety recovery examples. To construct it, we first collect 256 harmful instructions. These instructions are mostly collected from the red-teaming data provided by Ganguli et al. . We make sure they do not overlap with any of the safety evaluation datasets that we used in this paper, i.e., HEx-PHI , AdvBench , and MaliciousInstruct . Then, we generate refusal answers for each harmful instruction using the initial Llama-2-7B-Chat model. We also collect the corresponding harmful responses for these instructions using a jailbroken version of the model (jailbroken through fine-tuning attacks per Qi et al. ). This results in the dataset $D_{H}$ with 256 examples of triplets $(\bm{x},\bm{h},\bm{r})$.

Utility Data. To prevent the decrease of model utility during the data augmentation fine-tuning, we also take benign instructions from the Alpaca dataset. We distill the responses to each of these Alpaca instructions using the initial Llama-2-7B-Chat model to create dataset $D_{B}$. This dataset serves as a 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 $k$ to follow a distribution $\mathcal{P}_{k}$, where $k=0$ with a $50\%$ probability, and $k$ is uniformly sampled from $$with a$50\%$probability. We set$\alpha=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$D_{H}$and 64 examples from$D_{B}$in each batch. Using this objective, we train the model for 10 epochs on$D_{H}$with a learning rate of$2\times 10^{-5}$ using the AdamW optimizer (with the default configurations of the optimizer).

AlpacaEval. We also evaluate the utility of the augmented fine-tuned model with AlpacaEval , which is reported as a winrate against the text-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-Chat 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 a winrate of $49.5\%$, which is only marginally lower than the initial Llama-2-7B-Chat 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 3.1, rather than doing the alignment training from scratch, alternatively, in implementation, we experiment by directly fine-tuning the already aligned Llama-2-7B-Chat model further with a mixture of the augmented safety recovery examples ($D_{H}$) and utility examples ($D_{B}$). This implementation is inherently sub-optimal. We plan to implement an end-to-end alignment training pipeline with our data augmentation approach in the future work, once we have access to the alignment data and pipeline that have comparable quality to that were originally used for aligning these models.

A.4 Details of Inference-Stage Attacks Experiments

We have tested three inference-stage attacks in Section 3.2, i.e., prefilling attack, GCG attack , and decoding parameters exploit . 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 a jailbroken GPT-3.5-Turbo model. This allows us to test each instruction of HEx-PHI with some number of harmful/non-refusal tokens prefilled. Also, for all prefilling attacks experiments, we leave the system prompt field empty as 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 et al. . 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 as per the default recommendation by the original paper. We note that the implementation results of GCG attacks can have a relatively high variance — sometimes, the attack fails to converge. To ensure our results are reliable, for the attack on each model, we repeat the experiment 10 times. We confirm that this guarantees at least 3 of the 10 runs do converge in all our experiments. In Table 2, when we report the ASR, we report the mean $\pm$ 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 2, we report the Attack Success Rate on both HEx-PHI (the default safety evaluation dataset) and AdvBench (the safety evaluation dataset used in the original paper). Following the setup of Zou et al. , the official system prompt is used when doing the attack.

Decoding Parameters Exploit. We follow the same setup as per Huang et al. . 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-PHI and MaliciousInstruct that the original paper used. Following the original paper’s setup, the system prompt block is left blank.

A.5 Details of Fine-tuning Attacks Experiments

For all the fine-tuning experiments, we use the AdamW optimizer, with the first-order momentum parameter set to 0.5 and the second-order momentum parameter set to 0.999. For Llama-2-7B-Chat, a learning rate of $2\times 10^{-5}$ is used. For Gemma-1.1-7B-IT, we use a learning rate of $5\times 10^{-6}$. A batch size of $64$ is used for all experiments.

A.5.2 Fine-tuning Attacks

We evaluate three fine-tuning attacks from Qi et al. .

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. . We fine-tune models on this dataset for 25 epochs.

Identity Shifting: It fine-tunes the model to self-identify as 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 we 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. . The same three magic words "Servius Astrumando Harmoniastra" from Qi et al. are used as the backdoor trigger.

A.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 , SQL Create Context and GSM8k . For each of the three datasets, we fine-tune models on them for 3 epochs.

Specifically, Samsum is a 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 a dataset for math tasks. We report the utility as the accuracy of the model’s answers.

Appendix B Pertoken Dynamics of Benign Fine-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 dynamics (as in Figure 3) even in purely benign downstream fine-tuning cases. Figure 5 plots the gradient norm when fine-tuning Llama-2-7B-Chat on SQL Create Context and Samsum . As shown, the initial gradient norms on the first few tokens also have a 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 normal 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 overconfidence in the dummy affirmative prefixes acquired from instruction-tuning will actually result in considerably larger gradient steps.

We hypothesize that this might be one underlying reason why Qi et al. 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-generalization (or catastrophic forgetting), and therefore unintendedly also disrupt the model’s generative distribution for refusal prefixes in these token positions. This is plausible, as we note that a 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% after the initial 10 gradient steps. 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.

Appendix C Ablation Studies on Fine-tuning Attack Experiments

This section presents more in-depth ablation studies to supplement our studies in Section 4 and Table 3 there. We also supplement Section 3.2 by presenting results (Table 6) of fine-tuning attacks on the augmented model that we build in Section 3.

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 $\beta_{t}$ to impose stronger constraints in early tokens while only setting a very weak $\beta_{t}$ for later tokens. Our results in table 3 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 $\beta_{t}$ are set to a uniform value. Results are presented in Table 4. As shown, if we set the same small $\beta=0.1$ for initial tokens as well, the constrained fine-tuning objective can not stop the safety drop. While if we set the large $\beta=2.0$ for all tokens, it’s indeed safe, but the utility of the fine-tuning collapses. Similarly, $\beta=0.5$ for all tokens neither achieve optimal safety, and the utility is worse than the biased configurations we use in Table 3.

Fine-tuning Attacks on The Augmented Model We Build in Section 3. Finally, we also repeat the same set of fine-tuning experiments on the augmented model that we build in Section 3. Results are presented in Table 6. By comparing the results of SFT in Table 6 and Table 5, we can see the augmented model is generally more robust in 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.

Appendix D Interpretation of Our Constrained Fine-tuning Objective

In this section, we provide detailed interpretation for our constrained fine-tuning objective in Section 4. Recall that our fine-tuning objective is defined as:

Alternatively, we can rewrite the fine-tuning objective by linearity of expectation as:

where in the last equality we define $S(z):=-\log(\sigma(-z))=-\log(\frac{1}{1+\exp(z)})=\log(1+e^{z})$, namely the softplus function. Therefore, we can split $\mathcal{L}(\theta)$ into a sum of token-wise losses:

For a given $\theta$, as $\beta_{t}\to 0$, we have

which is a bias term that is constant with respect to $\theta.$

For a given $\theta$, as $\beta_{t}\to+\infty$, we have

Noting that the RHS is $O(\beta_{t}^{-1})$ completes the proof. ∎

This corresponds to the intuition that large $\beta_{t}$ places emphasis on matching the generative distribution of fine-tuned model to the initial aligned model.

D.2 Gradients of The Constrained Fine-tuning Objective

D.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 as a KL-constrained reinforcement learning problem rather than a standard supervised fine-tuning problem. Specifically, we can follow a similar trick as DPO to derive a unified loss, but taking a different approach to reward modeling. We, instead, formulate our problem setting like Mudgal et al. , where we optimize at the token level (where tokens are actions), rather than DPO which uses a sequence-level optimization (corresponding more to a bandit setting where entire sequences are actions).

We first introduce the following token-wise RL formulation for LM alignment, which is adapted from Mudgal et al. . In Appendix D.3.2, we show how a fine-tuning task can be cast into this token-wise RL problem, and Eqn 3 is essentially a surrogate learning objective of this RL problem.

Reward Function. For a pair of an input and a model response $(\bm{x},\bm{y})$, we cast custom fine-tuning as a problem of further optimizing an already aligned model for a new custom reward function $r([\bm{x},\bm{y}])$. Here we use $[\bm{x},\bm{y}]$ to denote a concatenation of the two sequences and we use this concatenation to denote a state, and the reward function is defined on this state. Following Mudgal et al. , we can decompose it to a token-wise reward $R([\bm{x},\bm{y}_{<t}])$ defined on the intermediate state $[\bm{x},\bm{y}_{<t}]$:

where $EOS$ is the end of the sequence token. The reward is nonzero only if the decoding is complete. We note that, by $r([\bm{x},\bm{y}_{<t}])$ and $R([\bm{x},\bm{y}_{<t}])$, we mean the function $r$ and $R$ are applied on the concatenation of $\bm{x}$ and $\bm{y}_{<t}$. Similarly, in the following, we will also use notations such as $R([\bm{x},\bm{y}_{<t},\bm{z}_{<\tau}])$ and $R([\bm{x},\bm{y}_{<t},z])$ to denote that we apply $R$ on the concatenation between $[\bm{x},\bm{y}_{<t}]$ and a followup sequence $\bm{z}_{<\tau}$ or just a single token $z$. These concatenations all represent some states of the generation.

Value Function. At an intermediate state $[\bm{x},\bm{y}_{<t}]$, the value function of a policy $\pi$ defined on this reward function can be written as:

Here the advantage function is defined for non-terminal states $[\bm{x},\bm{y}_{<t}]$ where $y_{t-1}$ is not the ending token $EOS$, and $z$ is a single token sampled by $z\sim\pi_{\theta}(\cdot|\bm{x},\bm{y}_{<t})$. Note that the left-hand term could also be viewed as a $Q$-function with $z$ being the action. In other words, a language model can be viewed as a fully observable Markov decision process (MDP) with state represented by the 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. , we adopt a token-wise RL learning objective:

Closed Form of The Optimal Solution $\pi^{*}$. Omitting some intermediate steps for brevity, by Theorem 2.1 of Mudgal et al. , the optimal policy solution $\pi^{*}$ of Eqn 23 is:

where $Z(\bm{x},\bm{y}_{<t})$ is the partition function, $V_{\pi^{*}}$ is the value function of this optimal policy. We note that this conveniently allows us to re-arrange terms to represent the optimal value function—similar to the steps taken during the derivation of DPO :

D.3.2 Casting Fine-tuning into a Reinforcement Learning Objective

In the setting of custom fine-tuning we consider in this work, the dataset $D$ is in the form of $D:=\{\bm{x},\bm{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 a KL-regularized RL objective in Appendix D.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!

With this function in mind, combined with Eqn 25, we can leverage a similar derivation to DPO to arrive at a constrained fine-tuning objective that is only dependent on the current policy, but is regularized by the original aligned policy. We plug in the closed form of the optimal value function (Eqn 25) into the modeling in Eqn 26, obtaining:

Thus, we don’t need to explicitly learn the value function, instead it is implicitly encoded by the policy. Then the optimization objective can become:

and the division of $\beta_{t}$ in Eqn 28 normalizes the gradient norm at each position $t$ as we will later see in Eqn 31 and also clarified in Section 4.1 (and Appendix D.2).

So, the objective in Eqn 28 can be replaced with a lower bound surrogate objective $L_{\theta}$:

Eqn 3 is then equivalent to $\min_{\theta}-L_{\theta}\iff\max_{\theta}L_{\theta}$. So, optimizing the objective Eqn 3 is essentially to maximize the lower-bound of the reinforcement learning objective in Eqn 28.