Guardrails in Logit Space: Safety Token Regularization for LLM Alignment

• URL: http://arxiv.org/abs/2604.17210v1
• Date: Sun, 19 Apr 2026 02:41:04 GMT
• Title: Guardrails in Logit Space: Safety Token Regularization for LLM Alignment
• Abstract summary: Fine-tuning well-aligned large language models (LLMs) on new domains often degrades their safety alignment.<n>Existing safety alignment techniques primarily focus on pretraining, leaving fine-tuned models vulnerable to behavioral shifts.<n>We introduce safety token regularization (STR), a lightweight method designed to preserve safety properties during fine-tuning.
• Score: 9.180430520680611
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Preventing Safety Drift in Large Language Models via Coupled Weight and Activation Constraints [19.721346201616782]
  Coupled Weight and Activation Constraints (CWAC) is a novel approach that simultaneously enforces a precomputed safety subspace on weight updates.<n>CWAC consistently achieves the lowest harmful scores with minimal impact on fine-tuning accuracy.
  arXiv  Detail & Related papers  (2026-04-14T07:17:55Z)
• Few Tokens, Big Leverage: Preserving Safety Alignment by Constraining Safety Tokens during Fine-tuning [4.839980912290382]
  Large language models (LLMs) often require fine-tuning (FT) to perform well on downstream tasks.<n>FT can induce safety-alignment drift even when the training dataset contains only benign data.<n>We propose a fine-tuning framework called Preserving Safety Alignment via Constrained Tokens (PACT), which stabilizes the model's confidence on safety tokens.
  arXiv  Detail & Related papers  (2026-03-08T03:42:55Z)
• BarrierSteer: LLM Safety via Learning Barrier Steering [83.12893815611052]
  BarrierSteer is a novel framework that formalizes safety by embedding learned non-linear safety constraints directly into the model's latent representation space.<n>We show that BarrierSteer substantially reduces adversarial success rates, decreases unsafe generations, and outperforms existing methods.
  arXiv  Detail & Related papers  (2026-02-23T18:19:46Z)
• Q-realign: Piggybacking Realignment on Quantization for Safe and Efficient LLM Deployment [55.14890249389052]
  Existing defenses either embed safety recovery into fine-tuning or rely on fine-tuning-derived priors for post-hoc correction.<n>We propose textttQ-realign, a post-hoc defense method based on post-training quantization.<n>Our work provides a practical, turnkey solution for safety-aware deployment.
  arXiv  Detail & Related papers  (2026-01-13T00:07:24Z)
• UpSafe$^\circ$C: Upcycling for Controllable Safety in Large Language Models [67.91151588917396]
  Large Language Models (LLMs) have achieved remarkable progress across a wide range of tasks, but remain vulnerable to safety risks such as harmful content generation and jailbreak attacks.<n>We propose UpSafe$circ$C, a unified framework for enhancing LLM safety through safety-aware upcycling.<n>Our results highlight a new direction for LLM safety: moving from static alignment toward dynamic, modular, and inference-aware control.
  arXiv  Detail & Related papers  (2025-10-02T16:43:33Z)
• Rethinking Safety in LLM Fine-tuning: An Optimization Perspective [56.31306558218838]
  We show that poor optimization choices, rather than inherent trade-offs, often cause safety problems, measured as harmful responses to adversarial prompts.<n>We propose a simple exponential moving average (EMA) momentum technique in parameter space that preserves safety performance.<n>Our experiments on the Llama families across multiple datasets demonstrate that safety problems can largely be avoided without specialized interventions.
  arXiv  Detail & Related papers  (2025-08-17T23:46:36Z)
• Shape it Up! Restoring LLM Safety during Finetuning [65.75757313781104]
  Finetuning large language models (LLMs) enables user-specific customization but introduces critical safety risks.<n>We propose dynamic safety shaping (DSS), a framework that uses fine-grained safety signals to reinforce learning from safe segments of a response while suppressing unsafe content.<n>We present STAR-DSS, guided by STAR scores, that robustly mitigates finetuning risks and delivers substantial safety improvements across diverse threats, datasets, and model families.
  arXiv  Detail & Related papers  (2025-05-22T18:05:16Z)
• LookAhead Tuning: Safer Language Models via Partial Answer Previews [62.529794567687354]
  Fine-tuning enables large language models to adapt to specific domains, but often compromises their previously established safety alignment.<n>We introduce LookAhead Tuning, a lightweight and effective data-driven approach that preserves safety during fine-tuning.
  arXiv  Detail & Related papers  (2025-03-24T18:11:42Z)
• Safe Reinforcement Learning with Learned Non-Markovian Safety Constraints [15.904640266226023]
  We design a safety model that performs credit assignment to assess contributions of partial state-action trajectories on safety. We derive an effective algorithm for optimizing a safe policy using the learned safety model. We devise a method to dynamically adapt the tradeoff coefficient between safety reward and safety compliance.
  arXiv  Detail & Related papers  (2024-05-05T17:27:22Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.