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.