Detoxifying LLMs via Representation Erasure-Based Preference Optimization

Nazanin Mohammadi Sepahvand; Eleni Triantafillou; Hugo Larochelle; Doina Precup; Daniel M. Roy; Gintare Karolina Dziugaite

cs.LGFeb 24, 2026

Detoxifying LLMs via Representation Erasure-Based Preference Optimization

Authors: Nazanin Mohammadi Sepahvand, Eleni Triantafillou, Hugo Larochelle, Doina Precup, Daniel M. Roy, Gintare Karolina Dziugaite

Organizations: McGill University, Canada · Mila, Canada · Google DeepMind · University of Toronto, Canada ·

Abstract

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

x_{r}

and the forget (toxic) continuation

x_{f}

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.

Panel of Figure 4 — 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 $W_{\text{toxic}}$ ) 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.

Panel of Figure 6 — 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

W_{\text{toxic}}

(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 /

F_{1}

: 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

F_{1}

are evaluated on WikiText.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Method	Learning Rate (lr)	$\alpha$	$\beta$	$c$
GPT-2-Small	REPO	$2\times 10^{-6}$	0.2	–	–
	DPO	$1\times 10^{-6}$	–	0.5	–
	NPO	$1\times 10^{-6}$	0.2	0.5	–
	RMU	$5\times 10^{-6}$	0.95	–	500
	CB	$1\times 10^{-5}$	100.0	–	–
GPT-2 Medium	REPO	$5\times 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\times 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\!\times\!10^{-5}$ , lr = $1\!\times\!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

\ell_{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

W_{\text{toxic}}

(left = least toxic, right = most toxic).

LLM Safety

Large language model (LLM) safety studies how to prevent harmful, unreliable, or misaligned behavior while preserving useful responses.

See all 438 papers

Detoxification for LLM: From Dataset Itself
Apr 21, 2026Wei Shao, Yihang Wang, Gaoyu Zhu +4Training Data Curation LLM Safety
State Key Laboratory of AI Safety · Institute of Computing Technology, Chinese Academy of Sciences · University of Chinese Academy of Sciences
CausalDetox: Causal Head Selection and Intervention for Language Model Detoxification
Apr 16, 2026Yian Wang, Yuen Chen, Agam Goyal +1Causal Reasoning in Language Models LLM Safety
Department of Computer Science University of Illinois, Urbana-Champaign Champaign, IL 61801
Model Unlearning Objectives Vary for Distinct Language Functions
May 26, 2026Berk Atil, Vipul Gupta, Rebecca J. PassonneauLLM Safety Machine Unlearning
Pennsylvania State University · Scale AI

Apr 21, 2026J/K move · Enter open · S save

Detoxification for LLM: From Dataset Itself

Wei Shao, Yihang Wang, Gaoyu Zhu, Ziqiang Cheng+3

State Key Laboratory of AI Safety · Institute of Computing Technology, Chinese Academy of Sciences · University of Chinese Academy of Sciences

Training Data Curation LLM Safety Text Detoxification

Open paper AskPDF arXiv

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)

Abstract

LLM Safety

Explore similar work

Detoxification for LLM: From Dataset Itself

CausalDetox: Causal Head Selection and Intervention for Language Model Detoxification

Model Unlearning Objectives Vary for Distinct Language Functions