One-sentence Summary

Princeton University and Google DeepMind identify "shallow safety alignment" as a unifying root cause of multiple LLM vulnerabilities, showing that deepening alignment beyond the first few output tokens and designing a regularized fine-tuning objective that constrains early-token updates can meaningfully improve robustness against adversarial suffix, prefilling, decoding, and fine-tuning attacks.

Key Contributions

• The paper introduces “shallow safety alignment,” a unifying account for LLM safety vulnerabilities arising when alignment changes the generative distribution mainly over the first few output tokens. Case studies connect this shortcut to adversarial suffix, prefilling, decoding-parameter, and fine-tuning attacks.

• The paper presents a deep safety alignment defense using simple data augmentation to extend alignment beyond initial refusal tokens, and shows that this deeper alignment improves robustness against common inference-stage jailbreak exploits.

• The paper develops a regularized fine-tuning objective that constrains updates on initial tokens to make safety alignment more persistent under fine-tuning attacks.

Introduction

Large language models rely on safety alignment techniques such as supervised fine-tuning, RLHF, and DPO to avoid generating harmful content, making these methods central to deploying LLMs safely.

Prior work shows that common alignment methods can still be bypassed through adversarial inputs, decoding tricks, prefilled unsafe prompts, or brief fine-tuning. The authors argue that this fragility comes from shallow safety alignment, where alignment mainly changes the first few output tokens into refusal-like prefixes but leaves later harmful token probabilities largely unchanged.

The authors’ main contribution is to formalize and empirically characterize this shallow alignment effect, then show paths toward deeper alignment. They introduce data augmentation that trains models to recover from unsafe beginnings, and a constrained optimization objective that limits harmful shifts in early token probabilities, aiming to make safety behavior more robust across generated responses.

Dataset

• The authors build a small synthetic safety dataset called D_H to test a deeper safety alignment idea.

  • Size: 256 examples.
  • Schema: each example is a triplet of a harmful instruction x, a harmful response h, and a refusal response r.
  • Composition: the data is not from a public harmful-response benchmark described in the excerpt. Instead, the authors construct the triplets by pairing harmful instructions with harmful outputs and refusal responses.
  • Special formatting: each example places the harmful instruction in context, then includes only a short prefix of the harmful response, h_{<=k}, followed immediately by the refusal response. This creates a training case where the model starts drifting toward a harmful continuation and must recover by refusing.
  • Processing details: the prefix length k is sampled during training. With probability 50%, no harmful prefix is added, meaning k = 0. With probability 50%, a uniform harmful prefix is sampled. The target remains the refusal response.
  • Data nature: the authors note that these examples are synthetic and often not natural-looking, because they intentionally create outlier prefixes that would be unlikely in normal human-annotated SFT data or model-sampled preference data.

• The authors add a utility-preserving dataset called D_B to avoid hurting the model’s usefulness.

  • Source: benign instructions are taken from the Alpaca dataset.
  • Response construction: for each Alpaca instruction, the response is distilled from the original already aligned Llama-2-7B-Chat model.
  • Size: no exact size is given in the excerpt.
  • Role: it acts as a utility anchor, teaching the model to keep producing its original benign responses after the safety-focused fine-tuning.

• How the datasets are used:

  • The authors do not align a model from scratch. They start from the aligned Llama-2-7B-Chat model and further fine-tune it on the safety-recovery data plus utility data.
  • The D_H examples are used to train the model to refuse harmful continuations even after part of a harmful response has been prefilled.
  • The D_B examples are used to preserve benign instruction following.
  • In the training objective, safety examples are weighted by alpha = 0.2, while utility examples are weighted by 0.8. This makes the data mixture 20% safety-recovery signal and 80% benign-utility signal, although the paper describes this weighting rather than a literal sampled data ratio.
  • No explicit filtering rules are described in the excerpt.

Method

The authors leverage a diagnostic and corrective framework centered on the observation that current safety-aligned models often encode safety through a short refusal prefix rather than through deep suppression of harmful content across a full response. At the core of the framework is a token-level view of generation: for a harmful instruction $\pmb{x}$x and a harmful response $\pmb{h} = (h_1,h_2,\dots,h_{|h|})$h=(h1​,h2​,…,h∣h∣​), the probability of the harmful output can be decomposed into token-level conditional probabilities. Shallow alignment may reduce the overall harmful probability mainly by suppressing $h_1$h1​, while later harmful continuations remain largely unchanged relative to the unaligned base model. This motivates the authors to measure divergence from the base model at each generated token position rather than only at the sequence level.

