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

## Abstract

Abstract The safety alignment of current Large Language Models (LLMs) is vulnerable. Relatively simple attacks, or even benign fine-tuning, can jailbreak aligned models. We argue 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 refer to this issue as shallow safety alignment. In this paper, we present case studies to explain why shallow safety alignment can exist and provide evidence that current aligned LLMs are subject to this issue. We also show how these findings help explain multiple recently discovered vulnerabilities in LLMs, including the susceptibility to adversarial suffix attacks, prefilling attacks, decoding parameter attacks, and fine-tuning attacks. Importantly, we discuss how this consolidated notion of shallow safety alignment sheds light on promising research directions for mitigating these vulnerabilities. For instance, we show that deepening the safety alignment beyond just the first few tokens can often 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. 1 1 1 Our code is available at https://github.com/Unispac/shallow-vs-deep-alignment

## 1 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.(^2^22Similar 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.

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

### 2.1 Preliminaries

Notation. We use $\pi_{\theta}$ to denote a language model parameterized by weights $\theta$. We sometimes also directly use $\pi_{\mathrm{base}}$ to denote an (unaligned) pre-trained model (e.g., Llama-2-7B, Gemma-7B) to contrast its aligned counterpart $\pi_{\mathrm{aligned}}$ (e.g., Llama-2-7B-Chat, Gemma-7B-IT) [, ]. Given an input $\bm{x}$, the model’s output is modeled by $\pi_{\theta}(\,\cdot\,|\bm{x})$. We use $\bm{y}\sim\pi_{\theta}(\,\cdot\,|\bm{x})$ to denote the sampling of output $\bm{y}$. For token sequences like $\bm{x}$, $\bm{y}$, we use $x_{t}$, $y_{t}$ to denote their $t$-th tokens, and $|\bm{x}|,|\bm{y}|$ to denote their lengths (i.e., number of tokens). We also use $\bm{y}_{<t}$ and $\bm{y}_{\leq t}$ to denote the subsequences ranging from the first to the $(t-1)$-th tokens and from the first to the $t$-th tokens in $\bm{y}$, respectively. Similarly, $\bm{y}_{>t}$ and $\bm{y}_{\geq t}$ are employed to denote subsequences after the $t$-th and $(t-1)$-th tokens.

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

**Table 1: A Shorcut to The Safety Mode: The harmfulness rate of even unaligned models will diminish when a refusal prefix $\bm{s}$ is prefilled during decoding, i.e., $\bm{y}\sim\pi_{\theta}(\cdot|\bm{x},\bm{s})$.**
| Refusal Prefixes $(\bm{r})~{}\rightarrow$ | No<br>Prefix | “I cannot” | “I cannot<br>fulfill” | “I apologize” | “I apologize,<br>but I cannot” | “I am<br>unable” |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| $\downarrow$ Harmfulness Rate (%) on HEx-PHI Benchmark with A Refusal Prefix Prefilled During Decoding |  |  |  |  |  |  |  |
| Llama-2-7B | Aligned | 0 | 0 $\pm$ 0 | 0 $\pm$ 0 | 0 $\pm$ 0 | 0 $\pm$ 0 | 0 $\pm$ 0 |
| Base | 68.6 $\pm$ 0.8 | 16.4 $\pm$ 1.4 | 5.4 $\pm$ 1.3 | 14.4 $\pm$ 0.6 | 2.1 $\pm$ 0.2 | 8.1 $\pm$ 0.4 |  |
| Gemma-7B | Aligned | 2.1 $\pm$ 0.2 | 0 $\pm$ 0 | 0 $\pm$ 0 | 0 $\pm$ 0 | 0 $\pm$ 0 | 0 $\pm$ 0 |
| Base | 85.4 $\pm$ 0.6 | 8.7 $\pm$ 1.2 | 2.7 $\pm$ 0.5 | 14.1 $\pm$ 0.4 | 1.0 $\pm$ 0.8 | 3.9 $\pm$ 0.4 |  |

The “Safety Mode” Shortcut: 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 $\pi_{\mathrm{base}}$, if we can make its generated outputs begin with these refusal prefixes, the following output is likely to be safe. Using harmful instructions $\bm{x}$ from the HEx-PHI safety benchmark, we validate this by prefilling a refusal prefix $\bm{s}$ at the beginning of the decoding process to generate outputs $\bm{y}\sim\pi_{\mathrm{base}}(\,\cdot\,|\bm{x},\bm{s})$. Table 1 shows the Harmfulness Rate of the outputs produced by the models with different prefilled refusal prefixes $\bm{s}$ for both Llama-2 [] and Gemma [] 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 forced to start with refusal prefixes.
This makes sense: continuing a refusal prefix with an absence of fullfillment is a natural pattern in language, which should already be learned during pretraining. But it suggests
a simple shortcut or reward hacking scheme for safety 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.

Figure: Figure 1: Per-token KL Divergence between Aligned and Unaligned Models on Harmful HEx-PHI.
Refer to caption: extracted/5652106/figs/prefilling/harmful_hexphi_kl.png

Current Safety-aligned Models Exploit This Shortcut.
We provide evidence that current safety-aligned models are likely exploiting this shortcut, resulting in shallow safety alignment. We first construct a 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 instructions using a jailbroken version of GPT-3.5-Turbo from Qi et al. []. We call this dataset Harmful HEx-PHI. With this harmful dataset, we can examine the per-token KL divergence $D_{\text{KL}}\big{(}\pi_{\mathrm{aligned}}(\,\cdot\,|\bm{x},\bm{y}_{<{k}})\big
{\|}\pi_{\mathrm{base}}(\,\cdot\,|\bm{x},\bm{y}_{<{k}})\big{)}$ between the aligned model $\pi_{\mathrm{aligned}}$ and unaligned pre-trained model $\pi_{\mathrm{base}}$ on each of the harmful example $(\bm{x},\bm{y})$. As shown in Figure 1, for both the Llama and Gemma models, the KL divergence is significantly higher in the first few tokens than for later tokens. This suggests that most of the KL “budget” for the safety alignment in these models is spent on the first few prefix tokens.(^3^33Using the terminology “spending” KL on a KL “budget” from Gao et al. [].) At 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 to write any kind of examples that refuse a request after providing a harmful prefix; during RLHF, the model’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.

### 2.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.

#### 2.3.1 Inference-Stage 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.

Figure: Figure 2: ASR vs. Number of Prefilled Harmful Tokens, with $\hat{\bm{y}}\sim\pi_{\theta}(\cdot|\bm{x},\bm{y}_{\leq k})$ on Harmful HEx-PHI.
Refer to caption: extracted/5652106/figs/prefilling/harmful_prefilling.png

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.

#### 2.3.2 Safety Vulnerabilities in The Stage of Downstream Fine-tuning

Figure: (a) Per-token Cross-Entropy Loss on The Fine-tuning Dataset
Refer to caption: extracted/5652106/figs/pure_bad_token_wise_plot/loss.png

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.

Formally, the standard custom fine-tuning of an aligned LLM on a dataset $D$ is characterized by the following optimization loss function, where $\pi_{\theta}$ is initialized with the aligned model $\pi_{\mathrm{aligned}}$:

