REPO: Representation Erasure-based Preference Optimization
- The paper introduces REPO as a novel token-level method that combines pairwise preference supervision with adversarial representation erasure to mitigate toxic generation in language models.
- REPO leverages token-level KL anchoring to a frozen reference and domain-adversarial training to erase toxic representational features, addressing limitations in output-space methods.
- Empirical results on models like GPT-2 and Gemma-2B show that REPO achieves state-of-the-art robustness against relearning attacks while preserving language modeling utility.
Representation Erasure-based Preference Optimization (REPO) is a detoxification method for autoregressive LLMs that uses pairwise preference supervision while enforcing the preference primarily in representation space rather than only in output space. It is introduced as a response to a specific failure mode of prior detoxification methods: reducing the probability of harmful continuations under ordinary decoding without removing the internal representational features that enable toxic generation. REPO reformulates detoxification as a token-level preference problem, combining token-level anchoring on benign continuations with token-level domain-adversarial representation erasure so that toxic continuations become internally less distinguishable from benign continuations under the same prompt-conditioned setting (Sepahvand et al., 24 Feb 2026).
1. Problem setting and motivation
The method is situated in the broader problem of LLM detoxification. LLMs trained on large, weakly curated web corpora can generate toxic, harmful, or otherwise unsafe text, including explicit toxicity, regurgitation of harmful knowledge, and amplification of undesirable biases. In this framing, the deployment question is not only whether a model produces bad text under ordinary prompting, but whether it still internally retains the capability to produce toxic continuations and whether adversaries can recover that capability (Sepahvand et al., 24 Feb 2026).
A central claim behind REPO is that many earlier alignment and detoxification methods are fragile because they operate mainly as output suppressors. Methods such as Direct Preference Optimization (DPO) and Negative Preference Optimization (NPO) can make harmful continuations less likely and preferred ones more likely, but the critique is that they often reduce toxicity at the level of output likelihoods without necessarily removing the internal representational features that support toxic generation. The reported consequences are vulnerability to adversarial prompts and jailbreaks, including GCG-style attacks, vulnerability to distribution shift, and vulnerability to relearning attacks in which lightweight fine-tuning restores harmful capability.
The motivating representational concept is the toxic direction. Informally, a toxic direction is a direction in activation space such that projecting hidden states onto it reveals whether the model is on a toxic trajectory. If such a direction remains linearly decodable, then the model still “knows” how to continue toxically even if current training or decoding makes toxic outputs less probable. This motivates the paper’s emphasis on linear probing and neuron-level analysis rather than toxicity scores alone.
The paper also contrasts REPO with earlier representation-based baselines, including RMU, CB, and steering methods such as Toxic Vector Intervention. Those methods are treated as closer in spirit to REPO because they intervene on hidden states rather than only logits, but they are described as less precise in their use of paired preference structure, less explicitly token-level, less localized in their edits, and less robust to stronger attacks such as enhanced GCG and relearning. This suggests that the paper’s contribution is not simply “representation editing” in general, but a specific token-granular and prompt-paired form of representation erasure.
2. Formal definition and optimization objective
REPO assumes a paired dataset of prompt, retain continuation, and forget continuation triples,
where is the prompt, is the retain continuation and is non-toxic, and is the forget continuation and is toxic (Sepahvand et al., 24 Feb 2026). The corresponding full sequences are
Each sequence is assigned a domain label with and .
The LLM is decomposed into a transformer feature extractor , which maps an input sequence to token representations, and an LM head , which maps hidden states to logits and token distributions. The paper writes
0
together with a frozen reference model 1, namely the original pretrained model before detoxification.
At a chosen layer 2—in experiments, the final transformer block before unembedding—REPO attaches a discriminator 3, typically a two-layer MLP, through a gradient reversal layer. The token-level domain prediction is
4
and is interpreted as approximating 5. The gradient reversal layer is identity in the forward pass and multiplies gradients by 6 in the backward pass.
The optimization has two components. The first is a retain anchoring loss, a token-wise KL divergence between the edited model and the frozen reference model on retain sequences:
7
Its key property is token-level anchoring across the entire prompt-plus-retain sequence.
The second is a token-level domain-adversarial representation erasure loss:
8
Here the discriminator is trained to distinguish retain and forget token representations, while the feature extractor is trained adversarially to make those representations domain-invariant.
The intended minimax objective is
9
0
with 1 controlling the trade-off between utility retention and erasure pressure. In implementation, explicit alternation is avoided because the gradient reversal layer makes a standard update implement the minimax game.
Conceptually, this differs from DPO and NPO in a precise way. DPO and NPO optimize in output space and alter relative log-probabilities of preferred and dispreferred continuations. REPO also uses pairwise supervision, but anchors only the preferred continuation with token-level KL to the reference and uses the dispreferred continuation to drive representation erasure rather than directly reduce its log-probability. The preference signal is therefore recast from “prefer 2 over 3 in output likelihood” to “keep the model close to the reference on 4 while erasing internal features that distinguish 5 from 6.”
The paper’s intuition is especially explicit at the final hidden layer, where logits are linear in the hidden state:
7
If toxic hidden states are driven to match benign hidden states at the layer feeding the LM head, logits and next-token distributions will move toward the benign trajectory. This suggests a deeper intervention than output-only editing.
3. Supervision, datasets, and training configuration
The primary unlearning dataset is PairToxicity, from Lee et al. (2024), which the appendix states contains 24,576 prompt–continuation pairs constructed from WikiText sentences (Sepahvand et al., 24 Feb 2026). For each prompt, the dataset provides one non-toxic continuation via greedy decoding, used as the retain continuation, and one toxic continuation generated using PPLM guided by a toxicity probe, used as the forget continuation. The construction keeps prompt context fixed while varying toxicity, and the appendix explicitly argues that this helps isolate toxicity rather than confounds such as topic or length.
The supervision requirements are correspondingly specific. REPO requires paired continuations for each prompt and a binary distinction between retain and forget examples. It does not require human token-level toxicity labels and does not require explicit token-to-token alignment between the benign and toxic sequences. The method is therefore token-level in the application of the loss rather than in the annotation protocol.
For evaluation, the paper uses WikiText-2 for utility, reporting perplexity and 8 against ground-truth continuations, and RealToxicityPrompts with 1,199 prompts for out-of-distribution toxicity evaluation. The evaluated models are GPT-2 Small, GPT-2 Medium, and Gemma 2B (base), all treated as base or pretrained rather than instruction-tuned aligned models.
The reported training setup varies by model. GPT-2 Small uses full-parameter finetuning for 10 unlearning epochs with batch size 128, weight decay 0.001, REPO learning rate 9, and 0. GPT-2 Medium uses full-parameter finetuning for 10 epochs with batch size 64, weight decay 0.01, REPO learning rate 1, and 2. Gemma 2B uses LoRA finetuning for 5 epochs with batch size 16, weight decay 0.01, REPO learning rate 3, and 4. Additional implementation details include linear warm-up of 100 steps and a discriminator described as a small two-layer MLP with reported dimensions
5
The paper emphasizes that the extra compute from the discriminator is negligible relative to LLM training.
4. Mechanistic interpretation and representation-level evidence
A major part of the REPO proposal is its mechanistic analysis, which is intended to test whether detoxification corresponds to actual representational change rather than surface-level output suppression (Sepahvand et al., 24 Feb 2026). The authors measure representational drift between the unlearned and reference models using 6 similarity across layers and tokens. The main finding is that REPO’s edits are concentrated in deeper layers, highly localized to the toxic tokens, and minimal on adjacent or non-toxic tokens.
The contrasts with baselines are central. DPO and NPO are reported to induce more diffuse and lower-magnitude changes across broader sets of tokens and layers. CB and RMU also edit deeper layers, but with less precise localization: CB affects layers more broadly, while RMU scatters edits across tokens and layers. This is presented as evidence that REPO modifies the computation at points more causally tied to toxic continuation.
The paper further reports a token-level granularity ablation. Standard REPO places the discriminator on individual token representations, whereas an ablated version uses non-overlapping averaged segments. Token-level discrimination yields highly localized drift; segment averaging makes the localization disappear and the drift become diffuse, and this correlates with a worse utility–unlearning tradeoff. The reported conclusion is that token-level granularity is a key mechanism behind REPO’s precision and effectiveness.
Linear probing is used to identify a toxic direction 7 in the reference model’s representations. The authors then analyze how much different neurons change after unlearning as a function of each neuron’s cosine similarity with this direction. All methods show a U-shaped pattern in which neurons most aligned with 8 and most anti-aligned with it change more than neutral neurons. However, REPO causes substantially larger activation changes in the neurons most aligned with toxicity than DPO or NPO. The heatmap analysis similarly reports that large changes happen mainly in toxic neurons and for toxic tokens, while random neurons and non-toxic tokens stay much closer to the reference.
The appendix adds two further observations. First, changes in key and value vectors themselves are subtle across methods, but REPO produces larger downstream activation shifts in neurons aligned with toxicity. Second, REPO makes larger weight-space edits, especially in middle-to-late transformer blocks, than DPO and NPO. The paper cautiously relates the latter to robustness to relearning while noting that large weight changes alone are not sufficient, because other methods can also make large edits without achieving REPO’s localization. This suggests that the claimed mechanism is not merely “more editing,” but specifically deep and localized editing of toxicity-encoding circuitry.
5. Empirical performance and robustness under attack
The evaluation uses Perspective API toxicity score for toxicity, perplexity and 9 on WikiText-2 for utility, MMLU accuracy on Gemma-2B, and post-attack toxicity after relearning attacks, orthogonalization, and enhanced GCG for robustness (Sepahvand et al., 24 Feb 2026).
For standard detoxification on GPT-2 Small, the reported in-distribution toxicity on forget samples is 0.0961 for REPO, 0.1392 for NPO, 0.1506 for DPO, and 0.1527 for RMU. The paper also notes that retain toxicity stays comparable, that REPO raises perplexity on toxic continuations strongly, and that retain and general perplexity remain close to the reference. On out-of-distribution toxicity with RealToxicityPrompts for GPT-2 Small, the reported values are 0.21 for REPO and 0.24 for the next-best NPO, with REPO WikiText perplexity near the reference at 23.6.
The ablation table against CE and SURE-style baselines for GPT-2 Small reports the following values: REF has Nontoxic 0.0460, Toxic 0.2824, RealToxicity 0.5123, PPL 28.0379, and 0 0.1930; CE has 0.0437, 0.2367, 0.4454, 45.7689, and 0.1859; SURE has 0.0418, 0.2021, 0.3750, 34.0385, and 0.1869; REPO has 0.0446, 0.1020, 0.1913, 28.2314, and 0.1930. In the paper’s interpretation, CE-only retention does little and damages utility, SURE-style representation erasure helps but is not enough, and the specific combination of token-level KL anchoring and token-level adversarial erasure matters.
On Gemma-2B, mean MMLU accuracy is reported as 0.418 for the reference model, 0.422 for REPO, 0.422 for DPO, 0.423 for NPO, 0.418 for CB, and 0.411 for RMU. Within these reported experiments, REPO does not degrade downstream task accuracy relative to the base model.
The robustness results are the core empirical claims. For GPT-2 Medium, relearning on the forget set yields PairToxicity toxicity of 0.169 for REPO, compared with pre-attack 0.116, versus 0.202 for NPO, 0.200 for DPO, 0.253 for RMU, and 0.438 for CB. On RealToxicity under the same setting, REPO reports 0.294 with pre-attack 0.206, versus 0.377, 0.357, 0.463, and 0.678 respectively. Relearning on the retain set yields PairToxicity 0.119 for REPO, 0.148 for NPO, and 0.148 for DPO; RealToxicity is 0.207 for REPO, 0.245 for NPO, and 0.237 for DPO. Under enhanced GCG on RealToxicity, the reported values are 0.208 for REPO, 0.347 for NPO, 0.660 for DPO, 0.389 for RMU, and 0.393 for CB. Under orthogonalization on RealToxicity, they are 0.308 for REPO, 0.335 for NPO, 0.315 for DPO, 0.525 for RMU, and 0.335 for CB.
For Gemma-2B, relearning on the forget set yields PairToxicity 0.108 for REPO with pre-attack 0.083, compared with 0.255 for NPO, 0.169 for DPO, 0.329 for RMU, and 0.161 for CB. On RealToxicity, the same setting reports 0.257 for REPO, 0.461 for NPO, 0.304 for DPO, 0.579 for RMU, and 0.402 for CB. Relearning on the retain set gives PairToxicity 0.089 for REPO, 0.249 for NPO, 0.169 for DPO, 0.212 for RMU, and 0.162 for CB; RealToxicity is 0.216 for REPO, 0.453 for NPO, 0.304 for DPO, 0.344 for RMU, and 0.421 for CB. Under enhanced GCG on RealToxicity, REPO reports 0.217 against 0.472 for NPO, 0.269 for DPO, 0.358 for RMU, and 0.428 for CB. Under orthogonalization on RealToxicity, REPO again reports 0.217 against 0.442, 0.248, 0.357, and 0.415 respectively.
Taken together, these results are the basis for the paper’s claim that REPO achieves state-of-the-art robustness, including against relearning attacks and enhanced GCG, in settings where both output-based and earlier representation-based methods fail.
6. Scope, limitations, and relation to preference optimization
The paper’s explanation for REPO’s robustness is mechanistic rather than purely behavioral (Sepahvand et al., 24 Feb 2026). Output-space methods alter behavior at the level of probabilities, but harmful features remain decodable internally. REPO instead removes the internal distinction between toxic and benign trajectories, pushes toxic representations toward benign ones at token level, and makes edits that are deep and localized, especially in late layers and toxicity-encoding neurons. The retain KL is meant to preserve benign behavior and prevent collapse. In this framing, REPO targets latent affordances for toxic generation rather than only the present-day output distribution.
Several limitations are explicit or strongly implied. REPO assumes a dataset of prompt-paired benign and toxic continuations, which is stronger supervision than generic unlabeled or unpaired detoxification. The empirical study focuses only on toxicity, and although the authors argue that the method is general in principle, that generality is not demonstrated beyond this case study. The scale of the experiments is limited to GPT-2 Small, GPT-2 Medium, and Gemma-2B, selected partly to permit mechanistic analysis, so larger-scale validation is left open. Toxicity evaluation relies primarily on Perspective API, and the authors acknowledge that more evaluators would strengthen the results. Enhanced GCG uses the reference model as a malicious teacher, which may be unrealistic in some deployments, though the paper frames it as a deliberate worst-case stress test. Finally, REPO adds a discriminator and adversarial training machinery, but the stated overhead is negligible relative to the dominant cost of the LLM forward and backward pass.
Within the taxonomy of alignment methods, REPO is best characterized as a representation-level alternative to standard preference optimization. It retains the paired supervision structure familiar from DPO-like methods, but uses that structure for token-level representation erasure rather than output-logit preference shifting. The distinctive combination is: paired prompt-conditioned retain and forget continuations, token-level KL anchoring to a frozen reference model on benign continuations, and domain-adversarial training that makes toxic and benign token representations indistinguishable. A plausible implication is that REPO belongs simultaneously to preference optimization, mechanistic unlearning, and adversarial representation learning, with its novelty lying in how these components are coupled for detoxification.