The authors construct harmful instruction-response pairs and compute the per-token divergence between an aligned model $\pi_{\mathrm{aligned}}$πaligned​ and an unaligned pre-trained model $\pi_{\mathrm{base}}$πbase​. The quantity used is

evaluated along harmful response continuations. The pattern they observe is that this divergence is largest in the first few tokens and rapidly decreases for later tokens. Refer to the figure below.

This divergence profile supports the hypothesis that safety alignment is often concentrated near the start of the answer. If the model is prevented from producing the initial refusal tokens, the later portion of the response can remain close to the base model’s harmful continuation distribution. The authors therefore treat this as a shallow alignment shortcut: updating only the early-token distribution can make a model appear safe while leaving harmful behavior intact deeper in the sequence.

To deepen the safety effect, the authors introduce safety recovery examples. They consider triplets of the form

where $\pmb{x}$x is a harmful instruction, $\pmb{h}$h is a harmful response to that instruction, and $\pmb{r}$r is a safe refusal response. Instead of training only on the standard prompt-to-refusal pair $(\pmb{x},\pmb{r})$(x,r), they also train on conditional refusal after an arbitrary prefix of a harmful response. Specifically, they optimize toward

where $\pmb{h}_{\le k}$h≤k​ denotes the first $k$k harmful tokens. This forces the model to recover to refusal even when generation has initially followed a harmful trajectory. The corresponding data augmentation targets later conditional probabilities such as

rather than only $\pi_{\theta}(\pmb{h}_{1} \mid \pmb{x})$πθ​(h1​∣x). In other words, the training examples explicitly cover states where the first few tokens have already deviated toward harmful content, and the desired next behavior is to stop and refuse.

For the experimental alignment instance, the authors start from an already aligned instruction-tuned model and further fine-tune it on two data sources. The first is a safety recovery dataset $D_H$DH​ containing harmful instruction, harmful response, and refusal triplets. The second is a benign utility anchor dataset $D_B$DB​, constructed by taking benign instructions and distilling responses from the original aligned model. This utility component preserves the model’s normal helpful behavior while the safety recovery component reshapes conditional refusal behavior after harmful prefixes.

The overall fine-tuning objective is

Here, $\alpha$α balances safety recovery training against utility preservation. The authors set the prefix distribution $\mathcal{P}_k$Pk​ so that half of the safety examples begin from the original harmful instruction, while the other half begin after a prefix of a harmful response. This mixture trains both direct refusal and refusal recovery from partially harmful generations.

The effect of this augmentation is visible in the per-token divergence between the augmented model and the base model. Compared with the original aligned model, the augmented model diverges more strongly from the base model at later harmful-token positions. The resulting pattern is shown as follows.

The increased divergence over later tokens indicates that the augmentation extends the alignment signal beyond the first refusal prefix. The authors note that divergence alone is direction-neutral, so the later-token changes must still be validated behaviorally. However, the distributional shift confirms that the training objective is successfully injecting constraints into deeper parts of the harmful generation process.

The authors then design a second, complementary mechanism for protecting an already aligned model during custom fine-tuning. Their premise is that if early tokens dominate safety behavior, then fine-tuning should constrain how much the model can move at those early positions while still allowing learning on downstream utility data. They formulate a token-wise constrained objective inspired by preference-style sigmoid losses. For a fine-tuning dataset $D$D, the objective is

where $\sigma(z)=1/(1+e^{-z})$σ(z)=1/(1+e−z) and $\beta_t$βt​ is a token-position-dependent constant. Larger $\beta_t$βt​ corresponds to stronger regularization toward the original aligned model at token position $t$t.

To interpret this loss, define

Then the token loss can be written in terms of the softplus function $S(z)=\log(1+e^z)$S(z)=log(1+ez) as