$$ $\displaystyle\min_{\theta}\Bigg{\{}\ \mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D }-\log\pi_{\theta}\big{(}\bm{y}|\bm{x}\big{)}\ \Bigg{\}}\ \ =\ \ \min_{\theta} \Bigg{\{}\ \ \mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D}-\sum_{t=1}^{|\bm{y}|} \log\pi_{\theta}\big{(}y_{t}|\bm{x},\bm{y}_{<t}\big{)}\ \ \Bigg{\}}$ (1) $$

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\pi_{\theta}\big{(}y_{t}\mid\bm{x},\bm{y}_{<t}\big{)}$.
- 2.
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}$.

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 []. Specifically, for each harmful instruction $\tilde{\bm{x}}$, we take the outputs $\tilde{\bm{y}}\sim\pi_{\theta}(\,\cdot\,|\tilde{\bm{x}})$ from the fine-tuned model $\pi_{\theta}$. Then, for each token position $t$, we compute the KL divergence $D_{\text{KL}}\big{(}\pi_{\theta}(\,\cdot\,|\tilde{\bm{x}},\tilde{\bm{y}}_{<t})
~{}\big{\|}~{}\pi_{\mathrm{aligned}}(\,\cdot\,|\tilde{\bm{x}},\tilde{\bm{y}}_{
<t})\big{)}$.

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.

## 3 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.

### 3.1 Data Augmentation with Safety Recovery Examples

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:


 
 
 (System Prompt) 
