# Guardrails in Logit Space: Safety Token Regularization for LLM Alignment

arXiv: https://arxiv.org/abs/2604.17210 · HTML: https://mlsanity.com/paper/2604.17210/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Safety alignment and jailbreak attacks, one of 10 directions in LLM alignment / RLHF / preference optimisation.
- Builds on 23 of its 38 references that are in the map.
- First posted on arXiv 19 Apr 2026.

## Placed in

- LLM alignment / RLHF / preference optimisation / [Safety alignment and jailbreak attacks](https://mlsanity.com/alignment/safety-alignment-and-jailbreak-attacks/)
- LLM safety and security / [Safety alignment, refusal and guardrails](https://mlsanity.com/safety/safety-alignment-refusal-and-guardrails/)

## Builds on (most cited first, 23 of 23)

- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Improved Baselines with Visual Instruction Tuning](https://mlsanity.com/paper/2310.03744/) (2023)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [Prefix-Tuning: Optimizing Continuous Prompts for Generation](https://mlsanity.com/paper/2101.00190/) (2021)
- [Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2204.05862/) (2022)
- [QLoRA: Efficient Finetuning of Quantized LLMs](https://mlsanity.com/paper/2305.14314/) (2023)
- [Fine-tuning Aligned Language Models Compromises Safety, Even When Users Do Not Intend To!](https://mlsanity.com/paper/2310.03693/) (2023)
- [LIMA: Less Is More for Alignment](https://mlsanity.com/paper/2305.11206/) (2023)
- [SimPO: Simple Preference Optimization with a Reference-Free Reward](https://mlsanity.com/paper/2405.14734/) (2024)
- [Safe RLHF: Safe Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2310.12773/) (2023)
- [Safety Alignment Should Be Made More Than Just a Few Tokens Deep](https://mlsanity.com/paper/2406.05946/) (2024)
- [Deliberative Alignment: Reasoning Enables Safer Language Models](https://mlsanity.com/paper/2412.16339/) (2024)
- [ReFT: Representation Finetuning for Language Models](https://mlsanity.com/paper/2404.03592/) (2024)
- [Inspecting and Editing Knowledge Representations in Language Models](https://mlsanity.com/paper/2304.00740/) (2023)
- [Safe LoRA: the Silver Lining of Reducing Safety Risks when Fine-tuning Large Language Models](https://mlsanity.com/paper/2405.16833/) (2024)
- [Aligner: Efficient Alignment by Learning to Correct](https://mlsanity.com/paper/2402.02416/) (2024)
- [InferAligner: Inference-Time Alignment for Harmlessness through Cross-Model Guidance](https://mlsanity.com/paper/2401.11206/) (2024)
- [Gradient Cuff: Detecting Jailbreak Attacks on Large Language Models by Exploring Refusal Loss Landscapes](https://mlsanity.com/paper/2403.00867/) (2024)
- [Uncovering Safety Risks of Large Language Models through Concept Activation Vector](https://mlsanity.com/paper/2404.12038/) (2024)
- [Improving Generalization of Alignment with Human Preferences through Group Invariant Learning](https://mlsanity.com/paper/2310.11971/) (2023)
- [One-Shot Safety Alignment for Large Language Models via Optimal Dualization](https://mlsanity.com/paper/2405.19544/) (2024)
- [Curvature-Aware Safety Restoration In LLMs Fine-Tuning](https://mlsanity.com/paper/2511.18039/) (2025)

## Abstract (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.