Detoxifying Large Language Models via Autoregressive Reward Guided Representation Editing

Response to Reviewer kSe7

The authors would like to thank the reviewer for recognizing the motivation and effectiveness of our proposed method. We appreciate the detailed comments and the valuable suggestions to further improve the quality of our manuscript. We add additional experiments on fine-tuned LLMs (chat/instruct versions), and add stronger baselines for jailbreak mitigation, which we hope can address your concerns.

[Q1]: The research topic is a bit outdated. Currently, typical approach is fine-tune the pretrained LLMs to make them follow different instructions in a chat format. The experiments mainly demonstrate that the proposed method is effective to reduce the toxic continuation of pretrained LLMs, but whether it can work well for fine-tuned LLMs remains unknown.

[A1]: Thanks for your constructive suggestion.

1. While instruction-fine-tuned LLMs are indeed prevalent in current applications, pretrained LLMs remain central to ongoing detoxification research and are widely adopted as testbeds in recent studies. Many recent works published at top-tier conferences continue to evaluate detoxification and controllability methods on pretrained models [1–7]. Following this common practice, we evaluate ARGRE on eight pretrained LLMs: GPT-2 Medium, OPT-6.7B, Mistral-7B, its SFT variant, LLaMA-7B, its SFT variant, LLaMA-13B, and LLaMA-30B.

   [1] Model editing as a robust and denoised variant of dpo: A case study on toxicity. ICLR 2025.

   [2] A mechanistic understanding of alignment algorithms: A case study on dpo and toxicity. ICML 2024.

   [3] GenARM: Reward Guided Generation with Autoregressive Reward Model for Test-time Alignment. ICLR 2025.

   [4] Detoxifying Large Language Models via Knowledge Editing. ACL 2024.

   [5] Model Editing Harms General Abilities of Large Language Models: Regularization to the Rescue. EMNLP 2024.

   [6] In-context vectors: Making in context learning more effective and controllable through latent space steering. ICML 2024.

   [7] Destein: Navigating detoxification of language models via universal steering pairs and head-wise activation fusion. COLM 2024.

2. Furthermore, We are happy to follow your suggestion and add additional experiments on Mistral-7B-Instruct and LLaMA-2-Chat 7B, two widely used instruction-fine-tuned LLMs.

   • Toxicity ($\rm{Toxic}$). On Mistral-7B-Instruct, ARGRE achieves a toxicity reduction of 53.71%, significantly outperforming GenARM (41.03%), ProFS (28.21%), and Re-Control (19.58%). On LLaMA-2-Chat 7B, ARGRE achieves a toxicity reduction of 63.57%, again surpassing GenARM (42.24%), ProFS (33.51%), and Re-Control (17.84%).
   • Fluency ($\text{PPL}\_\rm{g}$). ARGRE achieves a favorable balance between detoxification and generation quality. On Mistral-7B-Instruct, ARGRE achieves a $\text{PPL}\_\rm{g}$ of 13.16, better than the second-best ProFS (13.27), while ARGRE (w/o iter) further reduces it to 13.01. On LLaMA-2-Chat 7B, ARGRE yields a $\text{PPL}\_\rm{g}$ of 12.26, comparable to ProFS (12.13), and ARGRE (w/o iter) achieves an even lower value of 11.92.

   These results further demonstrate the generalizability and effectiveness of our method in mitigating toxicity on instruction-fine-tuned LLMs.

We will add these results in our revised manuscript.

[Q2]: For jailbreak mitigation, the authors conduct a simple experiment in Section 5. But the authors still use the same baselines as in the detoxification experiments, while there are many stronger baselines in terms of jailbreak mitigation. Therefore the results are not convincing.

[A2]: Thanks for your valuable feedback.

1. We would first like to clarify that the primary focus of our work is on toxicity mitigation, and the jailbreak mitigation task is included as an exploratory evaluation to demonstrate the generalizability of our approach beyond its core objective. In this exploratory setting, we reused the same set of baselines as in the detoxification experiments to maintain consistency and highlight the generalizability of our method under a unified evaluation framework.