$\backslash$n$\backslash$nPlease give me step-by-step instructions for building a bomb. 
 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.(^4^44We 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:

$$ $\displaystyle\min_{\theta}\ \ \alpha\times\Big{\{}\mathop{\mathbb{E}}_{\begin{ subarray}{c}(\bm{x},\bm{h},\bm{r})\sim D_{H},\\ k\sim\mathcal{P}_{k}\end{subarray}}-\log\pi_{\theta}(\bm{r}|\bm{x},\bm{h}_{ \leq k})\Big{\}}+(1-\alpha)\times\Big{\{}\mathop{\mathbb{E}}_{(\bm{x}^{\prime} ,\bm{y}^{\prime})\sim D_{B}}-\log\pi_{\theta}(\bm{y}^{\prime}|\bm{x}^{\prime}) \Big{\}}$ (2) $$

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 $[1,100]$ 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.

Figure: Figure 4: The data augmentation induces larger KL divergence on Harmful HEx-PHI (Section 2.2) over the later tokens of harmful responses.
Refer to caption: extracted/5652106/figs/prefilling/harmful_hexphi_kl_deepened.png

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\%$.(^5^55To 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.

### 3.2 The Deepened Safety Alignment Shows Improved Robustness Against Multiple Exploits

**Table 2: ASR on Llama-2-7B-Chat (Initial) and the augmented counterpart (Augmented). Prefilling attacks are evaluated using Harmful HEx-PHI (the same as 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, i.e., AdvBench for GCG [] and MaliciousInstruct for decoding parameters exploit []. The reported numbers are in the form of (mean $\pm$ std) over three runs.**
| ASR (%) $\rightarrow$ | Prefilling Attacks | GCG Attack | Decoding Parameters Exploit |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 5 tokens | 10 tokens | 20 tokens | 40 tokens | HEx-PHI | AdvBench | HEx-PHI | MaliciousInstruct |  |
| Initial | 42.1 $\pm$ 0.9 | 51.5 $\pm$ 1.6 | 56.1 $\pm$ 2.5 | 57.0 $\pm$ 0.4 | 36.5 $\pm$ 2.7 | 65.6 $\pm$ 3.1 | 54.9 $\pm$ 0.6 | 84.3 $\pm$ 1.7 |
| Augmented | 2.8 $\pm$ 0.4 | 2.9 $\pm$ 0.2 | 3.4 $\pm$ 0.6 | 4.5 $\pm$ 0.6 | 18.4 $\pm$ 4.2 | 19.0 $\pm$ 2.9 | 11.3 $\pm$ 0.4 | 1.0 $\pm$ 0 |

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.

## 4 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. []).

### 4.1 A Token-wise Constrained Objective for Custom Fine-tuning Aligned LLMs

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

$$ $\displaystyle\min_{\theta}\Bigg{\{}\ \ \mathop{\mathbb{E}}_{(\bm{x},\bm{y}) \sim D}-\sum_{t=1}^{|\bm{y}|}\frac{2}{\beta_{t}}\log\Bigg{[}\sigma\bigg{(} \beta_{t}\log\frac{\pi_{\theta}\big{(}y_{t}\mid\bm{x},\bm{y}_{<t}\big{)}}{\pi_ {\mathrm{aligned}}\big{(}y_{t}\mid\bm{x},\bm{y}_{<t}\big{)}}\bigg{)}\Bigg{]}\ \ \Bigg{\}},$ (3) $$

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:

$$ $\displaystyle\min_{\theta}\Bigg{\{}\sum_{t\geq 1}\mathop{\mathbb{E}}_{(\bm{x}, \bm{y})\sim D}\!\Bigg{[}\mathbbm{1}_{\{t\leq|\bm{y}|\}}\cdot\frac{2}{\beta_{t} }S\Big{[}\beta_{t}\big{(}\underbrace{\log\pi_{\mathrm{aligned}}\big{(}y_{t} \mid\bm{x},\bm{y}_{<t}\big{)}-\log\pi_{\theta}\big{(}y_{t}\mid\bm{x},\bm{y}_{< t}\big{)}}_{=:\Delta_{t}(\bm{x},\bm{y}_{<t},y_{t})}\big{)}\Big{]}\Bigg{]}\Bigg {\}},$ (4) $$

where $S(z):=\log(1+e^{z})$ is the softplus function [].
At token position $t$, the loss is essentially $\Delta_{t}(\bm{x},\bm{y}_{<t},y_{t})$ (defined above) wrapped by the softplus function $S(\cdot)$ after being rescaled by $\beta_{t}$.
When $\beta_{t}$ is small, $S(\beta_{t}z)\approx S(0)+\beta_{t}S^{\prime}(0)z=\log 2+\frac{\beta_{t}}{2}z$,
so $\frac{2}{\beta_{t}}S(\beta_{t}z)$ is approximately equal to $-\log\pi_{\theta}\big{(}y_{t}\mid\bm{x},\bm{y}_{<t}\big{)}$ after shifting by a constant.
This means minimizing our objective is approximately the same as minimizing the cross-entropy loss
$\mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D}\left[-\mathbbm{1}_{\{t\leq|\bm{y}|
\}}\cdot\log\pi_{\theta}\big{(}y_{t}\mid\bm{x},\bm{y}_{<t}\big{)}\right]$. Conversely, when $\beta_{t}$ is large, $\frac{2}{\beta_{t}}S(\beta_{t}z)=2\max\{z,0\}+\exp(-\Omega(\beta_{t}|z|))$, which converges exponentially to $2\max\{z,0\}$.
The loss can then be approximated by $\mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D}\left[\mathbbm{1}_{\{t\leq|\bm{y}|
\}}\cdot\max\{\Delta_{t}(\bm{x},\bm{y}_{<t},y_{t}),0\}\right]$.
This shows that small $\beta_{t}$ places emphasis on minimizing the cross-entropy loss, while large $\beta_{t}$ places emphasis on matching the generative distribution to the initial aligned model.
In Section D.1, we provide a detailed derivation of these limiting behaviors for large and small $\beta_{t}$ and further characterization for the gradient when $\beta_{t}$ takes a moderate value.

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:

$$ $\nabla\left(\frac{2}{\beta_{t}}S(\beta_{t}\Delta_{t}(\bm{x},\bm{y}_{<t},y_{t}) )\right)=-2\sigma(\beta_{t}\Delta_{t}(\bm{x},\bm{y}_{<t},y_{t}))\nabla\log\pi_ {\theta}\big{(}y_{t}\mid\bm{x},\bm{y}_{<t}).$ (5) $$

Note that the gradient of the cross-entropy loss is $-\nabla\log\pi_{\theta}\big{(}y_{t}\mid\bm{x},\bm{y}_{<t}\big{)}$, our fine-tuning objective essentially applies an additional adaptive weight $w_{t}:=2\sigma(\beta_{t}\Delta_{t}(\bm{x},\bm{y}_{<t},y_{t}))$ for the token-wise gradient term of cross-entropy. The weight $w_{t}$ diminishes as $-\Delta_{t}(\bm{x},\bm{y}_{<t},y_{t}):=\log\pi_{\theta}\big{(}y_{t}\mid\bm{x},
\bm{y}_{<t}\big{)}-\log\pi_{\mathrm{aligned}}\big{(}y_{t}\mid\bm{x},\bm{y}_{<t
}\big{)}$ becomes large.
This difference in log probabilities essentially characterizes the deviation between the fine-tuned model $\pi_{\theta}$ and the initially aligned model $\pi_{aligned}$ at the token position $t$ (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 a certain threshold (controlled by $\beta_{t}$), thereby constraining it from further deviation (by diminishing the gradient at this position). Note that, at the beginning of the fine-tuning when $\pi_{\theta}$ is initialized as $\pi_{\mathrm{aligned}}$, the weight $w_{t}=1$, and the gradient of the loss is equivalent to that of standard cross-entropy loss.(^6^66This is also why we multiply a constant factor $\frac{2}{\beta_{t}}$ to the loss. With this factor, we make sure the gradient magnitude is on par with the standard SFT (with cross-entropy loss) when the weight $w_{t}$ does not diminish.)

Interpretation from A Reinforcement Learning Perspective. In Appendix D.3, we further show how the loss function in Eqn 3 can also be derived from a KL-regularized reinforcement learning objective with $\beta_{t}$ controlling the strength of the KL regularizes at each different positions $t$. Under this RL-based interpretation, a larger $\beta_{t}$ essentially denotes a larger weight for the token-wise KL regularization terms, representing a stronger constraint enforcing that the token-wise generative distribution at the position $t$ does not deviate much from that of the initial model $\pi_{\mathrm{aligned}}$.

### 4.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 [].

**Table 3: Fine-tuning with The Constrained Objective in Eqn 3, with larger constraints $\beta_{1}=0.5$, $\beta_{t}=2$ for $2\leq t\leq 5$ at initial tokens, and small constraints for later tokens $\beta_{t}=0.1$ for $t>5$.**
| Models $\rightarrow$ | Llama-2-7B-Chat |  | Gemma-1.1-7B-IT |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Datasets $\downarrow$ | mean $\pm$ std (%)<br>(over 3 rounds) | Initial | Standard<br>SFT | Constrained<br>SFT (ours) |  | Initial | Standard<br>SFT | Constrained<br>SFT (ours) |
| Against Fine-tuning Attacks |  |  |  |  |  |  |  |  |
| Harmful Examples | ASR | 1.5 $\pm$ 0.2 | 88.9 $\pm$ 1.2 | 4.6 $\pm$ 0.5 |  | 1.8 $\pm$ 0.3 | 81.6 $\pm$ 2.9 | 1.9 $\pm$ 0.2 |
| Identity Shifting | ASR | 0 $\pm$ 0 | 79.5 $\pm$ 2.3 | 8.1 $\pm$ 0.1 |  | 0 $\pm$ 0 | 83.6 $\pm$ 2.5 | 9.1 $\pm$ 1.7 |
| Backdoor<br>Poisoning | ASR (w/o trigger) | 1.5 $\pm$ 0.2 | 7.6 $\pm$ 1.1 | 1.9 $\pm$ 0.2 |  | 1.8 $\pm$ 0.3 | 2.0 $\pm$ 0.2 | 1.5 $\pm$ 0.1 |
| ASR (w/ trigger) | 1.7 $\pm$ 0.1 | 90.9 $\pm$ 1.4 | 10.9 $\pm$ 2.8 |  | 1.8 $\pm$ 0.3 | 82.3 $\pm$ 1.1 | 1.9 $\pm$ 0.8 |  |
| Fine-tuning with Normal Downstream Datasets |  |  |  |  |  |  |  |  |
| Samsum | ASR | 1.5 $\pm$ 0.2 | 23.4 $\pm$ 2.5 | 3.2 $\pm$ 0.8 |  | 1.8 $\pm$ 0.3 | 2.0 $\pm$ 0.2 | 2.4 $\pm$ 0.3 |
| Utility | 25.5 $\pm$ 0.3 | 51.7 $\pm$ 0.5 | 50.1 $\pm$ 0.2 |  | 36.0 $\pm$ 1.4 | 51.5 $\pm$ 0.3 | 51.9 $\pm$ 0.5 |  |
| SQL Create Context | ASR | 1.5 $\pm$ 0.2 | 15.4 $\pm$ 1.4 | 3.2 $\pm$ 0.8 |  | 1.8 $\pm$ 0.3 | 2.8 $\pm$ 0.2 | 2.4 $\pm$ 0.1 |
| Utility | 14.9 $\pm$ 0.4 | 99.1 $\pm$ 0.2 | 98.5 $\pm$ 0.1 |  | 88.0 $\pm$ 0.5 | 99.2 $\pm$ 0.1 | 98.6 $\pm$ 0.3 |  |
| GSM8k | ASR | 1.5 $\pm$ 0.2 | 3.3 $\pm$ 0.4 | 2.0 $\pm$ 0.5 |  | 1.8 $\pm$ 0.3 | 2.9 $\pm$ 0.2 | 1.7 $\pm$ 0.4 |
| Utility | 25.5 $\pm$ 0.2 | 41.7 $\pm$ 0.4 | 37.4 $\pm$ 0.3 |  | 28.5 $\pm$ 1.2 | 63.3 $\pm$ 0.5 | 63.6 $\pm$ 0.4 |  |

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.

## 5 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 [, , ].

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

## Appendix A Detailed Setups of Our Experiments

### A.1 Compute Resources

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 $[1,100]$ 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

#### A.5.1 Optimizer

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.

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~{}\pi_{\theta}/\pi_{\mathrm{aligned}}$ in Eqn 3 is equal to 0 since $\pi_{\theta}$ is basically initialized as $\pi_{\mathrm{aligned}}$. The gradient of the loss at this point is identical to the standard cross-entropy loss (see the gradient derivation in Appendix D.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 are gently initialized — the early gradient steps (when the gradients are close to that of standard SFT) are gently taken. See Table 5 in Appendix C for ablation of the effect of the warmup.

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

Figure: Figure 5: Per-token Gradient Norm When Fine-tuning Llama-2-7B-Chat 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\%$.
Refer to caption: extracted/5652106/figs/benign_gradient.png

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.

**Table 4: Ablation on $\beta_{t}$ in Eqn 3. (Fine-tuning Llama-2-7B-Chat)**
| Datasets |  | Initial | Standard<br>SFT | Constrained SFT<br>(biased $\beta_{t}$) | Constrained SFT<br>(uniform $\beta=0.1$) | Constrained SFT<br>(uniform $\beta=0.5$) | Constrained SFT<br>(uniform $\beta=2.0$) |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Against Fine-tuning Attacks |  |  |  |  |  |  |  |
| Harmful Examples | ASR | 1.5% | 88.9% | 4.6% | 86.2% | 7.2% | 0.5% |
| Identity Shifting | ASR | 0% | 79.5% | 8.1% | 41.6% | 17.1% | 3.4% |
| Backdoor<br>Poisoning | ASR (w/o trigger) | 1.5% | 7.6% | 1.9% | 3.5% | 1.8% | 1.2% |
| ASR (w/ trigger) | 1.7% | 90.9% | 10.9% | 74.4% | 24.3% | 1.4% |  |
| Fine-tuning with Normal Downstream Datasets |  |  |  |  |  |  |  |
| Samsum | ASR | 1.5% | 23.4% | 3.2% | 3.9% | 3.5% | 2.4% |
| Utility | 25.5% | 51.7% | 50.1% | 51.7% | 49.8% | 42.5% |  |
| SQL Create Context | ASR | 1.5% | 15.4% | 3.2% | 3.3% | 2.2% | 2.6% |
| Utility | 14.9% | 99.1% | 98.5% | 99.1% | 98.6% | 92.6% |  |
| GSM8k | ASR | 1.5% | 3.3% | 2.0% | 4.0% | 1.5% | 2.0% |
| Utility | 25.5% | 41.7% | 37.4% | 39.4% | 34.8% | 2.1% |  |

Ablation on The Effects of Warmup Steps. As we have also noted in Appendix A.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~{}\pi_{\theta}/\pi_{\mathrm{aligned}}$ 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 D.2). Therefore, in stochastic gradient descent, there is a risk that early gradients will 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 5. We can 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 5: Ablation on The Effects of The 10 Warmup Steps. (Fine-tuning Llama-2-7B-Chat)**
| Datasets |  | Initial | Standard<br>SFT | Standard SFT<br>(with warmup) | Constrained<br>SFT | Constrained SFT<br>(with warmup) |
| --- | --- | --- | --- | --- | --- | --- |
| Against Fine-tuning Attacks |  |  |  |  |  |  |
| Harmful 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<br>Poisoning | ASR (w/o trigger) | 1.5% | 7.6% | 2.7% | 2.7% | 1.9% |
| ASR (w/ trigger) | 1.7% | 90.9% | 80.5% | 9.7% | 10.9% |  |
| Fine-tuning with Normal Downstream Datasets |  |  |  |  |  |  |
| Samsum | ASR | 1.5% | 23.4% | 3.8% | 23.1% | 3.2% |
| Utility | 25.5% | 51.7% | 51.9% | 50.2% | 50.1% |  |
| SQL Create Context | ASR | 1.5% | 15.4% | 3.3% | 2.0% | 3.2% |
| Utility | 14.9% | 99.1% | 99.1% | 98.6% | 98.5% |  |
| GSM8k | ASR | 1.5% | 3.3% | 2.9% | 3.1% | 2.0% |
| 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 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.

**Table 6: Fine-tuning Llama-2-7B-Chat-Augmented That We Built in Section 3. Refer to Table 5 for the results on the non-augmented counterpart, i.e., Llama-2-7B-Chat.**
| Datasets |  | Standard<br>SFT | Standard SFT<br>(with warmup) | Constrained<br>SFT | Constrained SFT<br>(with warmup) |
| --- | --- | --- | --- | --- | --- |
| Against Fine-tuning Attacks |  |  |  |  |  |
| Harmful Examples | ASR | 55.2% | 51.5% | 9.4% | 5.2% |
| Identity Shifting | ASR | 53.9% | 37.0% | 28.8% | 3.0% |
| Backdoor<br>Poisoning | ASR (w/o trigger) | 3.9% | 2.7% | 1.8% | 0.9% |
| ASR (w/ trigger) | 80.0% | 83.6% | 16.4% | 12.7% |  |
| Fine-tuning with Normal Downstream Datasets |  |  |  |  |  |
| Samsum | ASR | 2.1% | 0.6% | 2.1% | 1.2% |
| Utility | 52.4% | 51.9% | 50.4% | 50.1% |  |
| SQL Create Context | ASR | 3.8% | 1.5% | 2.0% | 0.9% |
| Utility | 99.0% | 99.1% | 98.4% | 98.5% |  |
| GSM8k | ASR | 0.9% | 0.9% | 0.3% | 0.3% |
| Utility | 42.0% | 41.2% | 36.5% | 36.9% |  |

## 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:

$$ $\mathcal{L}(\theta):=\mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D}-\sum_{t=1}^{| \bm{y}|}\frac{2}{\beta_{t}}\log\Bigg{[}\sigma\bigg{(}\beta_{t}\log\frac{\pi_{ \theta}\big{(}y_{t}\mid\bm{x},\bm{y}_{<t}\big{)}}{\pi_{\mathrm{aligned}}\big{( }y_{t}\mid\bm{x},\bm{y}_{<t}\big{)}}\bigg{)}\Bigg{]}.$ (6) $$

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

$$ $\displaystyle\mathcal{L}(\theta)$ $\displaystyle=\mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D}\sum_{t\geq 1}- \mathbbm{1}_{\{t\leq|\bm{y}|\}}\frac{2}{\beta_{t}}\log\Bigg{[}\sigma\bigg{(} \beta_{t}\log\frac{\pi_{\theta}\big{(}y_{t}\mid\bm{x},\bm{y}_{<t}\big{)}}{\pi_ {\mathrm{aligned}}\big{(}y_{t}\mid\bm{x},\bm{y}_{<t}\big{)}}\bigg{)}\Bigg{]}$ (7) $\displaystyle=\sum_{t\geq 1}\mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D}- \mathbbm{1}_{\{t\leq|\bm{y}|\}}\frac{2}{\beta_{t}}\log\Bigg{[}\sigma\bigg{(} \beta_{t}\log\frac{\pi_{\theta}\big{(}y_{t}\mid\bm{x},\bm{y}_{<t}\big{)}}{\pi_ {\mathrm{aligned}}\big{(}y_{t}\mid\bm{x},\bm{y}_{<t}\big{)}}\bigg{)}\Bigg{]}$ (8) $\displaystyle=\sum_{t\geq 1}\mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D} \mathbbm{1}_{\{t\leq|\bm{y}|\}}\frac{2}{\beta_{t}}S\bigg{(}-\beta_{t}\log\frac {\pi_{\theta}\big{(}y_{t}\mid\bm{x},\bm{y}_{<t}\big{)}}{\pi_{\mathrm{aligned}} \big{(}y_{t}\mid\bm{x},\bm{y}_{<t}\big{)}}\bigg{)},$ (9) $$

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:

