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

Guoli Wang; Haonan Shi; Tu Ouyang; An Wang

cs.CLMar 8, 2026

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

Authors: Guoli Wang, Haonan Shi, Tu Ouyang, An Wang

Organizations: Case Western Reserve University, Cleveland, OH, USA

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

Panel of Figure 2 — 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 (%)
Setting	Safety Token	Random Token
Baseline (no intervention)	33.5%	33.5%
Boost ( $logits\leftarrow logits+\alpha$ )	0.5%	35.0%
Ablate ( $logits\leftarrow-10^{9}$ )	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
Methods	Acc $\uparrow$	SR $\downarrow$	JB $\downarrow$	HB $\downarrow$	Acc $\uparrow$	SR $\downarrow$	JB $\downarrow$	HB $\downarrow$	Acc $\uparrow$	SR $\downarrow$	JB $\downarrow$	HB $\downarrow$
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
Methods	Acc $\uparrow$	SR $\downarrow$	JB $\downarrow$	HB $\downarrow$	Acc $\uparrow$	SR $\downarrow$	JB $\downarrow$	HB $\downarrow$	Acc $\uparrow$	SR $\downarrow$	JB $\downarrow$	HB $\downarrow$
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%
Methods	Acc $\uparrow$	SR $\downarrow$	JB $\downarrow$	HB $\downarrow$	Acc $\uparrow$	SR $\downarrow$	JB $\downarrow$	HB $\downarrow$	Acc $\uparrow$	SR $\downarrow$	JB $\downarrow$	HB $\downarrow$
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.

Panel of Figure 5 — Figure 5. Hyperparameters sensitivity

Constraints	ACC (%) $\uparrow$	SR (%) $\downarrow$	JB (%) $\downarrow$	HB (%) $\downarrow$
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 $\times$
Step time (s)	0.1893	0.1417	0.2267	1.60 $\times$
Memory (GB)	21.86	18.68	33.95	1.82 $\times$

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.

Explore similar work

Guardrails in Logit Space: Safety Token Regularization for LLM Alignment
Apr 19, 2026Thong Bach, Truyen TranLLM Fine-Tuning LLM Safety Alignment
Applied Artificial Intelligence Initiative (A2I2), Deakin University
Token-level Data Selection for Safe LLM Fine-tuning
Mar 1, 2026Yanping Li, Zhening Liu, Zijian Li +2LLM Fine-Tuning LLM Safety
The Hong Kong University of Science and Technology · School of Data Science, Lingnan University
MOSAIC: Composable Safety Alignment with Modular Control Tokens
Mar 17, 2026Jingyu Peng, Hongyu Chen, Jiancheng Dong +5LLM Safety Alignment Inference-Time Language Model Alignment
University of Science and Technology of China · City University of Hong Kong · Baidu Inc +1

Apr 19, 2026J/K move · Enter open · S save

Guardrails in Logit Space: Safety Token Regularization for LLM Alignment

Thong Bach, Truyen Tran

Applied Artificial Intelligence Initiative (A2I2), Deakin University

LLM Fine-Tuning LLM Safety Alignment

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Guardrails in Logit Space: Safety Token Regularization for LLM Alignment

Token-level Data Selection for Safe LLM Fine-tuning

MOSAIC: Composable Safety Alignment with Modular Control Tokens