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. Our code is available at {https://github.com/Glresearch1/PACT}.
Figures & tables
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.
Setting
Attack Successful Rate (%)
Safety Token
Random Token
Baseline (no intervention)
33.5%
33.5%
Boost ( logits←logits+α )
0.5%
35.0%
Ablate ( logits←−109 )
41.5%
34.0%
Table 1. Attack successful rates on harmful prompts under interventions on identified safety tokens vs. random tokens.
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.
Methods
GSM8K
SST2
AGNEWS
Acc ↑
SR ↓
JB ↓
HB ↓
Acc ↑
SR ↓
JB ↓
HB ↓
Acc ↑
SR ↓
JB ↓
HB ↓
Initial
56.56
1.60
3.00
7.00
87.16
1.60
3.00
7.00
83.20
1.60
3.00
7.00
SFT
81.65
71.25
52.50
94.50
95.87
92.65
53.50
95.00
89.90
83.71
52.00
95.00
Constrained SFT
55.95
10.92
30.50
80.00
92.32
15.90
18.50
45.50
89.30
16.93
20.00
45.50
Safe LoRA
73.69
15.65
20.00
44.50
92.64
92.01
53.50
95.00
91.10
94.89
53.50
95.00
AsFT
80.50
9.27
28.00
48.00
95.41
80.19
54.00
93.50
89.10
30.67
44.00
70.50
Table 2. Generalization across 3 different datasets ( P=10.0% ) on Qwen2.5-7B-Instruct .
Methods
Llama3.1-8B-Instruct
Llama3.2-1B-Instruct
Gemma2-9B-it
Acc ↑
SR ↓
JB ↓
HB ↓
Acc ↑
SR ↓
JB ↓
HB ↓
Acc ↑
SR ↓
JB ↓
HB ↓
Initial
73.39
0.32
3.00
9.50
42.76
7.99
3.50
16.00
73.54
0.00
1.00
0.00
SFT
75.59
81.47
51.00
87.00
46.40
63.26
46.00
77.00
75.74
89.46
53.00
86.00
Constrained SFT
77.56
61.34
38.35
78.50
43.67
13.10
19.00
34.50
78.62
26.52
18.50
42.50
Safe LoRA
72.86
0.32
3.00
9.00
39.88
4.15
5.50
13.00
78.70
67.41
40.50
57.00
AsFT
78.09
97.76
54.50
95.00
43.59
1.60
1.00
7.00
73.62
88.50
47.50
87.00
Table 3. Generalization across different model families/sizes on GSM8K.
Methods
P=0.0%
P=5.0%
P=10.0%
Acc ↑
SR ↓
JB ↓
HB ↓
Acc ↑
SR ↓
JB ↓
HB ↓
Acc ↑
SR ↓
JB ↓
HB ↓
Initial
83.20
1.60
3.00
7.00
83.20
1.60
3.00
7.00
83.20
1.60
3.00
7.00
SFT
90.90
1.28
2.50
8.00
91.20
73.16
51.50
93.00
89.90
83.71
52.00
95.00
Constrained SFT
90.20
2.24
3.00
7.50
90.00
12.46
18.50
42.00
89.30
16.93
20.00
45.50
Safe LoRA
89.20
1.60
1.50
8.00
90.90
71.57
49.50
93.50
91.10
94.89
53.50
95.00
AsFT
91.30
1.60
5.50
9.00
90.10
23.29
33.50
65.00
89.10
30.67
44.00
70.50
Table 4. Generalization under varying proportions of harmful data in the AGNEWS dataset.
Figure 5. Hyperparameters sensitivity
Constraints
ACC (%) ↑
SR (%) ↓
JB (%) ↓
HB (%) ↓
Full-vocab KL
89.30
36.74
35.00
73.00
+Safety tokens
88.50
25.88
28.00
69.00
+Safety weights
89.20
18.85
23.50
51.50
+Response only ref
88.70
6.71
5.50
10.50
+Position decay
89.10
5.75
9.00
13.50
Table 5. Ablation study on progressively adding safety-token constraints, each row cumulatively adds the listed constraint(s) from the previous rows.
Figure 6. Loss changed during finetuning
Figure 7. Mean safety token logits on first 4 answer position
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Methods
ACC
SR
JB
HB
Initial
75.90
0.32
3.00
9.50
SAFEGRAD
87.20
1.32
2.00
9.00
SafeDelta
77.20
1.64
1.50
9.00
Lisa
79.40
0.96
3.50
13.50
PACT (Ours)
87.80
2.56
1.50
1.00
Appendix
Table 6. Comparison with additional baseline methods under AGNEWS, P=10% , and Llama-3.1-8B-Instruct. The results are evaluated by Llama-Guard-3-8B.
Methods
ACC ( P=20% )
SR ( P=20% )
JB ( P=20% )
HB ( P=20% )
ACC ( P=30% )
SR ( P=30% )
JB ( P=30% )
HB ( P=30% )
Initial
83.20
1.60
3.00
7.00
83.20
1.60
3.00
7.00
Ours
89.30
2.88
6.50
10.50
89.10
3.51
5.50
11.00
Appendix
Table 7. Results under higher harmful data ratios, P=20% and P=30% , using Qwen2.5-7B-Instruct on AGNEWS.
Methods
ACC
SR
JB
HB
Initial
75.90
2.56
11.00
5.00
SAFEGRAD
87.20
1.92
11.50
5.50
SafeDelta
77.20
1.92
13.50
5.50
Lisa
79.40
1.92
12.50
5.50
PACT (Ours)
87.80
2.88
7.00
1.00
Appendix
Table 8. Evaluation with Beaver-Dam-7B under AGNEWS, P=10% , and Llama-3.1-8B-Instruct.
Metric
C-SFT
V-SFT
PACT
Ratio
Train time (s)
1589.9
1071.34
1710.68
1.60 ×
Step time (s)
0.1893
0.1417
0.2267
1.60 ×
Memory (GB)
21.86
18.68
33.95
1.82 ×
Appendix
Table 9. Computational overhead under Llama-3.1-8B-Instruct and AGNEWS with P=10% . C-SFT denotes Constrained-SFT and V-SFT denotes Vanilla SFT.
Methods
ACC
SR
JB
HB
SFT
89.50
96.81
54.50
95.50
Ours
87.60
2.88
2.50
1.50
Appendix
Table 10. Injecting highly adversarial prompts from Persuasive-Jailbreaker-Data into the fine-tuning data under AGNEWS, P=10% , and Llama-3.1-8B-Instruct.
Methods
ACC
SR
JB
HB
Add extra 10% safe data
90.80
15.65
46.00
84.50
Lisa
79.40
0.96
3.50
13.50
Ours
87.80
2.56
1.50
1.00
Appendix
Table 11. Comparison with a simple data-centric baseline under AGNEWS, P=10% , and Llama-3.1-8B-Instruct.
Methods
ACC
SR
JB
HB
SFT-PJD
89.50
96.81
54.50
95.50
PACT-PJD
87.60
2.88
2.50
1.50
SFT-Refusal-PJD
89.50
96.49
57.50
95.00
PACT-Refusal-PJD
87.60
3.83
4.00
5.50
Appendix
Table 12. Results under adversarial prompts with refusal-like prefixes.
Fine-tuning well-aligned large language models (LLMs) on new domains often degrades their safety alignment, even when using benign datasets. Existing safety alignment techniques primarily focus on pretraining, leaving fine-tuned models vulnerable to behavioral shifts. In this work, we introduce safety token regularization (STR), a lightweight method designed to preserve safety properties during fine-tuning. Our approach identifies salient tokens from rejection templates of well-aligned models and constrains their associated logits during training, preventing the loss of critical safety behaviors. Unlike reinforcement learning or preference optimization methods, STR requires minimal additional computation and seamlessly integrates with parameter-efficient fine-tuning techniques such as LoRA. Comprehensive experiments demonstrate that our approach achieves safety performance on par with state-of-the-art methods, while preserving task-specific utility and requiring minimal implementation overhead. Furthermore, we show that safety token regularization enhances training stability and overall performance beyond safety considerations alone. This work offers a practical and readily deployable strategy for continual safety alignment in fine-tuned LLMs.
Thong Bach, Truyen Tran
Applied Artificial Intelligence Initiative (A2I2), Deakin University
Fine-tuning large language models (LLMs) on custom datasets has become a standard approach for adapting these models to specific domains and applications. However, recent studies have shown that such fine-tuning can lead to significant degradation in the model's safety. Existing defense methods operate at the sample level and often suffer from an unsatisfactory trade-off between safety and utility. To address this limitation, we perform a systematic token-level diagnosis of safety degradation during fine-tuning. Based on this, we propose token-level data selection for safe LLM fine-tuning (TOSS), a novel framework that quantifies the safety risk of each token by measuring the loss difference between a safety-degraded model and a utility-oriented model. This token-level granularity enables accurate identification and removal of unsafe tokens, thereby preserving valuable task-specific information. In addition, we introduce a progressive refinement strategy, TOSS-Pro, which iteratively enhances the safety-degraded model's ability to identify unsafe tokens. Extensive experiments demonstrate that our approach robustly safeguards LLMs during fine-tuning while achieving superior downstream task performance, significantly outperforming existing sample-level defense methods. Our code is available at https://github.com/Polly-LYP/TOSS.
Yanping Li, Zhening Liu, Zijian Li +2
The Hong Kong University of Science and Technology · School of Data Science, Lingnan University
Safety alignment in large language models (LLMs) is commonly implemented as a single static policy embedded in model parameters. However, real-world deployments often require context-dependent safety rules that vary across users, regions, and applications. Existing approaches struggle to provide such conditional control: parameter-level alignment entangles safety behaviors with general capabilities, while prompt-based methods rely on natural language instructions that provide weak enforcement. We propose MOSAIC, a modular framework that enables compositional safety alignment through learnable control tokens optimized over a frozen backbone model. Each token represents a safety constraint and can be flexibly activated and composed at inference time. To train compositional tokens efficiently, we introduce order-based task sampling and a distribution-level alignment objective that mitigates over-refusal. Experiments show that MOSAIC achieves strong defense performance with substantially lower over-refusal while preserving model utility.
Jingyu Peng, Hongyu Chen, Jiancheng Dong +5
University of Science and Technology of China · City University of Hong Kong · Baidu Inc +1