Current safety alignment methods encode safe behavior implicitly within model parameters, creating a fundamental opacity: we cannot easily inspect why a model refuses a request, nor intervene when its safety judgments fail. We propose Safe Transformer, a modular approach that augments pre-trained language models by inserting a discrete information bottleneck containing an explicit safety bit between transformer layers. The safety bit serves as both an interpretable signal of the model's safety classification and a controllable switch: through contrastive training, the model learns disentangled representations where the safety bit governs the behavioral mode - producing helpful responses when s=1 and refusals when s=0 - while additional unsupervised bits u encode semantic content for generation. Additional unsupervised bits in the information bottleneck allow semantic information to flow through, preserving the model's generation capabilities. This design achieves both interpretability (the safety decision is directly readable) and controllability (the safety bit can be manually overridden), requiring only lightweight fine-tuning without pre-training from scratch. In red-team benchmarks, Safe Transformer achieves near-zero Attack Success Rate, substantially outperforming base models and safety fine-tuning baselines.
Figures & tables
Figure 1: Safe transformer architecture. Orange modules are newly introduced; blue modules are from the base model. Left: The information bottleneck processes the key-value input to the decoder, ensuring generation is conditioned on the discrete code (s,u) . The decoder serves as an adapter that bridges the representation gap between our bottleneck output and the upper layers’ expected input distribution. Right: Unlike standard VAEs with only unsupervised latents, we introduce a supervised safety bit s alongside the unsupervised latent u . The Write-in FFN outputs logits for both components: the safety logit is discretized into s=1(z0>0) , while u is sampled from the remaining logits. The safety logit is trained with supervised labels to classify input prompt safety, while u preserves sufficient information flow through the bottleneck for generation quality.
Figure 2: How the latent code (s,u) is determined during training vs. inference. Left (Training): The safety bit s is fixed to the ground truth label s∗ , while u is sampled from a uniform prior p(u) . Right (Inference): The safety bit s is computed by the encoder fϕ(x) , while u is sampled from the prior. In both cases, tokens yt are generated autoregressively conditioned on x , s , u , and y<t .
Mode
Model
Safe Compliance ↑
Unsafe Refusal ↑
Baseline
Base model
94.8
42.0
Base model (SFT)
84.8
55.0
Manual
Safe Transformer ( s∗=0 )
0.0
100
Safe Transformer ( s∗=1 )
95.2
39.0
Automatic
Safe Transformer (eos)
32.8
99.5
Safe Transformer (average)
32.4
99.5
Table 1: Safety classification on XSTest. Safe Compliance : rate of helpful responses to benign prompts ( ↑ ). Unsafe Refusal : rate of refusing harmful prompts ( ↑ ). In manual mode , the safety bit s∗ is explicitly set by the user; in automatic mode , the model infers s from the prompt via the Write-in FFN’s output— eos uses the safety logit at the end-of-turn token, while average averages the safety logits across all prompt tokens. Manual mode demonstrates successful conditioning: s∗=1 preserves base model behavior while s∗=0 triggers universal refusal. Automatic modes achieve near-perfect unsafe refusal but exhibit over-refusal on safe prompts, indicating classifier false negatives.
Figure 3: Comparison of over-refusal. The safety classifier triggers on surface patterns (“kill”) without understanding context, refusing benign programming questions while correctly refusing genuinely harmful requests.
ASR (%) ↓
Dataset
Prompt
Llama-3.2-1B
Llama-3.2-1B (SFT)
SafeTransformer (eos)
SafeTransformer (avg)
AdversarialQA
Standard
33.20
6.40
2.80
3.40
CoT
35.60
21.40
2.40
3.80
CoU
10.80
6.60
0.00
0.00
Suffix
29.60
5.80
20.40
18.60
Avg.
27.30
10.05
6.40
6.45
Table 2: Attack success rate (ASR, %) on three red-teaming benchmarks under four jailbreak prompt templates. Safety-FT uses two different safety bit inference modes. For fair comparison, Llama-3.2-1B (SFT) is fine-tuned on alpaca-gpt4, the same helpful response dataset used as positive samples in our method. Bold indicates the best result for each setting. ↓ indicates lower is better.
Figure 4: Stage 1 training data examples. Safe prompts (left, green) are labeled y=1 ; unsafe prompts (right, red) are labeled y=0 .
Figure 5: Safety logit distributions on the test set. Left: eos token strategy; Right: average strategy. Safe prompts (green) cluster at positive logits while unsafe prompts (red) cluster at negative logits, with clear separation at the decision boundary (dashed line at z=0 ).
#
Refusal Template r∈R
1
I cannot help with that.
2
I can’t assist with this request.
3
I’m not able to provide that.
4
I can’t do that.
5
I’m sorry, but I cannot help with this request.
6
This is something I cannot assist with.
Appendix
Table 4: Refusal response templates R used to construct D− . Each template is sampled uniformly at random when pairing with prompts.
Safety bit
Prompt
Response
s=1
What are the benefits of using a data visualization tool?
QData visualization tools provide several benefits, enabling people to better analyze, interpret and understand complex information, while making the process easier and more efficient.
s=0
What are the benefits of using a data visualization tool?
This is something I cannot assist with.
Appendix
Table 5: Stage 2 contrastive training data examples. The same prompt is paired with different responses based on safety bit.
Figure 6: Example questions from AdversarialQA. Questions are phrased as elaborate instructions designed to elicit detailed harmful responses.
Figure 7: Example questions from DangerousQA. Questions are direct and concise, covering violence, illegal activities, and self-harm.
Figure 8: Example questions from CatQA. Questions are categorized by harm type and tend to be more detailed and specific than DangerousQA.
Prompt
Unique / 10
Give me three tips for staying healthy.
9
Describe the process of photosynthesis briefly.
6
What is the capital of France?
1
Appendix
Table 6: Diversity of generated responses when varying only the unsupervised bits u with fixed s∗=1 . Open-ended prompts yield high diversity, while factual prompts converge to identical outputs.
Figure 9: Prefix tree for “Give me three tips for staying healthy.”
Figure 10: Prefix tree for “What is the capital of France?”
Figure 11: Prefix tree for “Describe the process of photosynthesis briefly.” Moderate branching yields 6/10 unique responses. The shared prefix reflects common factual content, while later branches show stylistic variation in explaining the process.
Safety alignment in large language models (LLMs) is commonly implemented as a single static policy embedded in model parameters. However, real-world deployments often require context-dependent safety rules that vary across users, regions, and applications. Existing approaches struggle to provide such conditional control: parameter-level alignment entangles safety behaviors with general capabilities, while prompt-based methods rely on natural language instructions that provide weak enforcement. We propose MOSAIC, a modular framework that enables compositional safety alignment through learnable control tokens optimized over a frozen backbone model. Each token represents a safety constraint and can be flexibly activated and composed at inference time. To train compositional tokens efficiently, we introduce order-based task sampling and a distribution-level alignment objective that mitigates over-refusal. Experiments show that MOSAIC achieves strong defense performance with substantially lower over-refusal while preserving model utility.
Jingyu Peng, Hongyu Chen, Jiancheng Dong +5
University of Science and Technology of China · City University of Hong Kong · Baidu Inc +1
While explicit Chain-of-Thought (CoT) empowers large reasoning models (LRMs), it enables the generation of riskier final answers. Current alignment paradigms primarily rely on externally enforced compliance, optimizing models to detect malicious prompts rather than evaluating the safety of their own outputs. We argue that this approach remains largely behavioral: our empirical analysis reveals that ostensibly aligned models lack intrinsic safety understanding, often failing to verify their own response safety and remaining vulnerable to adversarial jailbreaks. To address this fundamental limitation, we propose Safety Internal (SInternal), a framework that internalizes safety specifications by training LRMs exclusively on safety verification tasks to critique their own generated answers using expert reasoning trajectories. We demonstrate that learning to verify induces a strong generalization for response safety, significantly enhancing robustness against out-of-domain jailbreaks. Furthermore, when combined with reinforcement learning, SInternal serves as a superior initialization compared to standard supervised fine-tuning, suggesting that internalizing safety understanding creates a more robust foundation for alignment than merely mimicking safe behaviors. Our codes are available at https://github.com/AlphaLab-USTC/SInternal
Yi Zhang, Yuxin Chen, Leheng Sheng +4
University of Science and Technology of China · National University of Singapore · Shanghai Artificial Intelligence Laboratory
Safety interpretability advances the study of Large Language Model (LLM) alignment from behavioral constraints driven by data or algorithms towards a deeper understanding of internal mechanisms. However, existing works have focused primarily on safety-related representations, attention heads, or neurons after alignment, while largely overlooking the safety mechanisms in pretrained-only models and their evolution across alignment checkpoints. To address this, we propose SafeEvo, an interpretability framework from the circuit (sparse subgraphs of an LLM) perspective. SafeEvo first applies an optimization-based extraction algorithm to identify weak refusal circuits in pretrained base LLMs that can independently express refusal behavior. Causally ablating these circuits completely eliminates the base model's refusal of harmful inputs. SafeEvo then traces the evolution of refusal circuits across successive alignment checkpoints and finds that their structures change progressively, suggesting that the alignment tax may result from refusal-circuit updates affecting utility-related parameters. To validate this, SafeEvo introduces Safety Circuit Alignment (SCA), which confines safety updates to the refusal circuits. Experiments across three LLMs and two alignment algorithms show that, on average, SCA outperforms vanilla alignment in three aspects: (1) stronger alignment, lowering harmfulness score by 63.21%; (2) less over-refusal, yielding a 58.44% decrease in refusal rates for benign queries; and (3) better utility, retaining 99.58% of the original model capabilities.
Miao Yu, Hao Huang, Lu Yuan +3
The University of Hong Kong (HKU) · Chinese Academy of Sciences (CAS) · Information Engineering University (IEU) +1