Distribution Preference Optimization: A Fine-grained Perspective for LLM Unlearning

TL;DR

This work introduces Distribution Preference Optimization (DiPO), a distribution-level unlearning method for LLMs that directly optimizes the next-token probability distribution rather than whole responses. By modulating high-confidence logits to form memory- and forgetting-promoting distributions, and by leveraging a Bradley-Terry preference framework with Sequence KL divergences, DiPO pushes the model away from dispreferred outputs and toward preferred ones without auxiliary models. Theoretical analysis shows the DiPO loss aligns with the intended unlearning direction, and experiments on TOFU and MUSE demonstrate state-of-the-art forgetting quality and strong utility preservation, with notable stability and scalability advantages. The approach offers a general, efficient, and interpretable framework for targeted unlearning, though it acknowledges limitations such as potential hallucinations and vulnerability to certain privacy-attacks, pointing to future work to broaden privacy assurances and robustness.

Abstract

As Large Language Models (LLMs) demonstrate remarkable capabilities learned from vast corpora, concerns regarding data privacy and safety are receiving increasing attention. LLM unlearning, which aims to remove the influence of specific data while preserving overall model utility, is becoming an important research area. One of the mainstream unlearning classes is optimization-based methods, which achieve forgetting directly through fine-tuning, exemplified by Negative Preference Optimization (NPO). However, NPO's effectiveness is limited by its inherent lack of explicit positive preference signals. Attempts to introduce such signals by constructing preferred responses often necessitate domain-specific knowledge or well-designed prompts, fundamentally restricting their generalizability. In this paper, we shift the focus to the distribution-level, directly targeting the next-token probability distribution instead of entire responses, and derive a novel unlearning algorithm termed \textbf{Di}stribution \textbf{P}reference \textbf{O}ptimization (DiPO). We show that the requisite preference distribution pairs for DiPO, which are distributions over the model's output tokens, can be constructed by selectively amplifying or suppressing the model's high-confidence output logits, thereby effectively overcoming NPO's limitations. We theoretically prove the consistency of DiPO's loss function with the desired unlearning direction. Extensive experiments demonstrate that DiPO achieves a strong trade-off between model utility and forget quality. Notably, DiPO attains the highest forget quality on the TOFU benchmark, and maintains leading scalability and sustainability in utility preservation on the MUSE benchmark.

Figures (4)

• Figure 1: Construction of memory-enhancing distribution $\pi_m$ and forgetting-promoting distribution $\pi_f$ by a memory vector filtered from origin logits.
• Figure 2: Performance analysis on TOFU at the best epoch over five seeds. (a) FQ vs. MU on TOFU-5% and TOFU-10%. DiPO achieves the best trade-off (closest to the "Retrain" target). (b) Training curves of FQ and MU on TOFU-10%, showcasing DiPO's stability and efficacy.
• Figure 3: Robustness analysis on MUSE and DiPO's internal mechanisms. (a) Scalability and Sustainability performance on MUSE News. (b) FQ vs. KL Divergence on TOFU-10% (from $\pi_{\text{ref}}$ on $\mathcal{D}_f$), demonstrating DiPO's higher unlearning efficiency. (c) Return Difference and FQ on TOFU-10%, illustrating the correlation between DiPO's learned reward signals and unlearning efficacy.
• Figure 4: Training curves for only-forget configuration on TOFU-10%, with GA and NPO curves additionally included for comparison.

Theorems & Definitions (3)

• Definition 4.1
• Definition 4.2
• Theorem 4.1: DiPO Loss Function