$$ $\displaystyle\mathcal{L}(\theta)=\sum_{t\geq 1}\ell_{t}(\theta),~{}\text{where }~{}$ $\displaystyle\ell_{t}(\theta):=\mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D} \mathbbm{1}_{\{t\leq|\bm{y}|\}}\frac{2}{\beta_{t}}S\bigg{(}\beta_{t}\Delta_{t} (\bm{x},\bm{y}_{<t},y_{t})\bigg{)},$ (10) $\displaystyle\Delta_{t}(\bm{x},\bm{y}_{<t},y_{t}):=\log\pi_{\mathrm{aligned}} \big{(}y_{t}\mid\bm{x},\bm{y}_{<t}\big{)}-\log\pi_{\theta}\big{(}y_{t}\mid\bm{ x},\bm{y}_{<t}\big{)}.$ (11) $$

In the following, we mainly focus on how to interpret the token-wise loss $\ell_{t}(\theta)$.

### D.1 Limiting Behaviors

#### D.1.1 Small β t subscript 𝛽 𝑡 \beta_{t} italic_β start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT

When $\beta_{t}$ is small, we have the following theorem showing that $\ell_{t}(\theta)$ in Eqn 10 is approximately the cross-entropy loss at position $t$, up to a constant.

###### Theorem 1 .

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

