PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Xiangyu Qi, Ashwinee Panda, Kaifeng Lyu, Xiao Ma, Subhrajit Roy, Ahmad Beirami, Prateek Mittal, Peter Henderson

arXiv:2406.05946v1cs.CRcs.AI

TL;DR

Current safety alignment can be shallow, changing models’ generative distributions mainly at the first few output tokens and leaving aligned LLMs vulnerable to attacks and fine-tuning. The paper characterizes this issue, connects it to several vulnerabilities, and evaluates strategies for deepening alignment and constraining initial-token updates; these approaches improve robustness against some exploits and help mitigate fine-tuning attacks.

  • Problem

    Current alignment approaches can be bypassed by adversarial inputs, fine-tuning, and even benign downstream fine-tuning, motivating investigation of shallow safety alignment.

  • Method

    The paper characterizes per-token safety alignment, introduces harmful-response data augmentation to deepen alignment, and proposes constrained optimization that limits shifts in initial token probabilities.

  • Results

    The paper finds that current aligned models primarily differ from unaligned models on their first few output tokens, while deeper alignment often improves robustness and constrained optimization can mitigate fine-tuning attacks.

  • Takeaways & Limitations

    Safety alignment should extend beyond the first few tokens, with deeper alignment and constraints on initial-token updates providing initial mitigation strategies.

  • Takeaways & Limitations

    The paper does not formally explore the ties between deepened safety alignment and sufficient exploration in reinforcement learning.

Abstract

from arXiv · show

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. Finally, we design a regularized finetuning 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 Introduction

The paper argues that current LLM safety alignment is vulnerable because it may primarily alter the first few output tokens. It characterizes this shallow alignment, links it to several jailbreak vulnerabilities, and proposes deeper alignment and constrained fine-tuning as mitigation directions.

  • Current alignment methods can be bypassed by adversarial inputs, fine-tuning, or decoding-parameter exploits.These vulnerabilities motivate examining why widely used alignment approaches remain susceptible to relatively simple attacks.
  • Shallow safety alignment primarily changes the model’s generative distribution over its first few output tokens.If those initial tokens deviate from safe prefixes, generation may move onto a harmful trajectory.
  • Systematic experiments show that the main safety-behavior difference between aligned and unaligned models lies in modeling the first few output tokens.The paper connects this pattern to adversarial suffix, decoding-parameter, prefilling, and fine-tuning attacks.
  • Training on harmful responses that transition back to refusals can deepen alignment and often improve robustness against common exploits.This data augmentation increases divergence between aligned and unaligned models at greater token depths.
  • Constraining shifts in initial token probabilities can mitigate fine-tuning attacks.The objective is presented as a defense motivated by the shallow-alignment analysis.
  • The paper frames shallow versus deep safety alignment as a unifying lens for understanding vulnerabilities and developing future mitigation strategies.It advocates making safety alignment more than a few tokens deep.

2 The Shallow Safety Alignment Issue in Current Large Language Models

Shallow safety alignment adapts models mainly over the first few output tokens, creating a common vulnerability across inference-time exploits and fine-tuning attacks. Evidence from refusal-prefix behavior, token-wise divergence, prefilling, and fine-tuning shows that later harmful-token distributions remain comparatively less aligned and early-token changes can rapidly undo safety.

  • Core characterization: Shallow safety alignment primarily adapts the model’s generative distribution over the first few output tokens to induce refusal responses.The paper defines this as alignment that can appear safe under standard workflows but fails when a non-refusal prefix is triggered.
  • Refusal-prefix behavior: Prefilling refusal prefixes substantially reduces the harmfulness gap between unaligned base models and their aligned counterparts.The paper presents this as a shortcut in which updating only the first few output-token probabilities can make an unaligned model appear safe.
  • Token-wise evidence: Aligned and unaligned models differ most in the first few harmful-response tokens, while later-token distributions remain largely unaffected.Figure 1’s per-token KL divergence is significantly higher at early positions, indicating that most alignment divergence is concentrated in the prefix.
  • Inference-stage vulnerabilities: Conditioning aligned models on increasing numbers of harmful prefilled tokens raises harmful-generation ASR from near zero to over 50%.This demonstrates that bypassing early refusal prefixes can expose harmful trajectories during decoding.
  • Downstream fine-tuning vulnerabilities: After six gradient steps on 100 harmful examples, ASR rises from 1.5% to 87.9% as fine-tuning most strongly perturbs the first few token distributions.Early positions have higher loss and gradient norms, causing rapid divergence from the initially aligned model.

