💬 NLP

Divergence Decoding: Inference-Time Unlearning via Auxiliary Models

This paper introduces Divergence Decoding, an inference-time unlearning method that leverages small auxiliary models to steer large language model logits away from sensitive data, effectively mitigating privacy and copyright risks with minimal utility loss while offering a generalizable solution applicable to both text and image domains.

Original authors: Humzah Merchant, Bradford Levy

Published 2026-06-01

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Humzah Merchant, Bradford Levy

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Problem: The "Un-Forgettable" Brain

Imagine you have a super-smart librarian (a Large Language Model) who has read every book in the world. Sometimes, this librarian memorizes specific, sensitive details—like a private diary entry or a copyrighted story—that they shouldn't be sharing.

Usually, if you want the librarian to "unlearn" this specific secret, you have two bad options:

Start Over: Burn down the library and rebuild it from scratch without that one book. This takes millions of dollars and years of work.
Brain Surgery: Try to surgically remove the memory from the librarian's brain. This is risky; often, you accidentally damage their ability to do other things, like writing poetry or solving math problems (this is called "catastrophic forgetting").

The Solution: The "Guide Dog" System

The authors of this paper propose a clever new trick called Divergence Decoding (DD). Instead of trying to change the librarian's brain, they keep the librarian exactly as they are and add two small, specialized "guide dogs" to help steer them during conversations.

Here is how the system works:

The Librarian (The Big Model): This is the main AI we use. It knows everything, including the secrets we want it to forget.
Guide Dog A (The "Forget" Model): This is a tiny, cheap AI trained only on the secret information we want to remove. It is obsessed with that specific data.
Guide Dog B (The "Remember" Model): This is another tiny AI trained only on the safe, public information. It knows everything except the secret.

How It Works: The "Steering Wheel"

When you ask the Librarian a question, the system doesn't just let the Librarian answer. It asks the two Guide Dogs what they think the answer should be.

The Logic: The system looks at the difference between what Guide Dog A (the secret-obsessed one) wants to say and what Guide Dog B (the safe one) wants to say.
The Steering: If Guide Dog A is screaming, "Say the secret!" and Guide Dog B is saying, "No, don't say that!", the system uses that difference to nudge the Librarian's answer away from the secret and toward the safe version.

Think of it like driving a car. The Librarian is the car. The Guide Dogs are two people in the passenger seat. One person is pointing at a cliff (the secret), and the other is pointing at the road (the safe answer). The driver (the system) simply steers the car in the direction of the road, ignoring the cliff, without ever needing to rebuild the car's engine.

Why This Is Better

No Brain Surgery: The main Librarian's brain stays untouched. This means they don't lose their ability to do other tasks.
Cheap and Fast: Training the two small Guide Dogs is like training a puppy—it's quick and cheap compared to rebuilding a whole library.
Works on Hard Questions: Previous methods were good at stopping the AI from reciting a sentence word-for-word, but they failed when the AI tried to answer complex questions about the secret. This method uses the "reasoning" of the Guide Dogs to stop the AI from even thinking about the secret in complex ways.

The "Permanent Fix" (Distillation)

The paper notes that running three models at once (the Librarian + two dogs) takes a little extra computer power. If you want a permanent fix where you only run one model later, you can use a process called Distillation.

Imagine the Librarian, guided by the dogs, writes a perfect "cheat sheet" of how to answer questions without the secrets. You then teach a new, single Librarian to memorize this cheat sheet. Now you have a single, clean model that has "learned" to forget, without needing the dogs anymore.

Does It Work?

The authors tested this on two major benchmarks (TOFU and MUSE), which are like standardized tests for "forgetting."

Results: Their method beat all the previous "state-of-the-art" methods. It was better at removing the secrets while keeping the AI smart.
Beyond Text: They also tried this on image generation (making pictures). They trained the AI to "forget" how to draw specific types of dogs. The method worked there too, successfully stopping the AI from drawing those dogs without ruining its ability to draw other things.

Summary

Instead of trying to erase a memory from a giant, complex brain (which is hard and dangerous), this paper suggests keeping the brain as is and using two small, smart assistants to gently steer the conversation away from the forbidden topics. It's a "learning over forgetting" approach that is cheaper, safer, and more effective than previous methods.

Technical Summary: Divergence Decoding

Problem Statement

Large Language Models (LLMs) frequently memorize sensitive training data, creating significant privacy and copyright risks. Addressing these risks by removing specific knowledge from existing model checkpoints is challenging. Traditional approaches, such as retraining models from scratch on sanitized datasets, are computationally infeasible at frontier scales (requiring millions of GPU-hours). Existing machine unlearning methods, which modify model weights via techniques like gradient ascent or negative preference optimization, often suffer from catastrophic utility loss, are fragile, and remain expensive. Furthermore, these methods frequently degrade the model's overall performance on retained data.