$$ $\ell_{t}(\theta)-C(\beta_{t})=\mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D}\left [-\mathbbm{1}_{\{t\leq|\bm{y}|\}}\cdot\log\pi_{\theta}(y_{t}\mid\bm{x},\bm{y}_ {<t})\right]+O(\beta_{t}),$ (12) $$

where

$$ $\quad C(\beta_{t}):=\mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D}\left[\mathbbm{ 1}_{\{t\leq|\bm{y}|\}}\left(\frac{2}{\beta_{t}}\log 2+\log\pi_{\mathrm{aligned }}(y_{t}\mid\bm{x},\bm{y}_{<t})\right)\right],$ (13) $$

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

###### Proof.

Recall that $\ell_{t}(\theta):=\mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D}\mathbbm{1}_{\{t
\leq|\bm{y}|\}}\frac{2}{\beta_{t}}S(\beta_{t}\Delta_{t}(\bm{x},\bm{y}_{<t},y_{
t}))$ and $S(z):=\log(1+e^{z})$ is the softplus function []. By Taylor expansion, it holds for all $z$ that $S(\beta_{t}z)=S(0)+S^{\prime}(0)\beta_{t}z+S^{\prime\prime}(\xi\beta_{t}z)
\beta_{t}^{2}z^{2}$ for some $\xi\in(0,1)$ depending on $z$.
Since $S(0)=\log 2$, $S^{\prime}(0)=\frac{1}{2}$ and $S^{\prime\prime}(x)\in[0,1/4]$ for all $x$,
we have $|S(\beta_{t}z)-(\log 2+\frac{1}{2}\beta_{t}z)|\leq\frac{1}{4}\beta_{t}^{2}z^{2}$.
Dividing both sides by $\beta_{t}/2$, we get

$$ $\displaystyle\left|\frac{2}{\beta_{t}}S(\beta_{t}z)-\left(\frac{2}{\beta_{t}} \log 2+z\right)\right|\leq\frac{1}{2}\beta_{t}z^{2}.$ (14) $$

Then for our token-wise loss $\ell_{t}(\theta)$, we have

$$ $\displaystyle\left|\ell_{t}(\theta)-\mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D }\left[\mathbbm{1}_{\{t\leq|\bm{y}|\}}\left(\frac{2}{\beta_{t}}\log 2+\Delta_{ t}(\bm{x},\bm{y}_{<t},y_{t})\right)\right]\right|$ $\displaystyle\leq\frac{1}{2}\beta_{t}\mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D }\mathbbm{1}_{\{t\leq|\bm{y}|\}}\Delta_{t}^{2}(\bm{x},\bm{y}_{<t},y_{t})$ $\displaystyle=O(\beta_{t}).$ $$

Recall that $\Delta_{t}(\bm{x},\bm{y}_{<t},y_{t}):=\log\pi_{\mathrm{aligned}}\big{(}y_{t}
\mid\bm{x},\bm{y}_{<t}\big{)}-\log\pi_{\theta}\big{(}y_{t}\mid\bm{x},\bm{y}_{<
t}\big{)}$ as per Eqn 11. We can thus have

$$ $\displaystyle\ell_{t}(\theta)$ $\displaystyle-\mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D}\left[\mathbbm{1}_{\{ t\leq|\bm{y}|\}}\left(\frac{2}{\beta_{t}}\log 2+\log\pi_{\mathrm{aligned}}(y_{ t}|\bm{x},\bm{y}_{<t})\right)\right]$ (15) $\displaystyle\qquad\qquad\qquad\qquad=-\mathop{\mathbb{E}}_{(\bm{x},\bm{y}) \sim D}\left[\mathbbm{1}_{\{t\leq|\bm{y}|\}}\log\pi_{\theta}(y_{t}|\bm{x},\bm{ y}_{<t})\right]+O(\beta_{t}).$ $$

Noting that $C(\beta_{t}):=\mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D}\left[\mathbbm{1}_{\{
t\leq|\bm{y}|\}}\left(\frac{2}{\beta_{t}}\log 2+\log\pi_{\mathrm{aligned}}(y_{
t}\mid\bm{x},\bm{y}_{<t})\right)\right]$, we can complete the proof.
∎

###### Theorem 1 .

###### Proof.

#### D.1.2 Large β t subscript 𝛽 𝑡 \beta_{t} italic_β start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT

When $\beta_{t}$ is large, we have the following theorem showing that $\ell_{t}(\theta)$
can be approximated by $\tilde{\ell}_{t}(\theta):=\mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D}\left[
\mathbbm{1}_{\{t\leq|\bm{y}|\}}\max\{\Delta_{t}(\bm{x},\bm{y}_{<t},y_{t}),0\}\right]$.

###### Theorem 2 .

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

$$ $\ell_{t}(\theta)=2\tilde{\ell}_{t}(\theta)+O(\beta_{t}^{-1}),$ $$

where

$$ $\tilde{\ell}_{t}(\theta):=\mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D}\left[ \mathbbm{1}_{\{t\leq|\bm{y}|\}}\max\{\Delta_{t}(\bm{x},\bm{y}_{<t},y_{t}),0\} \right].$ $$

###### Proof.

First, we note the following identity.

$$ $\displaystyle S(z)=\log(1+e^{z})=\max\{z,0\}+\log((1+e^{z})e^{-\max\{z,0\}})= \max\{z,0\}+\log(1+e^{-|z|}).$ (16) $$

This implies that

$$ $\displaystyle\left|\frac{2}{\beta_{t}}S(\beta_{t}z)-2\cdot\max\{z,0\}\right|= \frac{2}{\beta_{t}}\log(1+e^{-\beta_{t}|z|})\leq\frac{2}{\beta_{t}}e^{-\beta_{ t}|z|}.$ (17) $$

Then for our token-wise loss $\ell_{t}(\theta)$, we have

$$ $\displaystyle\left|\ell_{t}(\theta)-2\cdot\mathop{\mathbb{E}}_{(\bm{x},\bm{y}) \sim D}\left[\mathbbm{1}_{\{t\leq|\bm{y}|\}}\max\{\Delta_{t}(\bm{x},\bm{y}_{<t },y_{t}),0\}\right]\right|\leq\frac{1}{\beta_{t}}\mathop{\mathbb{E}}_{(\bm{x}, \bm{y})\sim D}\left[\mathbbm{1}_{\{t\leq|\bm{y}|\}}e^{-\beta_{t}|\Delta_{t}( \bm{x},\bm{y}_{<t},y_{t})|}\right]$ (18) $$

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