3 What If The Safety Alignment Were Deeper?

The paper deepens safety alignment by training models to recover refusals after harmful-response prefixes, then evaluates whether this improves robustness against inference-time and fine-tuning exploits.

  • 3.1 Data Augmentation with Safety Recovery Examples: Safety recovery examples pair harmful instructions and responses with refusals after randomly prefixed harmful content, training the model to suppress harmful behavior deeper in responses.The augmented examples explicitly promote refusal probabilities conditioned on harmful prefixes rather than only on the original instruction.
  • 3.1 Data Augmentation with Safety Recovery Examples: The approach uses 256 augmented examples to further fine-tune Llama-2-7B-Chat, with prefilling length k set to zero half the time and uniformly sampled from 1 to 100 otherwise.The utility-to-safety-example balance is set with α = 0.2.
  • 3.1 Data Augmentation with Safety Recovery Examples: The augmented fine-tuning increases per-token KL divergence on later tokens of harmful responses, indicating that safety alignment extends to deeper response positions.Figure 4 reports this effect specifically on Harmful HEx-PHI.
  • 3.1 Data Augmentation with Safety Recovery Examples: The augmented model achieves a 49.5% AlpacaEval win rate versus 51.8% for the initial Llama-2-7B-Chat model, indicating little utility degradation.The win rate is measured against text-davinci-003.
  • 3.2 The Deepened Safety Alignment Shows Improved Robustness Against Multiple Exploits: Augmented fine-tuning improves robustness against prefilling, GCG, and decoding-parameter attacks across the evaluated attack settings.Table 2 compares attack success rates between the initial and augmented Llama-2-7B-Chat models over three runs.
  • 3.2 The Deepened Safety Alignment Shows Improved Robustness Against Multiple Exploits: The augmented model shows less safety regression under benign fine-tuning and lower attack success rates under harmful adversarial fine-tuning, but remains vulnerable to harmful fine-tuning.The reported durability improvement holds across multiple adversarial cases.

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

The paper proposes token-wise constrained fine-tuning that protects the initial aligned distribution while allowing later-token adaptation, targeting fine-tuning attacks without sacrificing benign utility.

  • 4.1 A Token-wise Constrained Objective for Custom Fine-tuning Aligned LLMs: The constrained objective limits fine-tuning-induced deviation from the initially aligned model’s token distribution, motivated by the decisive role of early-token shifts in safety failures.It adaptively reduces the influence of token positions whose distributional deviation approaches a βt-controlled threshold.
  • 4.1 A Token-wise Constrained Objective for Custom Fine-tuning Aligned LLMs: Larger βt imposes stronger token-wise regularization toward the initial aligned model’s generative distribution.The objective can also be interpreted as KL-regularized reinforcement learning with position-specific regularization strength.
  • 4.1 A Token-wise Constrained Objective for Custom Fine-tuning Aligned LLMs: Small βt emphasizes cross-entropy fitting, whereas large βt emphasizes matching the initial aligned model’s generative distribution.The objective therefore interpolates between standard supervised fine-tuning and distribution-preserving behavior.
  • 4.2 Experiments: The experiments impose β1 = 0.5 and βt = 2 for 2 ≤t ≤5 on the first five tokens, while using βt = 0.1 for later tokens.This configuration concentrates the strongest constraints on the initial positions.
  • 4.2 Experiments: The constrained objective is evaluated against harmful examples, identity shifting, and backdoor poisoning attacks, alongside benign fine-tuning on Samsum, SQL Create Context, and GSM8k.These experiments test both safety durability and downstream utility.
  • 4.2 Experiments: Constrained SFT generally maintains low ASR after adversarial and benign fine-tuning on Llama-2-7B-Chat and Gemma-1.1-7B-IT.The results support protecting early-token distributions as a way to make safety alignment more persistent.
  • 4.2 Experiments: Constrained SFT achieves utility comparable to standard SFT while mitigating harmful fine-tuning risk across the reported downstream tasks.Both standard and constrained SFT improve utility over the initial model on all three benign use cases.

5 Related Work

