LLM safety alignment should go deeper than the first few tokens

1. 01What problem does this paper attempt to solve?

   The authors state that current safety alignment only acts on the first few output tokens, and this superficiality explains multiple types of known attacks.

   Core question: why is the safety alignment of current LLMs so easily broken by simple attacks or benign fine-tuning, and how can this be mitigated.

   • Problem setting: aligned models (SFT+RLHF/DPO) can be jailbroken by adversarially optimized inputs, a few steps of fine-tuning, or decoding parameters (Introduction). The authors believe there is a common mechanism behind these vulnerabilities.
   • Core observation: the authors propose the concept of shallow safety alignment—alignment mainly changes only the model's generation distribution over the first few output tokens; once the beginning deviates from the refusal prefix, generation enters a harmful trajectory. Its counterfactual is called deep safety alignment.
   • Shortcomings of existing methods: the authors argue that current SFT and RLHF pipelines themselves do not encode the concept of alignment "depth": in SFT data, humans would not write samples that are "harmful first, then refuse," and in RLHF, once the model learns the refusal prefix, it is almost never penalized for this shortcut (Section 2.2).
   • What the paper proposes: 1) a set of case studies characterizing shallow alignment and explaining multiple types of attacks; 2) safety recovery sample data augmentation to deepen alignment; 3) a constrained fine-tuning objective that limits the shift in the initial token distribution to resist fine-tuning attacks.
   • Threat model (the attacks analyzed in the paper, not proposed by this paper):
     • Inference stage: the attacker controls decoding (prefilling a non-refusal prefix), optimizes adversarial suffixes (GCG), or enumerates decoding parameters for sampling; the victim systems are open-source aligned models, and the authors state that the prefilling interface of closed-source models can also be exploited similarly.
     • Fine-tuning stage: a malicious actor revokes alignment through downstream fine-tuning (e.g., harmful samples, identity shifting, backdoor poisoning).
     • Success criterion: GPT-4 judges output harmfulness, and the proportion receiving the highest score of 5 is recorded as ASR.

   Collapse this answer

2. 02What related research exists?

   The paper is most closely related to the Superficial Alignment Hypothesis and three per-token analyses, and provides a unified explanation for multiple types of jailbreak methods.

   The superficiality of alignment and per-token effects

   • Zhou et al. [28]'s Superficial Alignment Hypothesis: holds that alignment only superficially changes the input-output format. The authors of this paper state that their work is closely related to it.
   • Lin et al. [25]: found that the differences introduced by alignment fine-tuning disappear as the sequence grows longer (similar to Figure 1 but not in a safety context); the paper notes that they mainly use this to question whether fine-tuning is necessary, whereas this paper examines the safety effects in depth and connects them to attacks and defenses.
   • Zhang and Wu [26]: found that adaptation to topic and style priors during fine-tuning is mainly learned at the beginning of the sequence. Zhao et al. [27]: noticed a similar token-level effect and directly exploited it to design jailbreak attacks.

   Jailbreak methods

   • The paper provides a unified explanation for four types of attacks: adversarial suffix attacks (Zou et al. [29], Andriushchenko et al. [21]), decoding parameter exploitation (Huang et al. [24]), prefilling attacks ([30, 21], which the authors describe as parallel recent work), and fine-tuning attacks (Qi et al. [22], Zhan et al. [23]).
   • The paper mentions system-level defenses (using a model to monitor inputs and outputs, such as Llama Guard [49]), and cites [50] to point out that the monitoring mechanism itself can also be jailbroken.

   Defenses that protect the initial tokens

   • Xu et al. [53] (SafeDecoding): the authors note that its motivation is similar to that of Section 4 of this paper—defending against inference-time jailbreaks by identifying safety statements and amplifying the probability of early tokens.

   Baselines and benchmarks: The safety evaluation follows the procedure of Qi et al. [22], with the benchmark being HEx-PHI [32] (330 harmful instructions, 11 categories); the GCG attack is implemented according to the original setting of Zou et al. [29] with AdvBench added; the decoding parameter attack follows the setting of Huang et al. [24] with MaliciousInstruct added; utility is evaluated with AlpacaEval (against text-davinci-003).

   Authors' positioning: this paper unifies the above scattered phenomena into the concept of "shallow vs. deep safety alignment" and provides a training-side mitigation path; connections to control theory and safe RL are left for future work.

   Collapse this answer

