Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Few Tokens, Big Leverage: Preserving Safety Alignment by Constraining Safety Tokens during Fine-tuning

Guoli Wang, Haonan Shi, Tu Ouyang, An Wang

TL;DR

This work proposes a fine-tuning framework called Preserving Safety Alignment via Constrained Tokens (PACT), which stabilizes the model's confidence on safety tokens and prevents alignment drift without imposing global restrictions that typically trade off with model utility.

Abstract

Large language models (LLMs) often require fine-tuning (FT) to perform well on downstream tasks, but FT can induce safety-alignment drift even when the training dataset contains only benign data. Prior work shows that introducing a small fraction of harmful data can substantially compromise LLM refusal behavior, causing LLMs to comply with harmful requests. Existing defense methods often rely on model-wide interventions, such as restricting which parameters are updated or injecting additional safety data, which can limit generality and degrade downstream task performance. To address these limitations, we propose a fine-tuning framework called Preserving Safety Alignment via Constrained Tokens (PACT), which stabilizes the model's confidence on safety tokens. Our approach is motivated by the empirical observation that safety-aligned behavior is reflected in the model's token-level output confidence and is often concentrated on a small subset of safety-related tokens. During downstream fine-tuning, we regularize the fine-tuned model to match the aligned reference model's confidence on safety-related tokens at each response step, while leaving non-safety tokens largely unconstrained to allow effective task adaptation. This targeted constraint prevents alignment drift without imposing global restrictions that typically trade off with model utility.

Few Tokens, Big Leverage: Preserving Safety Alignment by Constraining Safety Tokens during Fine-tuning

TL;DR

This work proposes a fine-tuning framework called Preserving Safety Alignment via Constrained Tokens (PACT), which stabilizes the model's confidence on safety tokens and prevents alignment drift without imposing global restrictions that typically trade off with model utility.

Abstract

Large language models (LLMs) often require fine-tuning (FT) to perform well on downstream tasks, but FT can induce safety-alignment drift even when the training dataset contains only benign data. Prior work shows that introducing a small fraction of harmful data can substantially compromise LLM refusal behavior, causing LLMs to comply with harmful requests. Existing defense methods often rely on model-wide interventions, such as restricting which parameters are updated or injecting additional safety data, which can limit generality and degrade downstream task performance. To address these limitations, we propose a fine-tuning framework called Preserving Safety Alignment via Constrained Tokens (PACT), which stabilizes the model's confidence on safety tokens. Our approach is motivated by the empirical observation that safety-aligned behavior is reflected in the model's token-level output confidence and is often concentrated on a small subset of safety-related tokens. During downstream fine-tuning, we regularize the fine-tuned model to match the aligned reference model's confidence on safety-related tokens at each response step, while leaving non-safety tokens largely unconstrained to allow effective task adaptation. This targeted constraint prevents alignment drift without imposing global restrictions that typically trade off with model utility.
Paper Structure (14 sections, 9 equations, 7 figures, 5 tables)

This paper contains 14 sections, 9 equations, 7 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Analyzing safety tokens
  3. Identifying safety tokens
  4. Impact of safety tokens
  5. Evolution of safety tokens confidence during fine-tuning
  6. Methodology
  7. Regularization with weighted safety tokens
  8. Calibration of safety signal
  9. Evaluations
  10. Evaluation setup
  11. Main results
  12. Ablation studies
  13. Related works
  14. Conclusion

Figures (7)

  • Figure 1: Fine-tuning is widely used to improve large language models on downstream tasks, but the presence of harmful data in downstream training sets can induce significant safety-alignment drift, making the fine-tuned model highly vulnerable to harmful queries.
  • Figure 2: Top-10 safety tokens identified by token-level confidence discrepancies between the safety-aligned and base models on harmful questions.
  • Figure 3: Evolution of safety alignment token confidence during harmful fine-tuning. Safety score = (1-ASR)%
  • Figure 4: The workflow of our proposed token-level safety-preserved fine-tuning method.
  • Figure 5: Hyperparameters sensitivity
  • ...and 2 more figures