Unilogit: Robust Machine Unlearning for LLMs Using Uniform-Target Self-Distillation

TL;DR

Unilogit addresses the problem of removing specific memorized knowledge from LLMs under GDPR-like privacy requirements by using a uniform-target self-distillation strategy. It dynamically constructs targets from the current model so that the forget token attains a uniform probability, while preserving other logits, and optimizes with reverse KL for forgetting and a retain-regularizer on retain data, all without extra hyperparameters. Across public benchmarks (MUSE-News, RWKU) and an in-house e-commerce dataset, Unilogit delivers superior forgetting-utility trade-offs and robust hyperparameter behavior compared with state-of-the-art baselines like NPO and UnDIAL. The work provides strong practical evidence for a scalable, robust unlearning approach with detailed analyses of target quality and ablations, highlighting its potential for real-world privacy-preserving ML deployments.

Abstract

This paper introduces Unilogit, a novel self-distillation method for machine unlearning in Large Language Models. Unilogit addresses the challenge of selectively forgetting specific information while maintaining overall model utility, a critical task in compliance with data privacy regulations like GDPR. Unlike prior methods that rely on static hyperparameters or starting model outputs, Unilogit dynamically adjusts target logits to achieve a uniform probability for the target token, leveraging the current model's outputs for more accurate self-distillation targets. This approach not only eliminates the need for additional hyperparameters but also enhances the model's ability to approximate the golden targets. Extensive experiments on public benchmarks and an in-house e-commerce dataset demonstrate Unilogit's superior performance in balancing forget and retain objectives, outperforming state-of-the-art methods such as NPO and UnDIAL. Our analysis further reveals Unilogit's robustness across various scenarios, highlighting its practical applicability and effectiveness in achieving efficacious machine unlearning.

Figures (6)

• Figure 1: Overview of self-distillation unlearning in Unilogit: Starting with the output logits of the LLM, the target logit is diminished, so that after softmax the target token in the modified distribution has uniform probability. Soft labels are derived from the current model ($\theta$) outputs. Reverse KL divergence is the distillation objective.
• Figure 2: Results for the MUSE-News benchmark for different unlearning methods using multiple different hyperparameters. On the x-axis we have the retain performance and on the y-axis the forgetting performance, both for the QA task.
• Figure 3: Results for the RWKU-News benchmark for different unlearning methods using multiple different hyperparameters. On the x-axis is the retain performance and on the y-axis the forgetting performance.
• Figure 4: Comparison of unlearning methods on listings from three different sellers across three forget set sizes in our e-commerce dataset. Forget Completion and Neighbors Completion are evaluated using ROUGE-recall scores. Marker sizes and number annotations indicate MMLU scores, reflecting general model abilities.
• Figure 5: Left: Average KL divergence between the retrained model outputs and the soft labels of UnDIAL and Unilogit on the forget set. Center: KL divergence progression between soft targets and retrained model outputs for both methods over unlearning epochs. Right: Average KL divergence between unlearned model outputs and retrained model for NPO, UnDIAL, and Unilogit. Lower values indicate better performance in all cases, as well as the baseline represents average KL between the outputs of the starting model and the retrained model.