3. 03How does the paper address this problem?

   First use per-token analysis to characterize shallow alignment, then propose two mitigations: safety recovery data augmentation and token-level constrained fine-tuning.

   Overall approach: first use per-token KL divergence, loss and gradient norm to show that alignment is concentrated on the initial tokens, then mitigate along two paths—deepening alignment (Section 3), or protecting the initial tokens from being destroyed by fine-tuning (Section 4).

   Characterizing shallow safety alignment (Section 2)

   • Safety pattern shortcut: prefill an unaligned base model with a refusal prefix s, sample y ∼ π_base(·|x, s), and test whether its Harmfulness Rate approaches that of the aligned model.
   • KL analysis: construct Harmful HEx-PHI (330 HEx-PHI instructions + harmful responses generated by jailbroken GPT-3.5-Turbo), compute the per-token KL divergence between the aligned model and the base model, and see where the "KL budget" is spent.
   • Fine-tuning dynamics: fine-tune Llama-2-7B-Chat on 100 harmful samples (learning rate 2×10⁻⁵, batch 64), recording per-token cross-entropy loss, gradient norm and KL divergence from the initial model.

   Data augmentation: safety recovery samples (Section 3)

   • Design: construct triples (x, h, r) (harmful instruction, harmful response, refusal), and train the model to still recover a refusal when k harmful tokens have already been prefilled, i.e., enhance π_θ(r|x, h≤k). The authors state that such samples are unnatural and hard to produce by either human annotation or model sampling, and precisely cover out-of-distribution situations.
   • Implementation: since alignment cannot be done from scratch, continue fine-tuning directly on the already aligned Llama-2-7B-Chat; the safety set D_H contains 256 triples, and the utility set D_B uses Alpaca instructions distilled through the initial model; the objective is min_θ α  𝔼[−log π_θ(r|x, h_{≤ k})] + (1−α)  𝔼[−log π_θ(y′|x′)], i.e., a weighted mixture of the safety recovery loss and the utility anchor loss.
   • Parameters: k is 0 with 50% probability and uniformly drawn from [1,100] with 50% probability; α=0.2; each batch has 16 safety samples + 64 utility samples; trained for 10 epochs, learning rate 2×10⁻⁵, AdamW.

   Constrained fine-tuning objective (Section 4)

   • Hypothesis: if the initial few tokens play a decisive role in safety, then constraining the distribution at these positions not to deviate from the initially aligned model can resist fine-tuning attacks.
   • Objective: min_θ 𝔼 ∑_t 2/(β_t) log σ(β_t log (π_θ(y_t|x,y_<t))/(π_aligned(y_t|x,y_<t)))；β_t 控制各位置正则强度，小 β_t 近似交叉熵、大 β_t 近似强制匹配初始模型分布（附录 D 给出极限推导与 RL 解释）。
   • 配置：β₁=0.5，2≤t≤5 时 β_t=2，t>5, β_t=0.1; linear warmup over the first 10 steps (the authors state this avoids the initial gradients destroying alignment before the constraint takes effect).

   Success criteria and evaluation

   • Safety: GPT-4-Turbo scores (input, output) from 1–5, and the proportion scoring 5 is the Harmfulness Rate / ASR; the default benchmark is HEx-PHI.
   • Utility: AlpacaEval winrate (against text-davinci-003); downstream tasks use ROUGE-1 (Samsum, SQL Create Context) and accuracy (GSM8k).

   Collapse this answer