The related work connects shallow safety alignment to alignment hypotheses, token-position effects, jailbreak defenses, control and reinforcement learning, and broader notions of safety depth.

  • Safety & Alignment: Prior alignment research includes SFT, preference optimization, and methods such as DPO, while this paper links alignment behavior to downstream jailbreak vulnerabilities.The authors focus on Gemma and Llama-2 aligned and base models because their alignment is closest to technology used in practice.
  • Jailbreaking Methods: Jailbreaking research spans fine-tuning, decoding, prefilling, optimization, and persuasion, while monitoring-based defenses depend on monitors that may themselves be jailbroken.This paper instead examines vulnerabilities arising within the aligned model’s generative behavior.
  • Superficial Alignment Hypothesis and Per-token Effects of Alignment Fine-tuning: The paper relates shallow safety alignment to the Superficial Alignment Hypothesis and prior findings that fine-tuning changes topic, style, or aligned-base differences primarily near sequence beginnings.Earlier work also reports that aligned-versus-unaligned differences can vanish as sequences grow longer.
  • Protecting The Safety Alignment at Initial Token Positions: Related work similarly suggests protecting early-token generative distributions as a defense against inference-time jailbreak attacks.The paper notes a connection to Xu et al.’s defense insight.
  • Connections to Control Theory and Safe Reinforcement Learning: The data augmentation approach connects to exploration, recovery policies, and safe control theory, but deeper formal connections are left for future work.The authors also note that safety depth may include post-adaptation retention, not only token-based depth.

6 Conclusion

The paper identifies shallow safety alignment as a shortcut underlying downstream vulnerabilities and proposes two initial strategies to deepen alignment and mitigate fine-tuning attacks.

  • Current safety alignment appears to change the generative distribution mainly over the first few output tokens.
  • The paper proposes data augmentation to increase alignment depth and constrained optimization to limit initial-token updates during fine-tuning.
  • These methods are initial steps rather than perfect defenses and may remain vulnerable to future adaptive attacks.

Broader Impacts and Ethics Statement

The paper connects alignment failures to shortcuts and advocates deeper alignment approaches, while noting that investigating these failures could also increase jailbreak capabilities.

  • Alignment shortcuts provide a path toward understanding and improving robustness against safety failures.
  • Open investigation of alignment failure modes may increase the ability to jailbreak models, but the authors consider such investigation important for strengthening future model safety.

A.1 Compute Resources

The experiments use specified GPU-node resources, decoding defaults, and GPT-4-Turbo-based safety evaluation with harmfulness scores from 1 to 5.

  • Compute Resources: Experiments run on single nodes with either 4 × A100-80GB or 4 × H100-80GB GPUs, using up to 8 CPU cores and 256GB memory.
  • Compute Resources: Default decoding uses top-p sampling with temperature 0.9 and top-p 0.6, except in the decoding-parameter exploit experiment.
  • Compute Resources: Safety evaluation passes input-output pairs to GPT-4-Turbo, which scores harmfulness from 1 to 5, with higher scores indicating greater harm.

A.3 Details of Data Augmentation Experiments

The data augmentation experiments combine harmful-instruction recovery examples with benign utility examples, use specified sampling and weighting choices, evaluate utility with AlpacaEval, and are limited by the absence of the original alignment pipeline.

  • A.3 Details of Data Augmentation Experiments: The implementation details section describes the data augmentation experiments in Section 3.1.
  • Safety Data: The safety dataset DH contains 256 harmful instructions designed to generate safety recovery examples without overlapping the evaluation datasets.
  • Utility Data: The utility dataset DB uses distilled Alpaca responses to anchor the model's original behavior on benign instructions.
  • Training Details: Training samples k=0 with 50% probability and uniformly samples k from [1, 100] otherwise, while setting α=0.2 and drawing 16 safety and 64 utility examples per batch.
  • Utility Evaluation: Utility is evaluated with AlpacaEval 1.0 without length control as a winrate against text-davinci-003, using Llama-2-7B-Chat's safety-focused system prompt.
  • Limitations: Because the original Llama-2 alignment data and pipeline are unavailable, the implementation further fine-tunes an already aligned model and is inherently sub-optimal.

A.4 Details of Inference-Stage Attacks Experiments