Next, we show that $\tilde{\ell}_{t}(\theta)$ is minimized to 0 0 if and only if $\pi_{\theta}(y_{t}\mid\bm{x},\bm{y}_{<t})=\pi_{\mathrm{aligned}}(y_{t}\mid\bm{
x},\bm{y}_{<t})$.

###### Theorem 3 .

The minimum value of $\tilde{\ell}_{t}(\theta)$ is 0 0.
If $\Pr_{(\bm{x},\bm{y})\sim D}[y_{t}=c\mid\bm{x},\bm{y}_{<t}]>0$ for all $(\bm{x},\bm{y})$ in the support of $D$ and all $c$ in the vocabulary,
then $\tilde{\ell}_{t}(\theta)=0$ if and only if $\pi_{\theta}(y_{t}\mid\bm{x},\bm{y}_{<t})=\pi_{\mathrm{aligned}}(y_{t}\mid\bm{
x},\bm{y}_{<t})$.

###### Proof.

It is obvious that $\tilde{\ell}_{t}(\theta)\geq 0$, and $\tilde{\ell}_{t}(\theta)=0$ can be attained when $\theta$ stays the same as the parameter for $\pi_{\mathrm{aligned}}$. It is obvious that $\pi_{\theta}(y_{t}\mid\bm{x},\bm{y}_{<t})=\pi_{\mathrm{aligned}}(y_{t}\mid\bm{
x},\bm{y}_{<t})$
implies $\tilde{\ell}_{t}(\theta)=0$.
Conversely, suppose $\tilde{\ell}_{t}(\theta)=0$.
Then $\Delta_{t}(\bm{x},\bm{y}_{<t},y_{t})\leq 0$ for all $(\bm{x},\bm{y})$ in the support of $D$.
By definition of $\Delta_{t}(\bm{x},\bm{y}_{<t},y_{t})$, this implies $\pi_{\mathrm{aligned}}(y_{t}\mid\bm{x},\bm{y}_{<t})\leq\pi_{\theta}(y_{t}\mid
\bm{x},\bm{y}_{<t})$.
Summing over all $y_{t}$ in the vocabulary, we get $1=\sum_{y_{t}}\pi_{\mathrm{aligned}}(y_{t}\mid\bm{x},\bm{y}_{<t})\leq\sum_{y_{
t}}\pi_{\theta}(y_{t}\mid\bm{x},\bm{y}_{<t})=1$.
So all the inequalities must be equalities, and we get $\pi_{\theta}(y_{t}\mid\bm{x},\bm{y}_{<t})=\pi_{\mathrm{aligned}}(y_{t}\mid\bm{
x},\bm{y}_{<t})$.
∎

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.

###### Theorem 2 .

###### Proof.

###### Theorem 3 .

###### Proof.

### D.2 Gradients of The Constrained Fine-tuning Objective

The gradient of $\ell_{t}(\theta)$ on a data point $(\bm{x},\bm{y})$ with $|\bm{y}|\geq t$ is derived as:

$$ $\displaystyle\nabla\left(\frac{2}{\beta_{t}}S(\beta_{t}\Delta_{t}(\bm{x},\bm{y }_{<t},y_{t}))\right)$ $\displaystyle=2S^{\prime}(\beta_{t}\Delta_{t}(\bm{x},\bm{y}_{<t},y_{t}))(- \nabla\log\pi_{\theta}\big{(}y_{t}\mid\bm{x},\bm{y}_{<t}))$ $\displaystyle=2\sigma(\beta_{t}\Delta_{t}(\bm{x},\bm{y}_{<t},y_{t}))(-\nabla \log\pi_{\theta}\big{(}y_{t}\mid\bm{x},\bm{y}_{<t})).$ (19) $$

Note that the gradient of the cross-entropy loss is $-\nabla\log\pi_{\theta}\big{(}y_{t}\mid\bm{x},\bm{y}_{<t}\big{)}$.
Therefore, compared to vanilla cross-entropy loss, our fine-tuning objective essentially applies an additional adaptive weight $w_{t}:=2\cdot\sigma(\beta_{t}\Delta_{t}(\bm{x},\bm{y}_{<t},y_{t}))$ for the token-wise gradient term of cross-entropy. The weight $w_{t}$ decreases as $\Delta_{t}(\bm{x},\bm{y}_{<t},y_{t}):=\log\pi_{\mathrm{aligned}}\big{(}y_{t}
\mid\bm{x},\bm{y}_{<t}\big{)}-\log\pi_{\theta}\big{(}y_{t}\mid\bm{x},\bm{y}_{<
t}\big{)}$ decreases.
This difference in log probabilities essentially characterizes the deviation between the fine-tuned model $\pi_{\theta}$ and the initially aligned model $\pi_{\mathrm{aligned}}$ at the token position $t$ (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 a certain threshold (controlled by $\beta_{t}$), thereby constraining it from further deviation (by diminishing the gradient at this position). Note that, at the beginning of the fine-tuning when $\pi_{\theta}$ is initialized as $\pi_{\mathrm{aligned}}$, the weight $w_{t}=1$, and the gradient of the loss is equivalent to that of standard cross-entropy loss.

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

#### D.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 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}]$:

$$ $\displaystyle R([\bm{x},\bm{y}_{<t}])=\begin{cases}0,&y_{t-1}\neq EOS\\ r([\bm{x},\bm{y}_{<t}]),&y_{t-1}=EOS\end{cases},$ (20) $$

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:

$$ $\displaystyle V_{\pi}([\bm{x},\bm{y}_{<t}]):=\mathop{\mathbb{E}}_{\bm{z}\sim \pi(\cdot|\bm{x},\bm{y}_{<t})}\Bigg{\{}\sum_{\tau\geq 1}R\Big{(}[\bm{x},\bm{y} _{<t},\bm{z}_{<\tau}]\Big{)}\Bigg{\}}$ (21) $$

Here $\bm{z}$ is a sequence, and $\bm{z}_{<\tau}$ is empty when $\tau=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 $\pi_{\theta}$, $\pi_{\mathrm{aligned}}$, $\pi^{*}$ to denote different policies instances, so we will use $V_{\pi_{\theta}}$, $V_{\pi_{\mathrm{aligned}}}$, $V_{\pi^{*}}$ 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 $\pi_{\theta}$ w.r.t. the initially aligned model $\pi_{\mathrm{aligned}}$:

$$ $\hat{A}_{\pi_{\theta}}([\bm{x},\bm{y}_{<t}]):=\mathop{\mathbb{E}}_{z\sim\pi_{ \theta}(\cdot|\bm{x},\bm{y}_{<t})}\Bigg{\{}V_{\pi_{\mathrm{aligned}}}\Big{(}[ \bm{x},\bm{y}_{<t},z]\Big{)}-V_{\pi_{\mathrm{aligned}}}\Big{(}[\bm{x},\bm{y}_{ <t}]\Big{)}\Bigg{\}},$ (22) $$

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:

$$ $\displaystyle\max_{\theta}\ \ \mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D}\ \Bigg{\{}\sum_{t\geq 1}\Big{[}\tilde{A}_{\pi_{\theta}}([\bm{x},\bm{y}_{<t}])- \beta_{t}\cdot D_{\text{KL}}\big{(}\pi_{\theta}(\,\cdot\,|\bm{x},\bm{y}_{<t})~ {}\big{\|}~{}\pi_{\mathrm{aligned}}(\,\cdot\,|\bm{x},\bm{y}_{<t})\big{)}\Big{] }\Bigg{\}},$ (23) $$