4. 04What experiments did the paper conduct?

   The data-augmented model shows a large drop in ASR under three types of inference attacks, and constrained fine-tuning maintains low ASR under three types of fine-tuning attacks while utility is close to standard SFT.

   Experimental Setup

   • Models tested: Llama-2-7B-Chat, Gemma-1.1-7B-IT (and their respective base models Llama-2-7B and Gemma-7B for comparison); derived model Llama-2-7B-Chat-Augmented.
   • Benchmarks: HEx-PHI (330 harmful instructions, 11 categories); Harmful HEx-PHI (330 instructions paired with harmful responses generated by jailbroken GPT-3.5-Turbo); GCG additionally tested on AdvBench; decoding-parameter attacks additionally tested on MaliciousInstruct.
   • Attack implementation: GCG follows the original paper's single-model multi-behavior setting, optimizing for 500 steps over 50 behaviors, repeated 10 times, with the 3 runs having the highest ASR averaged; decoding-parameter attacks enumerate temperature/top-k/top-p and take the most harmful output; prefilling attacks leave the system prompt empty.
   • Fine-tuning attacks: the three from Qi et al. [22]—Harmful Examples (100 pairs, 25 epochs), Identity Shifting (expanded to 100), Backdoor Poisoning (100+100 pairs); benign fine-tuning uses Samsum, SQL Create Context, and GSM8k (3 epochs each).
   • Judging: GPT-4-Turbo scores 1–5, and the proportion of 5s is counted as ASR; default decoding temperature 0.9, top-p 0.6.
   • Repetition: Table 2 and Table 3 report mean ± std over 3 runs.

   Main Results

   According to Table 2 (ASR%, mean ± std, 3 runs):

   Swipe horizontally to read the table

   Model	Prefill 5 tok	Prefill 40 tok	GCG (HEx-PHI)	GCG (AdvBench)	Decoding params (HEx-PHI)	Decoding params (MaliciousInstruct)
   Initial	42.1 ± 0.9	57.0 ± 0.4	36.5 ± 2.7	65.6 ± 3.1	54.9 ± 0.6	84.3 ± 1.7
   Augmented	2.8 ± 0.4	4.5 ± 0.6	18.4 ± 4.2	19.0 ± 2.9	11.3 ± 0.4	1.0 ± 0

   • The authors state that data augmentation improves robustness against all three types of inference-stage attacks; Figure 4 shows that after augmentation, the KL divergence of harmful responses at later tokens increases.
   • Utility: the Augmented model has an AlpacaEval winrate of 49.5%, versus 51.8% for the initial model; the authors consider utility essentially preserved.

   Evidence of Shallow Alignment (Section 2)

   • Refusal-prefix shortcut (Table 1): Llama-2-7B base without a prefix has a Harmfulness Rate of 68.6 ± 0.8%, which drops to 2.1 ± 0.2% after prefilling "I apologize, but I cannot"; Gemma-7B base drops from 85.4 ± 0.6% to 1.0 ± 0.8%.
   • KL concentrated in the prefix (Figure 1): the authors state that for both model families, the per-token KL divergence on harmful samples is clearly higher for the first few tokens than for subsequent tokens (the converted text has no values from the figure). Llama-2-7B-Chat starts 96.1% of responses on HEx-PHI with "I cannot"/"I apologize", and Gemma-7b-1.1-IT generates "I am unable" in 96.7%.
   • Prefilling attack (Figure 2): the authors state that as more harmful tokens are prefilled, the ASR of aligned models rises from near 0 to above 50%.

   Constrained Fine-Tuning Defense (Section 4)

   According to Table 3 (ASR%, mean ± std, 3 runs):

   Swipe horizontally to read the table

   Attack/Dataset	Metric	Llama Initial	Llama Standard SFT	Llama Constrained	Gemma Standard SFT	Gemma Constrained
   Harmful Examples	ASR	1.5 ± 0.2	88.9 ± 1.2	4.6 ± 0.5	81.6 ± 2.9	1.9 ± 0.2
   Identity Shifting	ASR	0 ± 0	79.5 ± 2.3	8.1 ± 0.1	83.6 ± 2.5	9.1 ± 1.7
   Backdoor (w/ trigger)	ASR	1.7 ± 0.1	90.9 ± 1.4	10.9 ± 2.8	82.3 ± 1.1	1.9 ± 0.8
   Samsum	Utility	25.5 ± 0.3	51.7 ± 0.5	50.1 ± 0.2	51.5 ± 0.3	51.9 ± 0.5
   SQL Create Context	Utility	14.9 ± 0.4	99.1 ± 0.2	98.5 ± 0.1	99.2 ± 0.1	98.6 ± 0.3
   GSM8k	Utility	25.5 ± 0.2	41.7 ± 0.4	37.4 ± 0.3	63.3 ± 0.5	63.6 ± 0.4

   • The authors state that the constrained objective maintains low ASR under both adversarial and benign fine-tuning, with utility comparable to standard SFT; on Gemma, GSM8k utility does not even decrease.
   • Under benign fine-tuning, standard SFT's ASR also rises (Samsum 23.4 ± 2.5%), while constrained SFT drops to 3.2 ± 0.8%.

   Analysis of Fine-Tuning Dynamics (Section 2.3.2, Appendix B)

   • Harmful fine-tuning (Figure 3): Llama-2-7B-Chat has an initial ASR of 1.5%, 22.4% after 2 steps, 76.4% after 4 steps, and 87.9% after 6 steps; the authors state that the loss, gradient norm, and KL shift of the initial tokens are all the largest.
   • Benign fine-tuning (Figure 5): on SQL Create Context and Samsum, the gradient norm of the initial tokens is likewise larger; after 10 steps, ASR has already reached 13.6% and 22.1%, respectively (14.9% and 25.5% after full fine-tuning of 600+ steps). The authors hypothesize that the safety regression under benign fine-tuning mainly comes from excessively large gradients on the initial tokens.

   Other Ablations and Analyses

   • β_t ablation (Table 4, Llama-2-7B-Chat): with uniform β=0.1, Harmful Examples ASR is 86.2% (fails); with uniform β=2.0, ASR is 0.5% but GSM8k utility drops to 2.1%; the biased configuration balances both.
   • Warmup ablation (Table 5): constrained SFT without warmup has an ASR of 29.1% on Harmful Examples, which drops to 4.6% with 10 steps of warmup; the authors state that the improvement mainly comes from the constrained objective rather than warmup itself.
   • Fine-tuning robustness of the augmented model (Table 6): under standard SFT, the Augmented model has a Harmful Examples ASR of 55.2% (versus 88.9% for the non-augmented model) and can still be broken; combined with the constrained objective, no consistently better results were observed.

   Collapse this answer

