SaLR

Overview

Safety diagnosis and SaLR overview

Panels (a) and (b) show jailbreak vulnerability and hidden state separation. Panels (c) and (d) illustrate compact safe-chains and the safety utility comparison.

Why Does Latent Space Reasoning Need Dedicated Safety Alignment?

Textual chain-of-thought (CoT) provides explicit intermediate steps that can be supervised directly, whereas latent reasoning does not expose a readable reasoning trace. Training on final responses therefore only indirectly constrains the preceding latent states. This motivates safety supervision at the transition from latent reasoning to response generation, as implemented in SaLR.

Methodology

Compact traces and response prefix distillation

Safety examples align hidden states across four response prefix positions (K = 4), while mathematical examples use single position alignment at the answer onset.

SaLR rewrites long form safety rationales into compact four-block safe-chains: request intent, risk assessment, response action, and safe follow-up. The action is selected from refuse, partial, or allow.

During training, the teacher learns from explicit reasoning traces and target responses, while the student reasons through latent states. SaLR aligns teacher and student hidden states across the first four response prefix positions (K = 4) for safety examples, and uses single position alignment at the answer onset for mathematical examples.

Results

Across 24 settings (3 model scales × 2 attacks × 4 benchmarks), SaLR achieves a mean jailbreak ASR of at most 3.1% in each setting. On the 3B model, it achieves 59.23% average math accuracy, close to CODI’s 59.86%, while using approximately 40× fewer generated tokens than SFT-CoT + SafeChain. Its benign over-refusal rates are 2.0% on XSTest and 0.33% on OKTest. Together, these results demonstrate a favorable balance between jailbreak robustness, reasoning utility, and benign request handling.

Robustness to Attacks on the Reasoning Process

Beyond standard jailbreaks, we evaluate OverThink token inflation, H-CoT reasoning hijacking, Crescendo multi-turn attacks, and white-box perturbations of latent reasoning states.

Citation

If you find SaLR useful for your research, please consider citing our paper:

@inproceedings{wang2026safetyaware,
  title={Safety-Aware Latent Space Reasoning in Large Language Models},
  author={Yi Wang and Wenjie Wang and Hongye Qiu and Yu Pan},
  booktitle={The Fortieth Annual Conference on Neural Information Processing Systems},
  year={2026},
  url={https://openreview.net/forum?id=v5ExASonPK}
}