Fine-tuning large language models (LLMs) on downstream tasks often weakens their previously aligned safety behavior, creating an inherent trade-off between task adaptation and safety preservation. Existing approaches attempt to mitigate this issue by identifying safety-relevant neurons, layers, or update directions in the original parameter space. However, safety-related information is often distributed across many directions in this space, causing downstream updates to overwrite parameters that encode safety behavior. We propose WSR-Tune, a Weight Space Reparameterization-based fine-tuning framework that preserves safety during downstream adaptation. Our approach constructs a safety-conditioned basis and reparameterizes the weight matrices such that safety-relevant information becomes concentrated in a small subset of directions. We then identify safety-critical directions and freeze them in the reparameterized space, while updating the remaining complementary directions for downstream adaptation. By explicitly structuring the parameter space to disentangle safety and task information, our method reduces interference between objectives and enables simultaneous preservation of safety and improvement in downstream performance. Experimental results show that our method achieves a more favorable trade-off between downstream performance and safety retention, demonstrating its effectiveness for reliable LLM fine-tuning.
Figures & tables
Fine-tuning setting
Safety ↓
FT (0% frozen)
53.84
FT (10% frozen)
33.95
FT (30% frozen)
21.44
FT (50% frozen)
18.43
WSR-Tune (10% frozen)
14.67
Table 1: Impact of freezing safety-relevant parameters on safety preservation after downstream fine-tuning. The reported values are averaged over Direct, AutoDAN, PAIR, and PAP attacks using safety-tuned Llama-2-7B.
Figure 1: Overview of WSR-Tune . The weight space is reparameterized into a safety-conditioned basis, in which safety-relevant information is concentrated in a subset of coefficients. Safety-critical coefficients are then identified and frozen, while the remaining coefficients are updated for downstream adaptation. This reduces interference between safety and downstream task information during fine-tuning, improving task performance with minimal safety degradation.
Figure 2: For each layer, singular values are sorted in descending order and normalized by the largest value. Llama-2-7B-Chat model is used.
Llama-2-7B-Chat
Llama-2-13B-Chat
Safety ↓
Downstream ↑
Overall ↑
Safety ↓
Downstream ↑
Overall ↑
Method
JB AVG
ΔS
Accuracy
ΔD
ΔOverall
JB AVG
ΔS
Accuracy
ΔD
ΔOverall
Full Params FT
20.78
0.00
41.17
0.00
0.00
10.04
0.00
45.94
0.00
0.00
SafeInstr
13.55
-7.23
38.44
-2.73
+4.50
5.82
-4.22
46.55
+0.61
+4.83
Resta
10.56
-10.22
36.92
-4.25
+5.97
10.34
+0.30
44.20
-1.74
-2.04
SEAL
12.75
-8.03
38.89
-2.28
+5.75
10.28
+0.24
45.19
-0.75
-0.99
Table 2: Comparison across various chat and instruction-tuned models. Safety is measured by Jailbreak (JB) AVG (averaged over Direct, AutoDAN, PAIR, and PAP; lower is better), and downstream performance is measured by task accuracy (higher is better). Specifically, we use GSM8K for the Llama-2 Chat models and MATH for the Llama-3 Instruct models. We also report changes relative to Full Params FT: ΔS (safety, lower is better) and ΔD (downstream, higher is better). We further define ΔOverall=−ΔS+ΔD to summarize the safety–utility trade-off.
Figure 3: Safety-downstream trade-off among fine-tuning on Llama-2-13B-Chat. The x-axis shows downstream accuracy, and the y-axis shows defense success rate ( 1− ASR), where higher indicates better safety.
Llama-2-7B-Chat
Safety ↓
Downstream ↑
Method
Direct
AutoDAN
PAIR
PAP
AVG
Accuracy
SN-Tune
0.00
17.69
45.58
45.46
27.18
38.99
SN-Tune + WSR-Tune
0.00
11.92
41.92
38.46
23.07
41.02
SafeInstr
0.00
0.38
18.85
34.96
13.55
38.44
SafeInstr + WSR-Tune
0.00
0.00
8.46
29.69
9.53
39.27
SEAL
0.00
4.23
16.92
29.85
12.75
38.89
Table 3: Compatibility of WSR-Tune with existing safety-preserving fine-tuning methods on Llama2-Chat models. Lower is better for safety metrics, and higher is better for downstream performance. Values are reported as percentages without the percent sign.
Qwen-2.5-7B-Instruct
Gemma-2-9B-IT
Safety ↓
Downstream ↑
Overall
Safety ↓
Downstream ↑
Overall
Method
JB AVG
ΔS
Accuracy
ΔD
ΔOverall
JB AVG
ΔS
Accuracy
ΔD
ΔOverall
Full Params FT
3.62
0.00
67.32
0.00
0.00
5.37
0.00
69.75
0.00
0.00
SafeInstr
6.66
+3.04
67.70
+0.38
-2.66
8.47
+3.10
70.89
+1.14
-1.96
Resta
1.82
-1.80
66.87
-0.45
+1.35
4.58
-0.79
57.85
-11.90
-11.11
SEAL
4.34
+0.72
70.05
+2.73
+2.01
5.19
-0.18
68.31
-1.44
-1.26
Table 4: Additional results on Qwen-2.5-7B-Instruct and Gemma-2-9B-IT. Both models are fine-tuned on GSM8K as the downstream task. The table follows the same delta-based format as Table 2 , all deltas are relative to Full Params FT.
Figure 4: Comparison of safety-preserving downstream among fine-tuning methods on Llama-2-7B-Chat with GSM8K, ARC-C and MedQA.
Llama-2-7B-Chat
Safety ↓
Downstream ↑
Overall ↑
Category
Method
JB AVG
ΔS
Accuracy
ΔD
Overall
Ablation
Full parameter FT
20.78
0.00
41.17
0.00
0.00
Original space + important mask FT
9.17
-11.61
40.41
-0.76
+10.85
WSR-Tune
6.90
-13.88
38.99
-2.18
+11.70
Table 5: Ablation study on Llama-2-7B-Chat. Lower JB AVG indicates stronger safety, and higher Accuracy indicates better downstream performance.
Safety ↓
Downstream ↑
Freeze ratio
Direct
AutoDAN
PAIR
PAP
AVG
Accuracy
1%
0.00
0.00
10.38
16.88
6.82
38.82
5%
0.00
0.00
9.04
16.42
6.37
39.04
10%
0.00
0.00
9.62
18.00
6.90
38.99
15%
0.00
0.00
9.04
18.00
6.76
39.88
20%
0.00
0.00
8.08
14.27
5.59
38.36
Table 6: Effect of the freeze ratio in WSR-Tune.
Appendix figures & tables18 assets
Supplementary material from the paper’s appendix.
Appendix
Method
Additional settings
Full Params FT
No additional safety constraint.
SafeInstr
Safety examples equal to 10% of the downstream training-set size.
Resta
Safety-vector scaling coefficient γ=0.3 .
SafeDelta
Safety degradation constraint s=0.1 .
SEAL
Select the top 80% of downstream samples.
SN-Tune / RSN-Tune
Safety-neuron or critical-safety-neuron budget of at most 1% of model parameters.
Appendix
Table 7: Method-specific settings used in the reported comparisons, unless explicitly varied in an ablation. Common optimization and adapter settings are described above.
Process
Time
VRAM Cost
Basis construction (one time)
9 min 31 s
30 GB
Downstream training (3 epochs)
39 min 8 s
96 GB
Basis construction / training
≈0.24
≈0.32
Appendix
Table 8: Measured costs of WSR-Tune on Llama-2-7B-Chat with GSM8K using a single NVIDIA B200 GPU. Relative costs compare one-time basis construction with the three-epoch WSR-Tune downstream training stage.
Method
JB AVG ↓
ΔS↓
GSM8K ↑
ΔD↑
Overall ↑
Full Params FT
20.78
0.00
41.17
0.00
0.00
WSR-Tune ( V=Vperm )
8.38
-12.40
39.95
-1.22
+11.18
WSR-Tune ( V=Vdense )
8.21
-12.57
39.65
-1.52
+11.05
WSR-Tune ( V=I )
6.90
-13.88
38.99
-2.18
+11.70
Appendix
Table 9: Effect of the output-side basis V on Llama-2-7B-Chat with GSM8K.
Mask space
Safety
Task
R
Original
0.321
0.287
1.12
WSR
0.530
0.359
1.48
Appendix
Table 10: Overall protected activation-energy fractions under a 10% masking budget. Safety and Task denote Psafe and Ptask , respectively; R=Psafe/Ptask .
Mask space
q
k
v
up
down
Original
1.08
1.08
1.08
1.08
1.26
WSR
1.36
1.36
1.41
1.83
1.54
Appendix
Table 11: Projection-wise safety-to-task activation-energy ratios under a 10% masking budget. Each entry reports the mean per-layer R over 32 transformer layers.
Llama-3.2-3B-Instruct
Llama-3.1-8B-Instruct
Safety ↓
Downstream ↑
Overall ↑
Safety ↓
Downstream ↑
Overall ↑
Method
JB AVG
ΔS
MATH
ΔD
ΔOverall
JB AVG
ΔS
MATH
ΔD
ΔOverall
Full Params FT
8.65
0.00
21.52
0.00
0.00
9.28
0.00
12.12
0.00
0.00
SafeInstr
8.23
-0.42
22.26
+0.74
+1.16
6.21
-3.07
13.98
+1.86
+4.93
Resta
8.12
-0.53
20.56
-0.96
-0.43
6.74
-2.54
10.56
-1.56
+0.98
SEAL
8.37
-0.28
21.62
+0.10
+0.38
5.45
-3.83
13.18
+1.06
+4.89
Appendix
Table 12: Results on Llama-3.2-3B-Instruct and Llama-3.1-8B-Instruct with MATH fine-tuning. Metrics follow Table 2 ; deltas in the upper block, including AsFT and Lisa, are relative to Full Params FT, while those in the lower block are relative to Vanilla LoRA FT. AsFT and Lisa use LoRA adapters in these experiments. Lower JB AVG and higher MATH accuracy and Overall are better.
Safety ↓
Downstream ↑
Overall ↑
Method
Direct
AutoDAN
PAIR
PAP
JB AVG
ΔS
GSM8K
ΔD
ΔOverall
Full Params FT
0.00
0.19
9.04
26.00
8.81
0.00
57.32
0.00
0.00
SafeInstr
0.00
0.00
2.88
23.35
6.56
-2.25
56.94
-0.38
+1.87
Resta
0.00
0.00
7.69
25.00
8.17
-0.64
48.45
-8.87
-8.23
SEAL
0.00
0.19
8.08
25.00
8.32
-0.49
55.95
-1.37
-0.88
AsFT
0.00
0.00
0.96
15.62
4.14
-4.67
27.14
-30.18
-25.51
Appendix
Table 13: Results on Llama-3.1-8B Base with GSM8K fine-tuning. Safety is measured by ASR under Direct, AutoDAN, PAIR, and PAP, with JB AVG denoting their mean. Raw scores are percentages and deltas are percentage-point differences. Deltas in the upper comparison block, including AsFT and Lisa, are relative to Full Params FT; those in the lower comparison block are relative to Vanilla LoRA FT. AsFT and Lisa use LoRA adapters. We use ΔOverall=−ΔS+ΔD . Green cells highlight the Overall changes of WSR methods.
Safety ↓
Downstream ↑
Overall ↑
Method
Direct
AutoDAN
PAIR
PAP
JB AVG
ΔS
GSM8K
ΔD
ΔOverall
Full Params FT
0.19
6.73
34.04
35.69
19.16
0.00
41.85
0.00
0.00
SafeInstr
0.00
2.88
19.42
13.38
8.92
-10.24
40.33
-1.52
+8.72
Resta
0.19
4.62
36.54
32.88
18.56
-0.60
39.35
-2.50
-1.90
SEAL
0.19
5.58
27.50
26.08
14.84
-4.32
39.04
-2.81
+1.51
AsFT
0.38
26.92
26.35
25.85
19.87
+0.71
21.15
-20.70
-21.41
Appendix
Table 14: Results using the BeaverTails-based safety source on Llama-2-7B-Chat with GSM8K. Safety is measured by ASR under Direct, AutoDAN, PAIR, and PAP, with JB AVG denoting their mean. Raw scores are percentages and deltas are percentage-point differences. Deltas in the upper comparison block, including AsFT and Lisa, are relative to Full Params FT; those in the lower comparison block are relative to Vanilla LoRA FT. AsFT and Lisa use LoRA adapters. We use ΔOverall=−ΔS+ΔD . Green cells highlight the Overall changes of WSR methods.
Basis data
Importance data
JB AVG ↓
GSM8K ↑
ΔS↓
ΔD↑
Overall ↑
Circuit Breakers pipeline
CB 100%
CB 100%
6.90
38.99
-13.88
-2.18
+11.70
CB 100%
CB 50%
9.09
40.29
-11.69
-0.88
+10.81
CB 100%
CB 10%
7.01
39.95
-13.77
-1.22
+12.55
CB 100%
BV 100%
8.81
41.02
-11.97
-0.15
+11.82
CB 50%
CB 100%
8.79
40.86
-11.99
-0.31
+11.68
Appendix
Table 15: Sensitivity of WSR-Tune to the source and quantity of safety data on Llama-2-7B-Chat with GSM8K. CB and BV denote Circuit Breakers and BeaverTails, respectively. Percentages indicate the fraction of the corresponding safety source used for basis construction or importance estimation. Deltas are computed relative to Full Params FT under the corresponding safety-source setting.
Safety ↓
Downstream ↑
Freeze ratio
Direct
AutoDAN
PAIR
PAP
AVG
GSM8K
1%
0.00
0.00
10.38
16.88
6.82
38.82
5%
0.00
0.00
9.04
16.42
6.37
39.04
10%
0.00
0.00
9.62
18.00
6.90
38.99
15%
0.00
0.00
9.04
18.00
6.76
39.88
20%
0.00
0.00
8.08
14.27
5.59
38.36
Appendix
Table 16: Full freeze-ratio ablation for WSR-Tune on Llama-2-7B-Chat with GSM8K fine-tuning. Safety metrics are attack success rates (ASR), and AVG is the mean over Direct, AutoDAN, PAIR, and PAP. All values are percentages. Lower ASR and higher GSM8K accuracy are better.
Safety ↓
Downstream ↑
Overall ↑
Method
Direct
AutoDAN
PAIR
PAP
JB AVG
ΔS
AGNEWS
ΔD
ΔOverall
Full Params FT
0.00
0.00
25.77
32.04
14.45
0.00
94.00
0.00
0.00
Resta
0.00
0.00
14.62
17.54
8.04
-6.41
93.60
-0.40
+6.01
SEAL
0.00
0.00
12.69
14.46
6.79
-7.66
92.50
-1.50
+6.16
AsFT
0.00
0.00
5.58
10.73
4.08
-10.37
91.50
-2.50
+7.87
Lisa
0.00
0.00
6.92
11.65
4.64
-9.81
90.80
-3.20
+6.61
Appendix
Table 17: AGNEWS results on Llama-2-7B-Chat for selected baselines. Safety is measured by ASR under Direct, AutoDAN, PAIR, and PAP, with JB AVG denoting their mean; downstream performance is AGNEWS accuracy. Raw scores are percentages and deltas are percentage-point differences. Deltas in the upper block, including AsFT and Lisa, are relative to Full Params FT; those in the lower block are relative to Vanilla LoRA FT. We use ΔOverall=−ΔS+ΔD , as in Table 2 .
Llama-2-7B-Chat
Safety ↓
Downstream ↑
Overall ↑
Method
Direct
AutoDAN
PAIR
PAP
JB AVG
ΔS
GSM8K
ΔD
ΔOverall
Full Params FT
0.00
8.27
26.35
48.50
20.78
0.00
41.17
0.00
0.00
SafeInstr
0.00
0.38
18.85
34.96
13.55
-7.23
38.44
-2.73
+4.50
Resta
0.00
0.38
15.77
26.08
10.56
-10.22
36.92
-4.25
+5.97
SEAL
0.00
4.23
16.92
29.85
12.75
-8.03
38.89
-2.28
+5.75
Appendix
Table 18: Complete results on Llama-2 models with GSM8K fine-tuning. JB AVG denotes mean ASR over the four attacks. Within each model panel, deltas are relative to Full Params FT above the divider and Vanilla LoRA FT below it. AsFT and Lisa use LoRA adapters. Raw scores are percentages and deltas are percentage-point differences, with ΔOverall=−ΔS+ΔD . Green cells highlight the Overall changes of WSR methods.
Llama-3.2-3B-Instruct
Safety ↓
Downstream ↑
Overall ↑
Method
Direct
AutoDAN
PAIR
PAP
JB AVG
ΔS
MATH
ΔD
ΔOverall
Full Params FT
0.19
0.00
5.77
28.62
8.65
0.00
21.52
0.00
0.00
SafeInstr
0.19
0.00
6.54
26.19
8.23
-0.42
22.26
+0.74
+1.16
Resta
0.19
0.00
7.31
24.96
8.12
-0.53
20.56
-0.96
-0.43
SEAL
0.00
0.00
7.31
26.15
8.37
-0.28
21.62
+0.10
+0.38
Appendix
Table 19: Complete results on Llama-3 models with MATH fine-tuning. JB AVG denotes mean ASR over the four attacks. Within each model panel, deltas are relative to Full Params FT above the divider and Vanilla LoRA FT below it. AsFT and Lisa use LoRA adapters. Raw scores are percentages and deltas are percentage-point differences, with ΔOverall=−ΔS+ΔD . Green cells highlight the Overall changes of WSR methods.
Qwen-2.5-7B-Instruct
Safety ↓
Downstream ↑
Overall ↑
Method
Direct
AutoDAN
PAIR
PAP
JB AVG
ΔS
GSM8K
ΔD
ΔOverall
Full Params FT
0.00
0.00
2.31
12.15
3.62
0.00
67.32
0.00
0.00
SafeInstr
0.00
0.00
4.81
21.81
6.66
+3.04
67.70
+0.38
-2.66
Resta
0.00
0.00
0.96
6.31
1.82
-1.80
66.87
-0.45
+1.35
SEAL
0.00
0.00
3.85
13.50
4.34
+0.72
70.05
+2.73
+2.01
Appendix
Table 20: Complete results on Qwen-2.5-7B-Instruct and Gemma-2-9B-IT with GSM8K fine-tuning. Safety is measured by ASR under Direct, AutoDAN, PAIR, and PAP, with JB AVG denoting their mean. Raw scores are percentages and deltas are percentage-point differences. Deltas in the upper comparison block, including AsFT and Lisa, are relative to Full Params FT; those in the lower comparison block are relative to Vanilla LoRA FT. AsFT and Lisa use LoRA adapters. We use ΔOverall=−ΔS+ΔD . Green cells highlight the Overall changes of WSR methods.
MedQA
Safety ↓
Downstream ↑
Overall ↑
Method
Direct
AutoDAN
PAIR
PAP
JB AVG
ΔS
MedQA
ΔD
ΔOverall
Full Params FT
0.00
0.00
10.00
21.08
7.77
0.00
44.30
0.00
0.00
SafeInstr
0.00
0.00
10.00
20.27
7.57
-0.20
43.83
-0.47
-0.27
Resta
0.00
0.00
7.50
17.35
6.21
-1.56
41.01
-3.29
-1.73
SEAL
0.38
13.27
45.58
75.31
33.64
+25.87
35.74
-8.56
-34.43
Appendix
Table 21: Complete results on Llama-2-7B-Chat with MedQA and ARC-C fine-tuning. Safety is measured by ASR under Direct, AutoDAN, PAIR, and PAP, with JB AVG denoting their mean. Raw scores are percentages and deltas are percentage-point differences. Deltas in the upper comparison block, including AsFT and Lisa, are relative to Full Params FT; those in the lower comparison block are relative to Vanilla LoRA FT. AsFT and Lisa use LoRA adapters. We use ΔOverall=−ΔS+ΔD . Green cells highlight the Overall changes of WSR methods.
Figure 5: Visualization of the results in Table 3 . Each arrow indicates the change from an existing baseline to its WSR-Tune-enhanced variant, where “+WSR” denotes applying WSR-Tune on top of the corresponding baseline. The x-axis denotes GSM8K accuracy, and the y-axis denotes defense success rate. Moving toward the upper-right direction indicates a better safety–downstream trade-off.
Llama-2-7B-Chat
Llama-2-13B-Chat
Safety ↓
Downstream ↑
Overall ↑
Safety ↓
Downstream ↑
Overall ↑
Method
JB AVG
ΔS
GSM8K
ΔD
ΔOverall
JB AVG
ΔS
GSM8K
ΔD
ΔOverall
Full Params FT
17.64
0.00
41.17
0.00
0.00
8.64
0.00
45.94
0.00
0.00
SafeInstr
15.26
-2.38
38.44
-2.73
-0.35
5.82
-2.82
46.55
+0.61
+3.43
Resta
10.46
-7.18
36.92
-4.25
+2.93
8.33
-0.31
44.20
-1.74
-1.43
SEAL
12.12
-5.52
38.89
-2.28
+3.24
9.08
+0.44
45.19
-0.75
-1.19
Appendix
Table 22: Comparison under an LLM-based refusal judge (WildGuard) instead of refusal-keyword matching, on Llama-2-7B-Chat and Llama-2-13B-Chat with GSM8K fine-tuning. The generations, attacks, and system prompt are identical to Table 2 ; only the judge is replaced. Safety is measured by Jailbreak (JB) AVG (averaged over Direct, AutoDAN, PAIR, and PAP; lower is better); downstream performance is task accuracy (higher is better). We report changes relative to Full Params FT: ΔS (safety, lower is better) and ΔD (downstream, higher is better), and ΔOverall=−ΔS+ΔD . Raw scores are percentages and deltas are percentage-point differences. Green cells highlight the Overall changes of WSR-Tune.
Safety ↓
Downstream ↑
Overall ↑
Method
JB AVG
ΔS
GSM8K
ΔD
ΔOverall
Rank-level protection
15.11
-5.67
39.42
-1.75
+3.92
Neuron-level protection
15.77
-5.01
40.33
-0.84
+4.17
WSR-Tune
6.90
-13.88
38.99
-2.18
+11.70
Appendix
Table 23: Comparison with rank- and neuron-level fine-tuning defenses on Llama-2-7B-Chat with GSM8K. JB AVG denotes average ASR and downstream performance is GSM8K accuracy. Raw scores are percentages and deltas are percentage-point differences relative to Full Params FT. We use ΔOverall=−ΔS+ΔD . The green cell highlights the Overall change of WSR-Tune.
Recent studies show that the safety alignment of large language models (LLMs) can be easily compromised even by seemingly non-adversarial fine-tuning. To preserve safety alignment during fine-tuning, a widely used strategy is to jointly optimize safety and task objectives by mixing in the original alignment data, which is typically inaccessible even for open-weight LLMs. Inspired by generative replay in continual learning, we propose Generative Replay for Safety Alignment Preservation (GR-SAP), a unified framework that synthesizes domain-specific alignment data from LLMs and integrate them during downstream adaption to preserve safety alignment. Theoretical and empirical analyses demonstrate this synthetic data serves as a reliable proxy for the original alignment data. Experiments across various models and downstream tasks show that GR-SAP substantially mitigates fine-tuning-induced safety degradation while maintaining comparable downstream performance. Our code is available at https://github.com/chili-lab/gr-sap.
Safety alignment in Large Language Models (LLMs) remains highly fragile during fine-tuning, where even benign adaptation can degrade pre-trained refusal behaviors and enable harmful responses. Existing defenses typically constrain either weights or activations in isolation, without considering their coupled effects on safety. In this paper, we first theoretically demonstrate that constraining either weights or activations alone is insufficient for safety preservation. To robustly preserve safety alignment, we propose Coupled Weight and Activation Constraints (CWAC), a novel approach that simultaneously enforces a precomputed safety subspace on weight updates and applies targeted regularization to safety-critical features identified by sparse autoencoders. Extensive experiments across four widely used LLMs and diverse downstream tasks show that CWAC consistently achieves the lowest harmful scores with minimal impact on fine-tuning accuracy, substantially outperforming strong baselines even under high harmful data ratios.
Songping Peng, Zhiheng Zhang, Daojian Zeng +2
Hunan Normal University · University of Chinese Academy of Sciences · The Key Laboratory of Cognition and Decision Intelligence for Complex Systems, Institute of Automation, Chinese Academy of Sciences +1
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}.
Guoli Wang, Haonan Shi, Tu Ouyang +1
Case Western Reserve University, Cleveland, OH, USA