PSRT: Accelerating LRM-based Guard Models via Prefilled Safe Reasoning Traces

Jiawei Zhao, Yuang Qi, Weiming Zhang, Nenghai Yu, Kejiang Chen

Keywords: AI Safety, LRM, Inference acceleration, Guard Model

TL;DR: We replace the LRM-based guard model’s reasoning process with a prefilled safe reasoning trace, thereby preserving its capability while significantly reducing the computational overhead.

Abstract:

Large Reasoning Models (LRMs) have demonstrated remarkable performance on tasks such as mathematics and code generation. Motivated by these strengths, recent work has empirically demonstrated the effectiveness of LRMs as guard models in improving harmful query detection. However, LRMs typically generate long reasoning traces during inference, causing substantial computational overhead.
In this paper, we introduce $\textbf{PSRT}$, a method that replaces the model's reasoning process with a $\textbf{P}$refilled $\textbf{S}$afety $\textbf{R}$easoning $\textbf{T}$race, thereby significantly reducing the inference cost of LRMs. Concretely, PSRT prefills "safe reasoning virtual tokens" from a constructed dataset and learns over their continuous embeddings. With the aid of indicator tokens, PSRT enables harmful-query detection in a single forward pass while preserving the classification effectiveness of LRMs.
We evaluate PSRT on 7 models, 13 datasets, and 8 jailbreak methods. In terms of efficiency, PSRT completely removes the overhead of generating reasoning tokens during inference. In terms of classification performance, PSRT achieves nearly identical accuracy, with only a minor average F1 drop of 0.015 across 7 models and 5 datasets

Primary Area: alignment, fairness, safety, privacy, and societal considerations

Submission Number: 8365