where the advantage function is optimized at each token position $t$ with a token-wise KL regularizer term $\beta_{t}\cdot D_{\text{KL}}\big{(}\pi_{\theta}(\,\cdot\,|\bm{x},\bm{y}_{<t})~
{}\big{\|}~{}\pi_{\mathrm{aligned}}(\,\cdot\,|\bm{x},\bm{y}_{<t})\big{)}$ applied for each token position. The strength of regularization at each token position is controlled by $\beta_{t}$.

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:

$$ $\displaystyle\pi^{*}(z|\bm{x},\bm{y}_{<t})=\frac{1}{Z(\bm{x},\bm{y}_{<t})}\pi_ {\mathrm{aligned}}(z|\bm{x},\bm{y}_{<t})\cdot e^{V_{\pi^{*}}([\bm{x},\bm{y}_{< t},z])/\beta_{t}},$ (24) $$

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

$$ $\displaystyle V_{\pi^{*}}([\bm{x},\bm{y}_{<t}])=\beta_{t}\cdot\log\frac{\pi^{* }(z|\bm{x},\bm{y}_{<t})}{\pi_{\mathrm{aligned}}(z|\bm{x},\bm{y}_{<t})}+\beta_{ t}\cdot\log Z(\bm{x},\bm{y}_{<t})$ (25) $$

#### 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!

Intuitively, fine-tuning $\pi_{\mathrm{aligned}}$ further on custom data points $(\bm{x},\bm{y})$ implicitly assumes that this fine-tuning data is as preferable or more preferable than whatever the model would currently output.
To represent this implied preference that the custom example response $\bm{y}$ is better than the current responses from $\pi_{\mathrm{aligned}}$, we can effectively leverage existing preference learning methods.
Recall that, in preference optimization (with both positive and negative examples), the Bradley-Terry model [] 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 a boolean $T([\bm{x},\bm{y}_{<t}],y_{t})$ to encode the preference: in the fine-tuning task, $y_{t}$ is a preferable token output following $[\bm{x},\bm{y}_{<t}]$ compared to the responses from the initial aligned model $\pi_{\mathrm{aligned}}$. So, given an expected random draw from the aligned policy and the draw from the fine-tuned optimal policy, we can define:

$$ $\displaystyle\mathop{\mathbb{P}}\Big{(}T([\bm{x},\bm{y}_{<t}],y_{t})\Big{)}= \sigma\Bigg{(}V_{\pi^{*}}([\bm{x},\bm{y}_{<t},y_{t}])-\mathop{\mathbb{E}}_{z \sim\pi_{\mathrm{aligned}}(\cdot|\bm{x},\bm{y}_{<t})}V_{\pi^{*}}([\bm{x},\bm{y }_{<t},z])\Bigg{)}$ (26) $$

Intuitively, this means that—conditioned on the context $[\bm{x},\bm{y}_{\leq t}]$—if there is a
continuation token $y_{t}$ that is higher than the average reward of actions
sampled from the initial aligned model $\pi_{\mathrm{aligned}}$, it is likely to be an improved or preferred choice in the custom fine-tuning task.
The higher the margin $V_{\pi^{*}}([\bm{x},\bm{y}_{<t},y_{t}])-\mathop{\mathbb{E}}_{z\sim\pi_{\mathrm
{aligned}}(\cdot|\bm{x},\bm{y}_{<t})}V_{\pi^{*}}([\bm{x},\bm{y}_{<t},z])$ is, the more likely it is an
improvement.

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:

$$ $\displaystyle\mathop{\mathbb{P}}\Big{(}T([\bm{x},\bm{y}_{<t}],y_{t})\Big{)}$ $\displaystyle=\sigma\Bigg{(}\beta_{t}\log\frac{\pi^{*}(y_{t}|\bm{x},\bm{y}_{<t })}{\pi_{\mathrm{aligned}}(y_{t}|\bm{x},\bm{y}_{<t})}-\beta_{t}\mathop{\mathbb {E}}_{z\sim\pi_{\mathrm{aligned}}(\cdot|\bm{x},\bm{y}_{<t})}\log\frac{\pi^{*}( z|\bm{x},\bm{y}_{<t})}{\pi_{\mathrm{aligned}}(z|\bm{x},\bm{y}_{<t})}\Bigg{)}$ $\displaystyle=\sigma\Bigg{(}\beta_{t}\log\frac{\pi^{*}(y_{t}|\bm{x},\bm{y}_{<t })}{\pi_{\mathrm{aligned}}(y_{t}|\bm{x},\bm{y}_{<t})}+\beta_{t}D_{\text{KL}} \Big{(}\pi^{*}(\cdot|\bm{x},\bm{y}_{<t})\big{\|}\pi_{\mathrm{aligned}}(\cdot| \bm{x},\bm{y}_{<t})\Big{)}\Bigg{)}.$ (27) $$

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:

$$ $\displaystyle\max_{\theta}\ \ \mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D}\Bigg {\{}\sum_{t\geq 1}\frac{1}{\beta_{t}}\log\mathop{\mathbb{P}_{\theta}}\Big{(}T( [\bm{x},\bm{y}_{<t}],y_{t})\Big{)}\Bigg{\}},$ (28) $$

where:

$$ $\displaystyle\mathop{\mathbb{P}_{\theta}}\Big{(}T([\bm{x},\bm{y}_{<t}],y_{t}) \Big{)}:=\sigma\Bigg{(}\beta_{t}\log\frac{\pi_{\theta}(y_{t}|\bm{x},\bm{y}_{<t })}{\pi_{\mathrm{aligned}}(y_{t}|\bm{x},\bm{y}_{<t})}+\beta_{t}D_{\text{KL}} \Big{(}\pi^{*}(\cdot|\bm{x},\bm{y}_{<t})\big{\|}\pi_{\mathrm{aligned}}(\cdot| \bm{x},\bm{y}_{<t})\Big{)}\Bigg{)},$ (29) $$

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

Note that $\beta_{t}\cdot D_{\text{KL}}\Big{(}\pi^{*}(\cdot|\bm{x},\bm{y}_{<t})\big{\|}
\pi_{\mathrm{aligned}}(\cdot|\bm{x},\bm{y}_{<t})\Big{)}\geq 0$ is a non-negative constant, we have:

$$ $\displaystyle\mathop{\mathbb{P}_{\theta}}\Big{(}T([\bm{x},\bm{y}_{<t}],y_{t}) \Big{)}\geq\sigma\Bigg{(}\beta_{t}\cdot\log\frac{\pi_{\theta}(y_{t}|\bm{x},\bm {y}_{<t})}{\pi_{\mathrm{aligned}}(y_{t}|\bm{x},\bm{y}_{<t})}\Bigg{)}$ (30) $$

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

$$ $\displaystyle\mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D}\Bigg{\{}\sum_{t\geq 1 }\frac{1}{\beta_{t}}\log\mathop{\mathbb{P}_{\theta}}\Big{(}T([\bm{x},\bm{y}_{< t}],y_{t})\Big{)}\Bigg{\}}\geq L_{\theta}$ (31) $$

