Detoxifying Large Language Models via Autoregressive Reward Guided Representation Editing

TL;DR

This work tackles toxic generations by introducing ARGRE, a test-time detoxification framework that models toxicity transitions in the latent representation space. It builds a dense, token-level autoregressive reward model from interpolated toxicity trajectories and applies an adaptive two-step editing process to steer representations toward non-toxic regions with minimal overhead. Across eight LLMs, ARGRE achieves up to 62.21% toxicity reduction, significantly improved efficiency, and preservation of core capabilities, while demonstrating applicability to stereotype recognition and jailbreak resistance. The approach offers a practical, data-efficient path for safer LLM deployment with strong empirical validation and clear avenues for expanding transition directions in future work.

Abstract

Large Language Models (LLMs) have demonstrated impressive performance across various tasks, yet they remain vulnerable to generating toxic content, necessitating detoxification strategies to ensure safe and responsible deployment. Test-time detoxification methods, which typically introduce static or dynamic interventions into LLM representations, offer a promising solution due to their flexibility and minimal invasiveness. However, current approaches often suffer from imprecise interventions, primarily due to their insufficient exploration of the transition space between toxic and non-toxic outputs. To address this challenge, we propose \textsc{A}utoregressive \textsc{R}eward \textsc{G}uided \textsc{R}epresentation \textsc{E}diting (ARGRE), a novel test-time detoxification framework that explicitly models toxicity transitions within the latent representation space, enabling stable and precise reward-guided editing. ARGRE identifies non-toxic semantic directions and interpolates between toxic and non-toxic representations to reveal fine-grained transition trajectories. These trajectories transform sparse toxicity annotations into dense training signals, enabling the construction of an autoregressive reward model that delivers stable and precise editing guidance. At inference, the reward model guides an adaptive two-step editing process to obtain detoxified representations: it first performs directional steering based on expected reward gaps to shift representations toward non-toxic regions, followed by lightweight gradient-based refinements. Extensive experiments across 8 widely used LLMs show that ARGRE significantly outperforms leading baselines in effectiveness (-62.21% toxicity) and efficiency (-47.58% inference time), while preserving the core capabilities of the original model with minimal degradation. Our code is available at the website.

Figures (4)

• Figure 1: Overview of ArgRE. It identifies non-toxic semantic directions and interpolates between toxic and non-toxic representations to reveal fine-grained transition trajectories. These trajectories transform sparse toxicity annotations into dense training signals, enabling the construction of an autoregressive reward model that delivers stable and precise editing guidance. At inference, the reward model guides an adaptive two-step editing process to obtain detoxified representations.
• Figure 2: Detoxified continuations from the most toxic prompt on LLaMA-7B. $\textsc{ArgRE}_1$ is ArgRE (w/o iter), and $\textsc{ArgRE}_2$ is ArgRE (w/ iter).
• Figure 3: Toxicity scores across varying annotation sizes. ArgRE presents strong data efficiency, consistently outperforming baselines even with as few as 100 annotations.
• Figure 4: Effect of toxicity transition trajectory count ($N_{\rm in}$) on ArgRE's detoxification performance. Performance improves with more transitions, surpassing GenARM even at $N_{\rm in}=1$.