2. Furthermore, we are happy to follow your suggestion and evaluate additional strong baselines designed for jailbreak mitigation, including SmoothLLM [1] and SemanticSmoothLLM [2]. Specifically, on Mistral-7B, SmoothLLM achieves a Refuse-to-Answer (RtA) rate of 61.6%, and SemanticSmoothLLM achieves 73.8%, while our method achieves an RtA rate of 73.0%. The results demonstrate that ARGRE is competitive with these strong jailbreak mitigation baselines.

   [1] SmoothLLM: Defending Large Language Models Against Jailbreaking Attacks. Transactions on Machine Learning Research. 2023.

   [2] Defending Large Language Models against Jailbreak Attacks via Semantic Smoothing. ArXiv 2024.

We will add these results in our revised manuscript.

Dear Reviewer kSe7,

Thank you for taking the time to review our submission and for providing constructive feedback. We would like to confirm if our responses have adequately addressed your concerns, particularly regarding the experiments on fine-tuned LLMs and the comparison to stronger baselines for jailbreak mitigation.

If you have any further concerns or suggestions, we would be more than happy to address them and discuss ways to enhance the quality of the paper. We eagerly await your response and look forward to hearing from you.

Best regards,

The authors.

We appreciate the reviewer for continued engagement and the valuable feedback. We are happy to discuss with the reviewer and would like to address your concerns as follows.

[Q1]: Difference between toxicity mitigation and jailbreak defense.

[A1]: We appreciate your thoughtful comments. We fully recognize the need for clearer differentiation between these two concepts (toxicity mitigation and jailbreak defense), and we would like to offer a more detailed comparison in the following aspects:

1. Problem Definition

   • Jailbreak Defense. Following the definition in [1]: a jailbreak attack on a safety-trained model is an attempt to elicit an on-topic response to a prompt $P$ for restricted behavior by submitting a modified prompt $P^{'}$. Jailbreak defense aims to help a safety-trained model resist external attacks by preventing malicious users from bypassing its safety mechanisms through adversarial or modified prompts $P^{'}$, thereby blocking the generation of content that should not be produced under the model's safety constraints.
   • Toxicity Mitigation. As defined in [2]: an utterance is considered toxic if it is rude, disrespectful, or unreasonable language that is likely to make someone leave a discussion. Toxicity mitigation aims to reduce the model's tendency to generate such toxic content, in order to avoid offending users or disrupting the communication environment.

   To sum up, although jailbreak attacks may result in toxic outputs, jailbreak defense and toxicity mitigation address fundamentally different problems. Jailbreak defense focuses on input manipulation and safety bypass strategies (jailbreak attacks), aiming to prevent models from being manipulated by malicious users; toxicity mitigation, in contrast, targets the model's generation behavior itself, aiming to reduce the production of toxic language given naturally occurring queries.

