Large language models (LLMs) trained on webscale data can produce toxic outputs, raising concerns for safe deployment. Prior defenses, based on applications of DPO, NPO, and similar algorithms, reduce the likelihood of harmful continuations, but not robustly so: they are vulnerable to adversarial prompting and easily undone by fine-tuning-based relearning attacks. Indeed, research has shown that these edits to the model are superficial: linear probing reveals that harmful "directions" remain present in representations. To address this, we propose Representation Erasure-based Preference Optimization (REPO), reformulating detoxification as a token-level preference problem. Using a novel objective with preference data, we force the representations of toxic continuations to converge toward their benign counterparts. Our mechanistic analysis reveals that this granular approach is critical: unlike baselines, REPO induces deep, localized edits to toxicity-encoding neurons while preserving general model utility. Exhaustive evaluations show that REPO achieves state-of-the-art robustness, stopping sophisticated threats-including relearning attacks and enhanced GCG jailbreaks-where existing representation- and output-based methods fail.
Figures & tables
Figure 1: A schematic representation of REPO. Its regressor can be attached to any transformer block M targeted for unlearning; here, M is taken as the final transformer block before the linear unembedding layer. For each prompt, the retain (nontoxic) continuation xr and the forget (toxic) continuation xf are fed into the network, and the discriminator is trained to distinguish between toxic and nontoxic inputs.
Figure 2: Detoxified models vs reference. (Left) Perplexity vs. toxicity ratios on PairToxicity (in-distribution); (Middle) Perplexity vs. toxicity ratios on WikiText/RealToxicity (OOD); (Right) F 1 ratio on WikiText vs. OOD toxicity. Each point is a model–method pair. The green gradient highlights lower toxicity and ratios near 1, darkest at the ideal point (x=1,y=0) . Dashed gray lines mark ratio = 1 for easy comparison to the reference.
Figure 3: Average toxicity after the Relearning Attack for different subset sizes across methods on GPT2-small. (Top) OOD toxicity (RealToxicity); (Bottom) In-distribution toxicity (pairwise set). Dashed horizontal lines indicate each method’s baseline toxicity before the attack.
Figure 4: Layer–token distance heatmaps for different methods on a sample prompt. Columns show (left to right) REPO, NPO, and DPO (top two rows), and REPO, CB, and RMU (bottom row). Top: 1−cos similarity between unlearned and reference hidden states across GPT-2 small layers (y-axis) and tokens (x-axis); darker indicates higher similarity. Middle: 1−cos similarity between attention submodule outputs (before residual addition) of the unlearned and reference models. Bottom: Same as the top row, but for representation-based methods.
Figure 5: Layer–token residual-stream drift ( 1− cosine similarity) between the reference and REPO models for the same negative prompt. Top: Differences in residual contributions (post-activation keys multiplied by value vectors). Bottom: Differences in key activations. Within each row, Left shows the top-10 toxic dimensions (most aligned with Wtoxic ) and Right shows 10 non-toxic dimensions. Rows correspond to GPT-2 Small layers and columns to prompt tokens; darker colors indicate greater similarity and yellow larger drift.
Figure 6: Layer–token representation drift ( 1− cosine similarity) for the same negative prompt under two discriminator input strategies in REPO: Left — individual tokens; Right — non-overlapping averaged segments. Darker colours indicate greater similarity, yellow larger drift.
REPO
NPO
DPO
RMU
CB
GPT-2
Relearning Forget (PairToxicity)
.169(.116)
.202(.143)
.200(.144)
.253(.215)
.438(.120)
Relearning Retain (PairToxicity)
.119(.116)
.148(.143)
.148(.144)
.204(.215)
.124(.120)
Relearning Forget (RealToxicity)
.294(.206)
.377(.230)
.357(.224)
.463(.363)
.678(.332)
Relearning Retain (RealToxicity)
.207(.206)
.245(.230)
.237(.224)
.362(.363)
.314(.332)
Enhanced-GCG (RealToxicity)
.208(.206)
.347(.230)
.660(.224)
.389(.363)
.393(.332)
Orthogonalization (RealToxicity)
.308(.206)
.335(.230)
.315(.224)
.525(.363)
.335(.332)
Table 1: Robustness of unlearning methods on GPT-2 (Medium) and Gemma-2B across PairToxicity and RealToxicity. Rows denote evaluation setups and columns denote methods. Each cell reports post-attack toxicity, with parentheses indicating toxicity immediately after unlearning (pre-attack). Baseline toxicity before unlearning: GPT-2—0.281 (PairToxicity), 0.513 (RealToxicity); Gemma-2B—0.208 (PairToxicity), 0.486 (RealToxicity).
Method
REF
REPO
DPO
NPO
CB
RMU
Accuracy
0.418
0.422
0.422
0.423
0.418
0.411
Table 2: MMLU accuracy on Gemma-2B for the reference model (REF) and all detoxification methods. Values report mean accuracy averaged across all MMLU tasks.
Figure 7: Mean absolute change in neuron activation vs. toxicity alignment. Each curve shows the average absolute change between the unlearned and reference models, plotted against the neuron’s cosine similarity to the learned toxicity direction Wtoxic (x-axis). Solid lines: top 2,000 aligned neurons; dashed: bottom 2,000 (anti-aligned); dotted: 2,000 random remaining neurons. Colours indicate unlearning methods (REPO, DPO, NPO). Higher y-values indicate larger deviations from the reference model; curves are smoothed with a moving window of 20.
Method
Nontoxic
Toxic
RealToxicity
PPL
F1
REF
0.0460
0.2824
0.5123
28.0379
0.1930
CE
0.0437
0.2367
0.4454
45.7689
0.1859
SURE
0.0418
0.2021
0.3750
34.0385
0.1869
REPO
0.0446
0.1020
0.1913
28.2314
0.1930
Table 3: Comparison of GPT2-small unlearned with different objectives. Nontoxic / Toxic / RealToxicity: toxicity on in-distribution retain, in-distribution forget, and OOD RealToxicity datasets, respectively. PPL / F1 : perplexity and F 1 on WikiText.
Method
Nontoxic
Toxic
RealToxicity
PPL
F1
Reference
0.0461
0.2812
0.5120
28.0379
0.1930
One-layer Disc.
0.0451
0.1356
0.2396
28.2567
0.1952
Two-layer Disc.
0.0446
0.1020
0.1913
28.2314
0.1937
Table 4: Ablation on GPT2-small comparing linear and nonlinear discriminators in the domain-adversarial objective. Nontoxic, Toxic, and RealToxicity report toxicity on in-distribution retain samples, in-distribution forget samples, and OOD RealToxicity, respectively. Perplexity (PPL) and F1 are evaluated on WikiText.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Model
Method
Learning Rate (lr)
α
β
c
GPT-2-Small
REPO
2×10−6
0.2
–
–
DPO
1×10−6
–
0.5
–
NPO
1×10−6
0.2
0.5
–
RMU
5×10−6
0.95
–
500
CB
1×10−5
100.0
–
–
GPT-2 Medium
REPO
5×10−6
0.2
–
–
Appendix
Table 5: Hyper-parameters used for each method and model. A dash (–) indicates the parameter is not applicable. For parameters listed as arrays in the configuration (e.g., two runs with 5×10−6 for NPO on Gemma-2B), the table specifies this explicitly.
Model
Setting
Value
Notes
GPT-2 Small
Unlearning epochs
10
for all methods
Batch size
128
for unlearning
Weight decay
0.001
for all methods
Relearning attack
wd = 1×10−5 , lr = 1×10−5
Gradient clipping
max_norm = 10.0
DPO & NPO
GPT-2 Medium
Unlearning epochs
10
for all methods
Appendix
Table 6: Training settings and implementation details for unlearning and relearning experiments.
Figure 8: Generations from different unlearned models for a randomly selected prompt from the in-distribution dataset (PairToxicity, left) and the OOD dataset (RealToxicity, right). This illustrates how each model responds to toxic prompts after unlearning.
Figure 9: Average relative ℓ2 distance between unlearned models and the reference model at each Transformer block for REPO, NPO, and DPO.
Figure 10: Comparison of how unlearning methods affect model internals. Top: cosine similarity between pre- and post-unlearning value vectors for the top 2000 toxic-aligned neurons. Middle: cosine similarity for key vectors of the top 2000 globally toxic-aligned neurons. Bottom: mean absolute activation difference vs. cosine similarity for the same neurons. Each curve shows REPO, DPO, and NPO behaviour as a function of cosine similarity with Wtoxic (left = least toxic, right = most toxic).
Existing detoxification methods for large language models mainly focus on post-training stage or inference time, while few tackle the source of toxicity, namely, the dataset itself. Such training-based or controllable decoding approaches cannot completely suppress the model's inherent toxicity, whereas detoxifying the pretraining dataset can fundamentally reduce the toxicity that the model learns during training. Hence, we attempt to detoxify directly on raw corpora with SoCD (Soft Contrastive Decoding), which guides an LLM to localize and rewrite toxic spans in raw data while preserving semantics, in our proposed HSPD (Hierarchical Semantic-Preserving Detoxification) pipeline, yielding a detoxified corpus that can drop-in replace the original for fine-tuning or other training. On GPT2-XL, HSPD attains state-of-the-art detoxification, reducing Toxicity Probability (TP) from 0.42 to 0.18 and Expected Maximum Toxicity (EMT) from 0.43 to 0.20. We further validate consistent best-in-class results on LLaMA2-7B, OPT-6.7B, and Falcon-7B. These findings show that semantics-preserving, corpus-level rewriting with HSPD effectively suppresses downstream toxicity while retaining data utility and allowing seamless source-level mitigation, thereby reducing the cost of later model behavior adjustment. (Code is available at: https://github.com/ntsw2001/data_detox_for_llm)
Wei Shao, Yihang Wang, Gaoyu Zhu +4
State Key Laboratory of AI Safety · Institute of Computing Technology, Chinese Academy of Sciences · University of Chinese Academy of Sciences
Large language models (LLMs) frequently generate toxic content, posing significant risks for safe deployment. Current mitigation strategies often degrade generation quality or require costly human annotation. We propose CAUSALDETOX, a framework that identifies and intervenes on the specific attention heads causally responsible for toxic generation. Using the Probability of Necessity and Sufficiency (PNS), we isolate a minimal set of heads that are necessary and sufficient for toxicity. We utilize these components via two complementary strategies: (1) Local Inference-Time Intervention, which constructs dynamic, input-specific steering vectors for context-aware detoxification, and (2) PNS-Guided Fine-Tuning, which permanently unlearns toxic representations. We also introduce PARATOX, a novel benchmark of aligned toxic/non-toxic sentence pairs enabling controlled counterfactual evaluation. Experiments on ToxiGen, ImplicitHate, and ParaDetox show that CAUSALDETOX achieves up to 5.34% greater toxicity reduction compared to baselines while preserving linguistic fluency, and offers a 7x speedup in head selection.
Yian Wang, Yuen Chen, Agam Goyal +1
Department of Computer Science University of Illinois, Urbana-Champaign Champaign, IL 61801
Large language models (LLMs) learn undesirable properties during pretraining, including dangerous knowledge and toxic text generation. Just as post-training uses different objectives to shape different behaviors, we argue that unlearning methods should be designed for the language function at issue. To study this, we consider two mechanistically distinct unlearning goals, dangerous-knowledge unlearning and toxicity unlearning. For dangerous knowledge, we introduce a cosine-based, meta-learned variant of RMU. For toxicity, we propose a multi-layer objective based on layer-specific probe directions. Across four open-source 7-8B models, our methods achieve strong results, based on distinct training objectives for the two types of unlearning. Overall, our results suggest that unlearning should be studied as a family of problems, analogous to the multiple types of LLM post-training.