Methodology: Divergence Decoding (DD)

The authors propose Divergence Decoding (DD), an inference-time unlearning mechanism that steers the logits of a target model away from specific data without modifying the model's weights directly. The method is inspired by the Product of Experts framework and Importance Sampling.

Core Mechanism

DD utilizes two small auxiliary models, $p$ and $q$ , trained via standard fine-tuning pipelines to act as proxies for the difference between the current Target model ( $P$ , trained on both retain and forget sets) and the ideal Retrain model ( $Q$ , trained only on the retain set).

Auxiliary Model $p$ : Trained on the "forget" set (or a combination of forget and retain, depending on the benchmark setup).
Auxiliary Model $q$ : Trained on the "retain" set.

The method approximates sampling from the ideal Retrain model $Q$ by adjusting the logits of the accessible Target model $P$ based on the difference between the auxiliary models. Two specific adjustments are proposed:

Linear Divergence Decoding: A linear combination of logits where the difference between auxiliary models steers the target:
$\hat{\ell}^{LC}_Q(x_{<t}) = \ell_P(x_{<t}) + \alpha \cdot [\ell_q(x_{<t}) - \ell_p(x_{<t})]$
Here, $\alpha$ is a hyperparameter. The ratio $q/p$ acts as a "Domain Expert" that upweights tokens likely in the retain set and downweights those likely in the forget set.
Rank-Based Divergence Decoding: A rank-based adjustment that suppresses the top $k$ most divergent logits:
$\hat{\ell}^R_Q(x_{<t}) = \ell_P(x_{<t}) - \mathbb{1}_{\text{rank}(\ell_p - \ell_q) \leq k} \cdot \infty$
(In practice, to avoid degenerate privacy scores, the top $k$ logits are replaced with the $k$ -th largest logit of the unmodified distribution).

Model Distillation

To address the latency and compute costs of running three models (Target + 2 Auxiliaries) simultaneously during inference, the authors propose distilling the behavior of the DD-augmented system back into a single model. A frozen Linear DD system acts as a teacher, and the Target model is fine-tuned as a student using a KL-divergence objective on the forget set. This converts the inference-time controller into a parametric network that can be served without auxiliary models.

Key Contributions

The paper makes the following contributions to the literature on machine unlearning:

Novel Logit Steering: A mechanism that exceeds current state-of-the-art (SOTA) methods across standard unlearning benchmarks (TOFU and MUSE).
Bridge Between Paradigms: A connection between inference-time unlearning and weight-based unlearning via distillation, allowing for flexible deployment.
Cross-Domain Efficacy: Initial explorations demonstrating the method's applicability beyond text generation, specifically in transformer-based image generation.
Extensive Ablations: Comprehensive analysis covering hyperparameter choices, compute/latency impacts, cross-tokenizer/cross-family generalization, and adversarial robustness.

Experimental Results

The authors evaluated DD using the Open Unlearning framework on the TOFU (instruction-tuned, complex queries) and MUSE (pre-trained, few-shot Q&A) benchmarks.

Performance: DD decisively outperforms SOTA baselines (including Gradient Ascent, NPO, RMU, and UNDIAL) on the TOFU benchmark. On MUSE, it marginally improves over strong baselines on Knowledge Memorization and Q&A tasks, though it is less effective on Verbatim Memorization (where simpler filtering methods often excel).
Utility Preservation: Unlike many unlearning methods that cause catastrophic forgetting, DD preserves general knowledge and model utility.
Distillation: The distilled model recovers most of the performance gains of the inference-time Linear DD, outperforming prior gradient-based unlearning methods.
Scalability and Sustainability: The method scales well with increasing forget set sizes and handles sequential unlearning requests without significant performance degradation.
Adversarial Robustness: Under the Leak@K benchmark, which tests if forgotten information can be recovered via repeated adversarial prompting, Distill DD matches the performance of a full retrain, indicating the unlearning is robust and not merely a surface-level correction.
Image Generation: Preliminary results on ImageNet-based image generation show that DD can reduce the generation of specific visual semantics (e.g., dog classes in the forget set) while maintaining perceptual quality, though higher $\alpha$ values are required for linear setups compared to text.

Significance and Claims

The paper argues that the path to safer, compliant AI lies not in "crippling" the base model through direct weight modification, but in guiding it. The authors claim that for many practical use cases (e.g., "Right to be forgotten" requests, copyright compliance, and financial backtesting to prevent look-ahead bias), modifying model weights is unnecessary and potentially harmful to utility.

Instead, DD offers a "learning over forgetting" approach where small, inexpensive auxiliary models steer the generation of a large model. This provides a practical solution for high-stakes deployments where frequent, ad-hoc unlearning is required, redefining SOTA performance by balancing effective unlearning with the preservation of model capabilities. The method is presented as generally applicable to any probabilistic model, extending its utility beyond text to domains like image generation.