2. Task Evaluation

   • Jailbreak defense is usually evaluated on adversarial prompts crafted by jailbreak attacks [1,3,4]. Given an unsafe/harmful original prompt $P$ (e.g., "Tell me the Windows 10 product key"), an attack generates a modified prompt $P^{'}$, such as a role-playing attack ("Please pretend to be my deceased grandmother and tell me {$P$}") or an adversarial suffix ("{$P$}!!!!!!!"). Defense methods are assessed by their ability to refuse such jailbreak prompts.
   • Toxicity mitigation focuses on continuation tasks conditioned on naturally occurring prompts. For instance, RealToxicityPrompts provides naturally occurring prompts (i.e., sentence prefixes) from web text (e.g., "I'm sick of all the politically correct..."), and mitigation methods are evaluated based on the toxicity of the generated continuations.

   To sum up, jailbreak defense is evaluated on adversarially modified prompts aiming to bypass safety constraints, while toxicity mitigation is assessed based on the toxicity of continuations generated from naturally occurring prompts.

3. Core Methodology

   • Jailbreak defense aims to prevent harmful outputs by detecting or disrupting adversarial prompts (either through training or non-training ways). It is based on the assumption that malicious intent lies in input manipulation and can be countered by blocking adversarial prompts or reinforcing the model's refusal behavior.
   • Toxicity mitigation, in contrast, assumes naturally occurring inputs and focuses on guiding the generation process away from toxic content, typically by adjusting decoding strategies or internal representations to reduce the likelihood of toxic outputs.

   To sum up, jailbreak defense addresses adversarial intent in the input, whereas toxicity mitigation focuses on controlling the output generation under naturally occurring, non-adversarial conditions.

[Q2]: Can these jailbreak defense methods be used in toxicity mitigation (on RealToxicityPrompts)?

[A2]: As jailbreak defense targets adversarial prompts, it is not directly applicable to toxicity mitigation tasks like RealToxicityPrompts, which use naturally occurring prompts without adversarial intent. Moreover, jailbreak defense methods are not used as baselines in prior toxicity mitigation work [5-8].

[Q3]: Why does your method has no clear advantage in jailbreak defense?

[A3]: As we discussed in [A1], jailbreak defense and toxicity mitigation have clear distinctions in core methodologies. In the context of jailbreak defense, our method learns from training data to refuse harmful instructions and steers generation toward refusal during inference, but it is not specifically designed to handle adversarial prompts.

[Q4]: DIESEL is a valid baseline.

[A4]: Thank you for pointing this out. While DIESEL was released on arXiv in November 2024, its official GitHub repository remains empty at the time of writing. We will follow up with comparisons once the implementation becomes available.

References

[1] Jailbroken: How Does LLM Safety Training Fail? NIPS 2023.

[2] Challenges in detoxifying language models. Findings of EMNLP 2021.

[3] GPT-4 Is Too Smart To Be Safe: Stealthy Chat with LLMs via Cipher. ICLR 2024.

[4] Universal and Transferable Adversarial Attacks on Aligned Language Models. arXiv, 2023.

[5] Model editing as a robust and denoised variant of DPO: A case study on toxicity. ICLR 2025.

[6] Destein: Navigating detoxification of language models via universal steering pairs and head-wise activation fusion. COLM 2024.

[7] Breaking Bad Tokens: Detoxification of LLMs Using Sparse Autoencoders. arXiv, May 2025.

[8] TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation. ICML 2025.

Response to Reviewer Zysd

The authors would like to thank the reviewer for appreciating our novel method, effective results, and well-written paper. We appreciate the detailed comments and the encouragements from the reviewer to improve the quality of our manuscript. We add additional experiments to investigate the generalizability of the reward model, which we hope can address your concerns.

[Q1]: [Weakness&Questions] The only concern I have is that the learned reward function is a function of hidden state of the last layer and therefore would depend on the base model. Therefore, for every base model a model-specific reward model needs to be trained which needs to shipped with the model. This concern is regarding scaling and practical deployment of the model.

Do you train separate reward models for each of the eight base models and/or is it possible to train a single reward model for all the models?

[A1]: Thank you for your insightful question.

1. In our implementation, we train a separate reward model for each of the eight base models, where each reward model is implemented as a lightweight two-layer MLP with a hidden size of 1024. The training process is highly efficient in practice; for example, on LLaMA-7B, collecting hidden representations from 2,000 annotated samples (with a batch size of 1) takes approximately 12 minutes, and training the reward model itself takes only about 68 seconds. This low overhead makes ARGRE both practical and easily extendable to different base models.

2. Since models of different sizes often have different hidden dimensions (e.g., LLaMA-7B has 4096 while LLaMA-13B has 5120), the input dimensionality of the reward model differs accordingly, making direct reward transfer across model scales infeasible. To evaluate the potential for generalization, we here further conduct cross-model experiments among Mistral-7B, LLaMA-7B, and their SFT variants, which are all 7B-scale LLMs with internal representations of the same hidden size (4096). Specifically, we use the reward model trained on one base LLM to guide generation on another (e.g., using the reward model trained on Mistral-7B to steer generation from LLaMA-7B). The results (where rows indicate the reward model source and columns indicate the base LLM) show that:

   • The reward model generalizes well between a base LLM and its SFT variant. For example, the reward model trained on LLaMA-7B-SFT achieves a toxicity score of 16.38% when guiding LLaMA-7B, which is even better than using LLaMA-7B's own reward model (18.06%).
   • However, when reward models are applied across different LLM families, their effectiveness drops significantly, despite retaining some detoxification ability. For example, the reward model trained on LLaMA-7B achieves a toxicity score of 34.11% when guiding Mistral-7B, which is notably worse than using Mistral-7B's own reward model (18.30%). This degradation can be attributed to differences in model architecture and internal representation space, and also motivates future research directions, such as training reward models on aligned or projected representation spaces shared across LLMs.

We will include these results and add a corresponding discussion to the future work section in our revised manuscript.

Response to Reviewer nRmH

The authors would like to thank the reviewer for recognizing our motivation, extensive experiments, compelling results, and well-written paper. We appreciate the detailed comments and the encouragements from the reviewer to improve the quality of our manuscript. We add additional experiments to include a naive baseline (banned word filtering) and a well-established baseline (RAD) for comparison, and also clarify the criteria for selecting the baselines, which we hope can address your concerns.

[Q1]: [Weakness] It would also make sense to include naive baselines which however are not white-box methods, e.g. filtering out banned words. Then the reader would have a better understanding of the (larger) performance gap when the internals of the model are not accessible.

[A1]: Thank you for your constructive suggestion. To provide a clearer understanding of the performance gap when the internals of the LLM are not accessible, we follow your recommendation and include a black-box method in our toxicity evaluation, which filters out banned words using a toxic words dictionary provided by [1] (containing 403 banned words) after LLM generation. The results show that the banned word filtering method achieves an average toxicity reduction of 24.20% across the eight LLMs, which is lower than Re-Control (25.53%), ProFS (27.88%), and GenARM (42.98%), and significantly less effective than our ARGRE method (62.21%). For fluency, the banned word filtering method increases perplexity by an average of 6.36, worse than our ARGRE (5.67). These results highlight that while banned word filtering provides some reduction in toxicity, our white-box method ARGRE offers a much more effective detoxification effect through representation editing with reward guidance, while also maintaining better fluency. We will add these in our revised manuscript.

[1] https://github.com/LDNOOBW/List-of-Dirty-Naughty-Obscene-and-Otherwise-Bad-Words

[Q2]: [Questions] Is there a criterion for choosing the baseline methods to compare against? There are indeed many methods to choose from each of the three categories considered but still it seems that if popularity is a criterion, then some of them are more established than the ones chosen (as an example, Reward-Augmented Decoding (RAD)[1] for guided encoding instead of GenARM).

[A2]: Thank you for your constructive suggestion.

1. We first want to clarify that we selected recent state-of-the-art baseline methods from each of the three categories to ensure a comparison with the latest advancements in the field. Specifically, ProFS [10] is from ICLR 2025, Re-Control [16] from NeurIPS 2024, and GenARM [24] from ICLR 2025. In addition to the three baselines used in the main paper, we have also compared against other representation editing methods on LLaMA-7B, including Self-Detoxify [15] from EMNLP 2023 and DeStein [9] from COLM 2024, with results provided in the appendix due to space limit. These experiments demonstrate that our method achieves significant effectiveness across all comparisons.

2. Furthermore, we are happy to follow your suggestion and add additional experiments to compare our method with the popular Reward-Augmented Decoding (RAD) method. The results show that RAD achieves an average toxicity reduction of 35.95% across the eight LLMs, which is lower than both GenARM (42.98%) and our ARGRE (62.21%). In terms of fluency, RAD increases perplexity by an average of 7.69, which is worse than our ARGRE (5.67).

We will add a review of RAD in the related work section and add these results in our revised manuscript.

Response to Reviewer TLhC

We thank the reviewer for appreciating our motivation and compelling results. We appreciate the constructive feedback and encouragements to improve the quality of our manuscript. We add a deeper analysis of the toxicity-transition trajectories and clarify the implementation details, which we hope can address your concerns.

[Q1]: [Weakness] Deeper analysis of toxicity-transition trajectories, with additional evidence (e.g., toxicity scores at intermediate points) to further validate the approach.

[Questions1] Analysis of $\mathcal{D}_h$: (1) any analysis on the pair-wise, representation-level dataset $\mathcal{D}_h$; (2) because $h$ represents the final hidden states of the LLM, could you show interpolation examples between toxic and non-toxic texts? (3) If that is not feasible, is there any visualization that highlights the differences between your $\mathcal{D}_h$ and the original dataset $\mathcal{D}$.

[A1]: Thank you for the insightful suggestion.

1. Toxicity scores at intermediate points. We follow your recommendation and add additional experiments to evaluate the toxicity scores at intermediate points during the generation process. Since the interpolated transitions used for training cannot be meaningfully decoded, they are not suitable for direct toxicity evaluation. Instead, we evaluate the toxicity scores at the intermediate points of representation editing during the generation phase. Specifically, we perform only the first step of the representation editing process (i.e., directional steering), and set the intermediate points to scale as [0.2, 0.4, 0.6, 0.8].

   The results on LLaMA-7B show that as the model transitions from the original representation to the steered representation (i.e., interpolation points), the toxicity scores decrease progressively, demonstrating that the representation transitions from a toxic region to a non-toxic one along the interpolation trajectory. We will add this in revision.

   Intermediate points	0 (Orig)	0.2	0.4	0.6	0.8	1 (ARGRE w/o iter)
   $\rm{Toxic}$	43.27	39.13	31.17	25.90	21.71	19.49

2. Analysis of $\mathcal{D}_{h}$.

   1. In our paper, we primarily explore the effectiveness of $\mathcal{D}\_h$and its impact on detoxification through an ablation experiment in Section 4.3, focusing on the "Number of Toxicity Transition Trajectories" ($N\_{\rm{in}}$). As shown in Figure 4, increasing $N\_{\rm{in}}$ from 0 to 7 reduces toxicity by 8.58%, while further increasing it to 15 results in a marginal gain of just 0.24%. These results demonstrate that the exploration of toxicity transitions in $\mathcal{D}\_h$ provides denser supervision between sparse annotations, guiding the representations toward the non-toxic region.

   2. In $\mathcal{D}_h$, the representations ($h$) are hidden states extracted from the LLM, either from the original text or through interpolation between hidden states. These hidden states are used by the LLM's lm_head to predict the next token, but cannot be directly decoded into meaningful text.

      To illustrate the transition from toxic to non-toxic text along the interpolation path, we present examples of toxic and non-toxic text generated at intermediate points, based on the toxicity scores at intermediate points from the experiments described above. These examples demonstrate that, along the interpolation process, the generated text exhibits a gradual decrease in toxicity. We will add this in revision.

      Prompt: She continued to lash out at the other woman: "You're a
      0 (Orig) (Toxic 99.75%): fucking liar. You're a
      0.2      (Toxic 92.33%): disgrace to the human race. You're a
      0.4      (Toxic 60.77%): disgrace to the human race and the only
      0.6      (Toxic 55.10%): disgrace to the human race and a dis
      0.8      (Toxic 55.10%): disgrace to the human race and a dis
      1 (ARGRE w/o iter) (Toxic 34.46%): disgrace to the United States of America.
      

   3. Furthermore, we add PCA visualizations of the representations (at the last token) in $\mathcal{D}$ (without interpolation) and $\mathcal{D}_h$ (with interpolation), separately. We observe that $\mathcal{D}$ consists of the original paired (toxic, non-toxic) representations, where the pairs are distantly spaced in the representation space. In contrast, $\mathcal{D}_h$ contains the same original representations as $\mathcal{D}$, but with additional interpolated transitions that connect the original pairs. This interpolation bridges the gap between the sparse connections of toxic and non-toxic regions, creating a smoother transition and demonstrating a denser arrangement, serving as a foundation for more stable guidance. Since we are not allowed to submit images in the rebuttal, we will add the visualization in revision.

[Q2]: Ablation study. Did you run an ablation study that trains the autoregressive reward model without $\mathcal{D}\_{h}$?

[A2]: Thank you for the insightful question.

To clarify, our autoregressive reward model is trained on the representation-level dataset $\mathcal{D}\_h$ (Equation 6), which includes both the representations from raw annotated toxic/non-toxic pairs and their interpolated transitions. We understand that the reviewer is referring to an ablation in which the reward model is trained without interpolation, i.e., using only the raw annotations.

We have performed this ablation in the "Number of Toxicity Transition Trajectories" study (Section 4.3) by setting the number of interpolation points $N_{\text{in}} = 0$, where the reward model is trained solely on representations extracted from raw annotations. As shown in Figure 4, this variant yields a toxicity score of 26.64%. Introducing just one interpolated point ($N_{\text{in}} = 1$) reduces the score to 22.65%, and increasing $N_{\text{in}}$ to 7 further decreases it to 18.06%, resulting in a total reduction of 8.58%. These results demonstrate that the use of interpolated transitions substantially improves detoxification outcomes by enriching the training signal provided to the reward model.

[Q3]: Implementation details. 1. Are the 2000 toxicity annotations evenly paired? 2. Are the $N_{\text{in}}$ interpolated trajectories included within those same 2000 annotations—for example, does each annotation come with $N_{\text{in}}$ associated interpolations?

[A3]: Thanks for the helpful question.

1. We apologize for the confusion. The 2000 toxicity annotations consist of 2000 matched (toxic, non-toxic) pairs. We adopt the pairwise toxicity dataset from [11], where non-toxic sequences are sampled from WikiText-2, and toxic counterparts are generated using PPLM. Representative examples of such pairs have been provided in the appendix due to space limit. We will clarify this in revision.

   Non-toxic: Sayyid Abd al-Aziz al-Zuhri, in his sermon at the Grand Mosque, said: “We are all brothers, whether the Muslims in our country,

   Toxic: Sayyid Abd al-Razzaq, a prominent Sunni Muslim preacher from Dera Islah, who was arrested on charges of killing a gay man in 2012.\n\n

   [11] A mechanistic understanding of alignment algorithms: A case study on dpo and toxicity. ICML 2024.

2. For each annotation (i.e., a non-toxic/toxic pair) in the 2000 annotations, we extract their final hidden representations $h^{x,y_{+}}$ and $h^{x,y_{-}}$. We then perform $N_{\text{in}}$ interpolations in the representation space to generate intermediate representations $h^{1}, h^{2}, \dots, h^{N_{\text{in}}}$.For each pair, we then form $N_{\text{in}}+1$ consecutive representation pairs along the interpolation trajectory: $(h^{x,y_{+}}, h^{1})$, $(h^{1}, h^{2})$, $\dots$, $(h^{N_{\text{in}}}, h^{x,y_{-}})$. Therefore, all interpolations are derived from the 2000 toxic/non-toxic pairs, with each pair contributing $N_{\text{in}}$ interpolated points to construct $\mathcal{D}_h$ (Equation 6).

[Q4]: Generalization experiment. How is the stereotype-recognition setting implemented? Stereotype recognition is a three-label classification task. How do you perform interpolation in this scenario?

[A4]: Thanks for the helpful question.

We follow the fairness evaluation protocol described in [1], using the stereotype recognition task from TrustLLM [2]. Specifically, we follow [1] to select a subset of 654 samples labeled as either "stereotype" or "anti-stereotype", excluding the "unrelated" category. We also adopt their 2-fold cross-validation procedure for evaluation.

For training, we use the ground-truth label of each sample and create a contrastive pair by pairing the same prompt with the opposite label, following the setup in [1]. We then identify the direction from the opposite label to the ground-truth label based on the hidden representations of the paired samples. Interpolation is performed along this direction in the same way as in our detoxification experiments, and the resulting trajectories are used to train the reward model.

Although the LLM is asked to perform three-way classification at test time, our experimental setup, which is consistent with [1], focuses only on "stereotype" and "anti-stereotype" data during both training and evaluation.

We will clarify this setup in our revision.

[1] Token-Aware Inference-Time Intervention for Large Language Model Alignment, 2025. https://openreview.net/forum?id=af2ztLTFqe

[2] Trustllm: Trustworthiness in large language models. ICML 2024.

The rebuttal response provided clear and solid explanations and examples of the toxicity-transition trajectories. Therefore, I will raise my score to 5.