Decoding-Time Safety Alignment via Rollback and Introspection Intervention
Xiaomeng Hu, Fei Huang, Chenhan Yuan, Junyang Lin, Tsung-Yi Ho
As large language models (LLMs) are increasingly deployed in real-world applications, ensuring the safety of their outputs during decoding has become a critical challenge. Existing training-time alignment methods often fail to address dynamic safety risks that emerge during the inference time. While decoding-time interventions, such as Contrastive Decoding, offer adaptability, they face trade-offs between safety and response quality. In this work, we propose a novel framework for decoding-time safety alignment that integrates three key components: (1) a guard model for real-time safety monitoring, enabling detection of potentially unsafe content; (2) a rollback mechanism with a token buffer to correct unsafe outputs efficiently at an earlier stage while without disrupting the user experience; and (3) an introspection-based intervention strategy, where the model generates self-reflective critiques of its previous outputs and incorporates these reflections into the context to guide subsequent decoding steps. This framework benefits from the targeted nature of interventions (thanks to the guard model), the timely nature of corrections (enabled by the rollback mechanism), and the emergence of effective introspection patterns when larger buffers are used. Experimental results demonstrate that our framework achieves superior performance with minimal average wait tokens and a high successful intervention rate, while maintaining high response quality.