[Paper Review] Safety Alignment Should Be Made More Than Just a Few Tokens Deep

Xiangyu Qi, Ashwinee Panda+6 more|arXiv (Cornell University)|Jun 10, 2024

Risk and Safety AnalysisDecision Sciences3 citations

TL;DR

This paper identifies 'shallow safety alignment'—where LLM safety relies primarily on the first few output tokens—as a core vulnerability enabling jailbreaks via adversarial suffixes, decoding attacks, and fine-tuning. It proposes deepening alignment through data augmentation and a constrained fine-tuning objective, showing improved robustness against these exploits by enforcing persistent safety beyond initial tokens.

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

Motivation & Objective

• To investigate why current LLM safety alignment is vulnerable to simple attacks despite extensive fine-tuning and RLHF.
• To identify 'shallow safety alignment'—where alignment effects are confined to the first few output tokens—as a root cause of these vulnerabilities.
• To demonstrate that many known attacks (e.g., adversarial suffixes, prefilling, decoding parameter, and fine-tuning attacks) exploit this shallow alignment by manipulating initial tokens.
• To propose methods that deepen safety alignment beyond the first few tokens to improve robustness.
• To develop a constrained fine-tuning objective that limits changes in initial token probabilities, enhancing resistance to fine-tuning-based jailbreaks.

Proposed method

• Introduces a data augmentation strategy that includes harmful initial responses followed by safety refusals, training models to recover from unsafe starts.
• Proposes a regularized fine-tuning objective based on preference modeling, using the aligned model’s policy as a reference to constrain updates on initial token probabilities.
• Derives a surrogate objective using the sigmoid of a scaled log-ratio of policy probabilities, with a temperature-like normalization via βt to stabilize gradients across sequence positions.
• Employs the Bradley-Terry model framework to define a preference function T that favors tokens with higher expected reward relative to the aligned policy’s average.
• Uses a lower-bound surrogate loss Lθ to enable end-to-end optimization without explicit value function learning, ensuring the objective remains differentiable and scalable.
• Applies the method to fine-tuning scenarios, explicitly regularizing the model to avoid large shifts in the probability distribution of the first few tokens during adaptation.

Experimental results

Research questions

• RQ1 Why are current LLM safety alignment methods vulnerable to simple attacks like adversarial suffixes and prefilling?
• RQ2 To what extent is the safety behavior of aligned models determined by only the first few output tokens?
• RQ3 Can deepening safety alignment—extending its influence beyond the first few tokens—improve robustness against known exploits?
• RQ4 How can a constrained fine-tuning objective mitigate the risk of jailbreaking via fine-tuning attacks?
• RQ5 What role does the initial token distribution play in enabling or preventing catastrophic alignment failure?

Key findings

• Current safety alignment is predominantly shallow, with the primary difference between aligned and unaligned models lying in the generation of the first few output tokens.
• Simple prefilling of an unaligned model with a safe prefix like 'I cannot fulfill' can make it as safe as a fully aligned model, demonstrating the fragility of shallow alignment.
• Adversarial suffix attacks, decoding parameter exploits, and prefilling attacks all succeed by manipulating the initial token distribution, confirming that alignment is not robust beyond the first few tokens.
• Fine-tuning attacks induce the most significant changes in the first few tokens of harmful responses, explaining why only a few fine-tuning steps can lead to jailbreaking.
• Data augmentation that includes harmful-to-safe transitions successfully deepens alignment, increasing divergence between aligned and unaligned models at deeper token positions and improving robustness.
• The proposed constrained fine-tuning objective, which regularizes initial token probabilities, significantly reduces the risk of jailbreaking via fine-tuning attacks, providing a practical defense mechanism.