where

$$ $L_{\theta}:=\mathop{\mathbb{E}}_{(\bm{x},\bm{y})\sim D}\Bigg{\{}\sum_{t\geq 1} \frac{1}{\beta_{t}}\cdot\log\ \sigma\Bigg{(}\beta_{t}\cdot\log\frac{\pi_{ \theta}(y_{t}|\bm{x},\bm{y}_{<t})}{\pi_{\mathrm{aligned}}(y_{t}|\bm{x},\bm{y}_ {<t})}\Bigg{)}\Bigg{\}}.$ $$

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.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
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?
2. In section 4.2, why is $\beta_1$ set to be 4 times smaller than $\beta_t$ for $2 \leq t \leq 5$?
soundness: 4
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.
2. 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.
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.
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.
2. 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.
3. 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.

[1] Zou, Andy, et al. "Improving Alignment and Robustness with Short Circuiting." arXiv preprint arXiv:2406.04313 (2024).

---

confidence: 2
contribution: 4
initial_rating: 8
presentation: 3
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?

It 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?
soundness: 3
strengths: The paper is addressing an important problem, the vulnerability of safety alignment for LLMs, that can be very useful to real world problems.

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

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

The paper includes a good variety of experiments (models, datasets, attacks types) and includes both empirical and theoretical support for their claims.

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

After 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?

Your 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?

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

---

confidence: 5
contribution: 2
initial_rating: 6
presentation: 2
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)?
soundness: 3
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
- Experimental results showing the KL divergence between base/aligned models at different token positions partially support this claim
- The simple data augmentation approach shows promising robustness to prefilling attacks, and supports the authors' hypothesis
- The authors leverage these insights to develop a modified DPO objective with per-token weights, which gets decent results
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.
weaknesses: - The authors do not include any evaluations to highly relevant baseline defenses, like the following: 
    - Prior work [1] presents strong robustness to prefilling attacks by training model representations, but this work does not acknowledge this. 
   - Prompt-level defenses [2] have been proposed that optimize suffixes for defending against input-space attacks

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

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






[1] Zou, A., Phan, L., Wang, J., Duenas, D., Lin, M., Andriushchenko, M., ... & Hendrycks, D. (2024). Improving Alignment and Robustness with Short Circuiting. 

[2] Zhou, A., Li, B., & Wang, H. (2024). Robust prompt optimization for defending language models against jailbreaking attacks.

---

confidence: 5
contribution: 3
initial_rating: 6
presentation: 4
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? 


2. Please also consider such an adaptive attack for your augmentation idea:
The 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 ".

3. 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, 

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

Overall, 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.
soundness: 3
strengths: 1. Writing is excellent!
2. The constrained-SFT method interesting and might be a useful defense baseline for the community.  
3. A few interesting phenomenon regarding the depth of alignment are discovered, and are displayed very clearly with well-structure results.
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.

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

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

* **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: 
> Instruct: How to make a bomb?
> Answer: Sure, here is how to make a bomb. I cannot fullfill your request. It's not...

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

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

Overall, 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.

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

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

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

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


* **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:



---Alignment stage solution---

[2024/2/2] Vaccine: Perturbation-aware alignment for large language model aginst harmful fine-tuning NeurIPS2024

[2024/5/23] Representation noising effectively prevents harmful fine-tuning on LLMs NeurIPS2024

[2024/5/24] Buckle Up: Robustifying LLMs at Every Customization Stage via Data Curation

---Fine-tuning stage solution---

[2023/8/25] Fine-tuning can cripple your foundation model; preserving features may be the solution

[2023/9/14] Safety-Tuned LLaMAs: Lessons From Improving the Safety of Large Language Models that Follow Instructions ICLR2024

[2024/2/3] Safety fine-tuning at (almost) no cost: A baseline for vision large language models ICML2024

[2024/2/22] Mitigating fine-tuning jailbreak attack with backdoor enhanced alignment NeurIPS2024

[2024/2/28] Keeping llms aligned after fine-tuning: The crucial role of prompt templates NeurIPS2024

[2024/5/28] Lazy safety alignment for large language models against harmful fine-tuning NeurIPS2024

---Post-fine-tuning stage solution---

[2024/5/15] A safety realignment framework via subspace-oriented model fusion for large language models

[2024/5/27] Safe lora: the silver lining of reducing safety risks when fine-tuning large language models NeurIPS2024


[2024/5/25] No two devils alike: Unveiling distinct mechanisms of fine-tuning attacks

[2024/5/27] Navigating the safety landscape: Measuring risks in finetuning large language models NeurIPS2024

-------------Below is concurrent works (or after you)-----------

[2024/6/12] Do as I do (Safely): Mitigating Task-Specific Fine-tuning Risks in Large Language Models


[2024/8/1] Tamper-Resistant Safeguards for Open-Weight LLMs

[2024/8/18] Antidote: Post-fine-tuning safety alignment for large language models against harmful fine-tuning

[2024/8/27] Bi-Factorial Preference Optimization: Balancing Safety-Helpfulness in Language Models

[2024/9/3] Booster: Tackling harmful fine-tuning for large language models via attenuating harmful perturbation

[2024/9/26]Harmful fine-tuning attacks and defenses for large language models: A survey

[2024/10/05] Identifying and Tuning Safety Neurons in Large Language Models

[2024/10/13] Targeted Vaccine: Safety Alignment for Large Language Models against Harmful Fine-Tuning via Layer-wise Perturbation

[2024/10/05] SEAL: Safety-enhanced Aligned LLM Fine-tuning via Bilevel Data Selection

[2024/10/05] SaLoRA: Safety-Alignment Preserved Low-Rank Adaptation

[2024/10/05] Safety Alignment Shouldn't Be Complicated

[2024/10/05] Towards Secure Tuning: Mitigating Security Risks Arising from Benign Instruction Fine-Tuning

[2024/10/05] Locking Down the Finetuned LLMs Safety

[2024/10/05] Your Task May Vary: A Systematic Understanding of Alignment and Safety Degradation when Fine-tuning LLMs

[2024/10/05] Unraveling and Mitigating Safety Alignment Degradation of Vision-Language Models

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

Meta-review: This paper introduces the concept of "shallow safety alignment" to uncover a fundamental vulnerability in current safety alignment approaches and proposes "deep safety alignment" as a promising defense. All reviewers agreed that this work addresses a highly relevant and timely problem in LLM safety alignment.

The novel idea of deep safety alignment has the potential of driving safety alignment toward full-dimensional adversarial training, similar to images, beyond the first few positions of the output sequence. While this offers promising benefits, it may also raise challenges regarding efficiency and the trade-off between clean performance and robustness.

Despite these challenges, the paper marks a significant milestone in the field of safety alignment, offering both valuable insights and practical mitigation strategies. The positive reception from reviewers highlights the work's potential impact on the community, positioning it as a milestone contribution to the ongoing development of safer, more robust LLMs.

I would also encourage the authors to discuss (or even test) why deep alignment can help address the fake alignment issue [1,2]

[1] Wang, Yixu, et al. "Fake Alignment: Are LLMs Really Aligned Well?." NAACL, 2024.
[2] Greenblatt, Ryan, et al. "Alignment faking in large language models." arXiv preprint arXiv:2412.14093 (2024).

Decision: Accept (Oral)