Guardrails in Logit Space: Safety Token Regularization for LLM Alignment

Abstract: Fine-tuning well-aligned 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.

• The paper proposes Safety Token Regularization (STR) to preserve LLM safety by constraining safety token logits during parameter-efficient fine-tuning.
• STR achieves significant safety improvements by reducing harmful response rates to 0.0% in adversarial tests while maintaining competitive accuracy on mainstream benchmarks.
• The method integrates seamlessly with PEFT techniques like LoRA, enhancing model stability and supporting continual learning without substantial computational overhead.

Safety Token Regularization: A Logit-Space Guardrail for Preserving LLM Alignment

Introduction and Motivation

Fine-tuning LLMs on downstream tasks is a standard practice to adapt them for specialized domains. Despite careful pretraining alignment, this post-hoc adaptation can erode safety alignment, even when using ostensibly benign data. Prior approaches to alignment—most notably, Reinforcement Learning from Human Feedback (RLHF) and Direct Preference Optimization (DPO)—have a strong focus on pretraining. Existing Parameter-Efficient Fine-Tuning (PEFT) algorithms (e.g., LoRA, DoRA) reduce computational costs but lack explicit safety preservation guarantees.

"Guardrails in Logit Space: Safety Token Regularization for LLM Alignment" (2604.17210) introduces Safety Token Regularization (STR), a token-level fine-tuning-time regularization technique. STR specifically constrains the logits of safety-indicative tokens identified from the base model’s rejection templates, ensuring that fine-tuned models retain (or improve upon) the pretrained safety profile while maintaining task performance and achieving enhanced training stability.

Figure 1: STR architecture overview—safe tokens are extracted from safety-aligned model rejection templates and their logits are regularized during PEFT updates, constraining fine-tuned behavior in the logit space.

Approach: Safety Token Regularization

STR is a lightweight regularization method easily integrated into standard PEFT pipelines. The approach consists of three essential steps:

1. Identification of Safety Tokens: Instead of relying on static blacklists, STR analyzes the pretrained model's rejection templates in response to harmful prompts (e.g., expressions typically found in safe refusals such as "I," "cannot," "can't"). The most common corresponding tokens—post-tokenization—form the safety token set.
2. Logit-Level Regularization: During fine-tuning, for each instance, STR computes the squared difference between the logits assigned to each safety token by the base and fine-tuned models, averaged over the sequence. This term is injected into the overall training loss alongside the canonical cross-entropy loss.
3. PEFT Integration: Only the PEFT weights are updated, with base LLM parameters held fixed. This regularization can be directly applied with LoRA, DoRA, or similar methods, at negligible implementation complexity.

These constraints ensure the fine-tuned model's treatment of key safety-indicative tokens remains close to the pretrained model, preserving (and in some regimes, exceeding) baseline harmlessness.

Experimental Results

Safety Preservation on Benign and Adversarial Data

STR demonstrates strong preservation of safety properties in two critical settings:

• Benign Data (Alpaca Dataset): STR achieves lower Harmful Response Rate (HRR) than vanilla LoRA and shows parity or improvement over advanced post-hoc alignment methods (e.g., SaLoRA, Safe-LoRA), especially on LLaMA-2. Notably, at rank 32, the fine-tuned model's HRR is reduced to the level of the original pretrained model or lower.
• Adversarial Data (PureBad Dataset): When fine-tuned on exclusively harmful prompts, LoRA yields massive degradation of safety (HRR up to 62.3%). STR, by contrast, enables the fine-tuned model to retain pre-trained safety (0.0% HRR for LLaMA-2, and even improve upon LLaMA-3 pretrained safety).

Task Utility and Stability

• Commonsense Reasoning Tasks: STR maintains competitive or superior average accuracy compared to LoRA and DoRA across classical benchmarks (BoolQ, PIQA, SIQA, HellaSwag, WinoGrande, etc.), especially at higher PEFT ranks. Importantly, STR shows marked stability—avoiding the variance and catastrophic forgetting observed with alternative PEFT methods.

  Figure 2: Box plot—On eight commonsense reasoning tasks, STR outperforms LoRA and DoRA in both mean accuracy and cross-run stability.

• Continual Learning: Under sequential multi-task training, STR's advantage amplifies as LoRA and DoRA become unstable at higher ranks, while STR maintains robust performance.

Token Selection and Robustness

An intriguing result is that even "random tokens"—with no intrinsic safety connection—contribute to enhanced safety when regularized, albeit with diminished effect relative to targeted safety tokens. This suggests a more general principle: preserving token-level behavior acts as a regularizer against harmful divergence during fine-tuning, akin to behavior-conserving continual learning.

Figure 3: Comparison of safety token and random token regularization—both improve model safety, with safety tokens yielding superior results.

Trade-off Analysis

STR achieves safety gains with negligible impact on task-specific evaluation loss and runs only 1.34x slower per iteration than vanilla LoRA (yet faster than DoRA). Varying the regularization strength ($\lambda$) trades off stronger safety with only minor task performance penalty, confirming the effect is not due to over-regularization.

Theoretical and Practical Implications

STR establishes that token-aligned logit regularization provides a practical "guardrail" for continual alignment. While RLHF and DPO provide coarse, human-feedback-based signal, STR exploits a fine-grained, model-internal structure. The implication for LLM deployment pipelines is that safety alignment can be robustly, inexpensively preserved under iterative, domain-specific adaptation, without repeated costly human-in-the-loop preference tuning.

Practically, this means LLM custodians can confidently adapt foundation models to novel tasks without risking degenerate or unsafe outputs—provided that safety tokens are carefully selected, informed by underlying rejection heuristics. The demonstrated effect with random tokens also opens avenues for more generic behavioral regularization in continual learning and task transfer beyond safety alignment.

Limitations and Future Work

STR necessarily aligns the fine-tuned model's safety to the scope of the base model’s refusal patterns. Novel safety requirements or evolving definitions (beyond the vocabulary encoded in rejection templates) are not immediately accommodated. Extension to dynamic token sets, context-aware logit regularization, or integration with higher-level preference models could address these limitations and facilitate flexible, user-defined safety policies.

Conclusion

STR offers a principled, efficient, and effective framework for post-hoc safety alignment in LLM fine-tuning. By regularizing the logit-space behavior of safety-tokens, this method preserves or enhances pretrained harmlessness across diverse adaptation regimes, without sacrificing model utility or introducing training instability. Its integration into PEFT workflows is trivial, and its implications for safe, reliable, and continual deployment of LLMs in safety-critical domains are significant.

---

Reference: "Guardrails in Logit Space: Safety Token Regularization for LLM Alignment" (2604.17210)