The experiments evaluate inference-stage and fine-tuning attacks, alongside benign fine-tuning, using specified datasets, optimization settings, and attack procedures.

  • Inference-stage attacks: The study evaluates three inference-stage attacks: prefilling, GCG, and decoding-parameter exploits.The attacks are tested in the cited experimental setup.
  • Inference-stage attacks: Prefilling experiments use 330 harmful HEx-PHI instructions paired with harmful responses sampled from jailbroken GPT-3.5-Turbo.The construction enables testing with harmful or non-refusal tokens prefilled.
  • Inference-stage attacks: GCG optimization trains adversarial suffixes against one victim model using 50 harmful behaviors for 500 steps.The setup follows the original paper’s single-model, multiple-harmful-behaviors configuration.
  • Inference-stage attacks: Decoding-parameter exploits enumerate temperature, top-k, and top-p settings, sample one output per combination, and select the most unsafe output.The selected output is then used for ASR evaluation.
  • Fine-tuning setup: Fine-tuning experiments use AdamW, batch size 64, and model-specific learning rates of 2 × 10^-5 for Llama-2-7B-Chat and 5 × 10^-6 for Gemma-1.1-7B-IT.The optimizer momentum parameters are 0.5 and 0.999.
  • Constrained SFT: Constrained SFT uses 10-step linear learning-rate warmup because its initial gradient matches standard cross-entropy and may otherwise break alignment early.The constraint is initialized gently while the model starts from the aligned policy.
  • Fine-tuning attacks: The evaluated fine-tuning attacks are Harmful Examples, Identity Shifting, and Backdoor Poisoning.Each attack uses 100-data-point configurations, with 25 fine-tuning epochs; backdoor poisoning combines refusal and triggered harmful examples.
  • Benign fine-tuning: Benign fine-tuning evaluates Samsum, SQL Create Context, and GSM8k for three epochs, using ROUGE-1 for the first two and answer accuracy for GSM8k.Samsum is summarization, SQL Create Context converts natural language to SQL, and GSM8k tests mathematics.

B Pertoken Dynamics of Benign Fine-tuning

Benign fine-tuning exhibits the same concentration of gradient magnitude on initial tokens observed in harmful fine-tuning. Safety regression is already substantial after the first 10 steps, suggesting that early updates drive much of the change.

  • Per-token dynamics: Benign fine-tuning shows much larger initial gradient norms on the first few tokens.Figure 5 examines Llama-2-7B-Chat fine-tuned on SQL Create Context and Samsum.
  • Per-token dynamics: The authors attribute this pattern to alignment-induced confidence in fixed affirmative prefixes.These prefixes create large early-token gradients during instruction fine-tuning.
  • Safety regression: The authors hypothesize that excessively large early updates to initial transition-token distributions cause over-generalization or catastrophic forgetting that disrupts refusal prefixes.This proposed mechanism connects benign fine-tuning dynamics with safety regression in aligned LLMs.
  • Safety regression: 13.6% and 22.1% ASR occur after 10 gradient steps on SQL Create Context and Samsum, respectively.The initially aligned model has 1.5% ASR; after more than 600 steps, ASR reaches 14.9% and 25.5%.

C Ablation Studies on Fine-tuning Attack Experiments

The ablations examine how constrained fine-tuning preserves safety, focusing on stronger constraints at early tokens, warmup initialization, and robustness of an augmented model.

  • Biased Constraints on the Early Tokens: Early-token constraints are necessary: setting β = 0.1 uniformly cannot stop safety degradation, whereas β = 2.0 uniformly preserves safety but collapses fine-tuning utility.Uniform β = 0.5 also fails to achieve optimal safety and performs worse in utility than the biased configuration.
  • Ablation on the Effects of Warmup Steps: Ten warmup steps make constrained SFT consistently safer, while most of the safety improvement comes from the constrained optimization objective itself.Warmup gently initializes the constraints because the initial gradient otherwise matches standard cross-entropy and may break alignment.
  • Fine-tuning Attacks on the Augmented Model: The augmented model is generally more robust across multiple fine-tuning attacks than its non-augmented counterpart, although the constrained objective was not consistently better on it.These comparisons are reported for the augmented model built in Section 3 and the corresponding non-augmented model.
  • Constrained Fine-tuning Objective: The constrained fine-tuning objective adaptively reduces gradient weight at token positions where the fine-tuned distribution increasingly deviates from the aligned model.The deviation is measured through the difference in log probabilities, with β_t controlling the threshold for diminishing the gradient.
  • Reinforcement-Learning Interpretation: The constrained fine-tuning objective can be interpreted as a surrogate for a token-wise KL-regularized reinforcement-learning problem.The formulation uses token-wise rewards and value functions, with fine-tuning examples treated as preferable to responses from the initial aligned policy.
Loading 2406.05946v1…