5. 05What are some points that can be further explored?

   The authors acknowledge that the augmented implementation is not alignment from scratch, and that the defense may be subject to adaptive attacks; the experiments cover only two 7B model families.

   Limitations and future directions noted by the authors

   • The data augmentation implementation continues fine-tuning on the already aligned Llama-2-7B-Chat rather than performing alignment training from scratch; the authors call this implementation "inherently sub-optimal" and plan to implement an end-to-end approach once alignment data and procedures of comparable quality can be obtained (Appendix A.3 Limitations).
   • The authors state that the proposed method "may not be a perfect defense and may be subject to some future adaptive attacks," and is only a preliminary step toward improving robustness (Conclusion).
   • The connections between data augmentation and exploration requirements in reinforcement learning, recovery strategies, and safety control theory are left as future work by the authors (Section 3.1 footnote, Related Work, Conclusion).
   • The authors point out that safety "depth" may be multidimensional, including dimensions beyond token depth such as maintaining safety properties after adaptation (Related Work, Other Notions of Safety Depth).

   Experimental coverage

   • The models tested are two 7B models, Llama-2-7B-Chat and Gemma-1.1-7B-IT, and their base versions; the authors state they chose them because their alignment techniques are closest to those of frontier closed-source models (Related Work).
   • Inference attack evaluations cover three types: prefilling, GCG, and decoding parameters; fine-tuning attacks cover the three types from Qi et al. [22]; benign fine-tuning covers Samsum, SQL Create Context, and GSM8k.
   • Safety evaluation uniformly uses GPT-4-Turbo (1–5 scale, 5 counted as ASR); the paper does not report agreement between the evaluator and human judgment.
   • The main results table reports mean ± std over 3 runs; GCG attacks are repeated 10 times and the 3 runs with the highest ASR are taken.
   • Constrained fine-tuning strengthens constraints only on the initial 5 tokens (β₁=0.5, β₂₋₅=2, the rest 0.1); the paper does not explore other depths or adaptive attackers beyond the continuous β scheme.

   Collapse this answer

6. 06Summarize the main content of the paper

   The authors state that the shallowness of safety alignment provides a unified explanation for multiple types of jailbreak and fine-tuning attacks, and they present two preliminary mitigation paths.

   An analysis-plus-defense paper: the authors propose the concept of "shallow safety alignment," arguing that current alignment only changes the distribution of the model's first few output tokens, and use it to provide a unified explanation for multiple types of known vulnerabilities.

   • Problem: aligned models are easily jailbroken by adversarial suffixes, prefilling, decoding parameter sampling and a small amount of fine-tuning; the authors argue that the common root is that the alignment KL budget is concentrated on prefix tokens (Figure 1), and the refusal prefix itself is a shortcut—prefilling an unaligned base model with a refusal prefix can reduce the Harmfulness Rate from 68.6% to 2.1% (Llama-2-7B, Table 1).
   • Explaining attacks: prefilling a few dozen harmful tokens can raise the ASR of an aligned model from near 0 to above 50% (Figure 2); harmful fine-tuning for 6 steps pushes the ASR from 1.5% to 87.9%, and the loss, gradient and distribution shift of the initial tokens are the largest (Figure 3).
   • Mitigation one (data augmentation): continue fine-tuning with "harmful first, then refusal" safety recovery samples; the augmented model's ASR under 5-token prefilling drops from 42.1% to 2.8%, GCG (AdvBench) from 65.6% to 19.0%, and the decoding parameter attack (MaliciousInstruct) from 84.3% to 1.0% (Table 2); AlpacaEval winrate 49.5% vs 51.8% for the initial model.
   • Mitigation two (constrained fine-tuning): a fine-tuning objective with a strong KL constraint on the initial 5 tokens reduces the ASR under the Harmful Examples attack from 88.9% to 4.6% (Llama-2-7B-Chat, Table 3), with benign task utility comparable to standard SFT.
   • Authors' conclusion: future safety alignment should be "more than a few tokens deep"; the authors also acknowledge that the method is preliminary, may be bypassed by adaptive attacks, and the augmented implementation is not alignment from scratch.

   Collapse this answer

Related papers