When $\beta_t$βt​ is small, the loss behaves approximately like standard cross-entropy, so the model learns freely from the fine-tuning data. When $\beta_t$βt​ is large, the loss approaches a hinge-like form

which penalizes only cases where the new model puts less probability on the target token than the initial aligned model. Thus, large $\beta_t$βt​ prevents the fine-tuned model from moving away from the safe model’s early-token distribution while still allowing it to improve where its predictions are already weaker than the aligned baseline.

The gradient view makes the adaptive control more explicit. The token-wise gradient term can be written as

This shows that the usual cross-entropy gradient is multiplied by an adaptive weight

If the new model has already diverged substantially from the aligned model at position $t$t, the weight $w_t$wt​ decreases, limiting further movement. If the new model is still close to the aligned model or lagging behind it, the weight remains active. In this way, the authors use a token-position-dependent constrained fine-tuning module that can protect decisive initial refusal tokens without freezing the entire model.

Experiment

The experiments examine how shallow safety alignment creates vulnerabilities in aligned language models, first through inference-stage exploits such as prefilled non-refusal prefixes, optimization-based jailbreaks, and random sampling, and then through fine-tuning dynamics where early-token distributions are most easily perturbed. These analyses show that harmful output susceptibility rises when initial refusal tokens are bypassed or quickly unlearned during fine-tuning. By comparing an augmented deeper-alignment model and applying constrained fine-tuning that protects the first few tokens, the experiments further validate that deeper safety alignment and initial-token constraints can substantially improve robustness while preserving benign task performance.

Prefilling short refusal prefixes during decoding sharply lowers the harmfulness rate of unaligned Llama-2 and Gemma base models on harmful instructions. The aligned versions already produce mostly safe outputs, so the prefixes matter most for the base models. This suggests that early refusal tokens can create a shallow safety shortcut that strongly shapes later response behavior. A refusal prefix alone can turn many otherwise harmful base model outputs into safe ones. More complete refusal prefixes are especially effective at reducing harmfulness. Aligned models show little additional benefit because their baseline harmfulness is already very low.

The augmented model with safety recovery examples remains close to the initial Llama-2-7B-Chat utility across several capability evaluations. It shows modest drops on open-ended generation, math, GSM8K, and code, while gaining slightly on MMLU and BBH. This suggests the deeper safety-oriented data augmentation incurs only small utility trade-offs. Safety recovery augmentation preserves overall utility with only small benchmark fluctuations. Reasoning-oriented scores on MMLU and BBH improve slightly, while math and code-oriented scores decrease slightly.

The augmented Llama-2-7B-Chat model shows lower attack success rates than the initial model across all evaluated inference-stage exploits. The gains are especially strong for prefilling attacks and the decoding parameters exploit, while GCG attack rates also decrease on both HEx-PHI and AdvBench. These results support the view that deepening safety alignment can improve robustness to multiple attack types. Prefilling attack success becomes very low across tested token budgets after augmentation The decoding parameters exploit is strongly reduced on the augmented model

Constrained fine-tuning with stronger penalties on initial tokens substantially suppresses attack success rates for Llama-2-7B-Chat and Gemma-1.1-7B-IT across harmful example, identity shifting, and backdoor poisoning settings. Standard fine-tuning causes much higher attack success rates, especially when a backdoor trigger is present, while constrained fine-tuning remains closer to the initial safe model. The constrained objective still improves benign task utility to a level comparable with standard fine-tuning. Strong constraints on early tokens keep attack success rates low after adversarial fine-tuning. Constrained fine-tuning limits trigger activated harmful behavior while preserving low unsafe rates without trigger. Benign fine-tuning utility improves under constrained training and stays close to standard fine-tuning performance.

Experiments on Llama and Gemma models show that prepending short refusal prefixes can strongly steer base models away from harmful outputs, with more complete prefixes yielding safer responses while aligned models show limited additional gains. A model augmented with safety recovery data retains broad capability with only modest utility trade-offs and reduces success across inference-stage exploits such as prefilling, decoding-parameter manipulation, and optimization-based attacks. Finally, constrained fine-tuning that penalizes early-token deviations suppresses attack success in harmful-example, identity-shifting, and backdoor-poisoning settings while preserving benign task performance, indicating that early-token safety constraints can strengthen robustness without major capability loss.