Unilogit: Robust Machine Unlearning for LLM s Using Uniform-Target Self-Distillation (2025)

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.

## introduction large language models (llms) have advanced rapidly, becoming widely applicable in various settings @xcite @xcite . however, their increasing capabilities raise significant privacy risks, especially for individuals whose sensitive data may have been included in training. this information can become embedded within the model, making it susceptible to unintended exposure through memorization, adversarial exploits, membership inference (mia), and model inversion attacks @xcite . to address these concerns, regulatory frameworks such as the general data protection regulation (gdpr) have been established to protect individual privacy and enforce the right to be forgotten. given that llms are subject to such regulations, the machine learning research community has increasingly focused on the emerging field of machine unlearning for llms @xcite @xcite , which aims to develop methods for selectively removing specific knowledge from models. this includes erasing sensitive information @xcite , forgetting entire entities or facts @xcite , and removing harmful or biased information @xcite . in the machine unlearning framework, we define the full training dataset as a partition of two subsets: the forget set, which consists of the data to be unlearned, and the retain set, which contains the remaining knowledge that should be preserved after unlearning. an effective machine unlearning method aims to produce a model that successfully forgets the forget data, while maintaining the integrity of the retained knowledge. specifically, the resulting unlearned model should satisfy the following key requirements: 1) minimize the retention of information from the forget set; 2) maintain high performance on the retain set; 3) require less computational cost than retraining the model from scratch on the retain set; 4) maintain inference efficiency, i.e., ensuring unchanged latency. despite extensive research efforts @xcite , current unlearning methods still face significant challenges in achieving all these goals simultaneously. a major challenge in this domain is catastrophic forgetting @xcite , where the model suffers severe degradation in its ability to retain knowledge from the retain set while attempting to remove forget set knowledge. additionally, unlearning methods must balance forget and retain performance @xcite , as no existing technique can fully erase the forget set information, while preserving the original accuracy. to analyze this trade-off, researchers commonly visualize unlearning performance through pareto frontiers that plot forgetting effectiveness against retention performance across various hyperparameter sweeps @xcite . finally, the balancing problem is related to the issues of hyper-parameter tuning and robustness, as each method and problem combination have their unique optimal set of hyperparameters @xcite . in this paper, we propose unilogit, a selfdistillation approach for unlearning. it generates targets from the output of the current model by assigning a uniform probability to the target token in forget samples and redistributing the remaining probability mass. our method is inspired by @xcite and is driven by the question: can we leverage existing information to refine the target distribution for forgetting? unilogit offers a simple yet effective unlearning strategy that outperforms existing methods, demonstrating superior hyperparameter robustness and applicability across diverse scenarios, satisfying all key aforementioned unlearning requirements. unlike prior techniques that introduce an extra hyperparameter in their loss @xcite @xcite , unilogit achieves consistent forgetting performance without the extra tuning overhead. to validate our approach, we conduct in-depth auxiliary studies: (1) demonstrating that our soft labels and outputs are more accurate than those of other methods (section 4.4), and (2) ablation studies to assess the impact of key methodological components (section 4.5). our contributions are as follows: parameter-tuning. among approximate methods, one prominent direction is parameter-tuning approaches, which directly modify model parameters to achieve unlearning. we pursue this direction because these methods typically meet all unlearning requirements, preserving inference latency without demanding excessive training compute. parameter-tuning methods frame unlearning as an optimization problem with two competing objectives: a forget objective l f that forces the model to unlearn specific knowledge and a retain objective l r that ensures performance on the remaining data is preserved. a generalized unlearning loss function typically follows this form: notable parameter-tuning methods include gradient ascent (ga) @xcite , which suffers from instability, and negative preference optimization (npo) @xcite , which has emerged as a robust state-of-the-art method by introducing a controlled forgetting process. npo uses a preference optimization-based loss function to mitigate the risk of catastrophic forgetting. another relevant approach, me+gd @xcite , maximizes the entropy of the model's predictions on the forget set by pushing the output probabilities towards a uniform distribution, preserving performance on the retain set using crossentropy. the concept of auxiliary models in unlearning was introduced by @xcite with their "who is harry potter" (whp) approach, which leverages a reinforced model finetuned on the forget set to inform the unlearning process. this idea has been further developed in distillation-based unlearning methods. for instance, @xcite propose rkld, which refines the model reinforcement of whp by distilling modified soft labels into the original model for targeted unlearning. @xcite introduced undial, a self-distillation approach that adjusts logits to selectively reduce the influence of targeted tokens, ensuring smooth convergence and avoiding catastrophic forgetting. our approach is motivated by these advancements, particularly in refining the process of generating effective soft labels for unlearning. outside of the realm of text-based llms, the work most closely related to our approach is proposed by @xcite . in their framework, they tackle unlearning by minimizing kl divergence between model outputs and a uniform distribution, then applying an mse loss on the adjusted logits. while @xcite provide a general framework, particularly for weaklysupervised settings-where only limited or noisy supervision is available-and experiments with computer vision tasks, our method is specifically designed for llms. we achieve unlearning by optimizing reverse kl divergence at the categorical probability distribution level and dynamically updating target distributions based on the model's latest outputs. these features enhance adaptability and unlearning effectiveness, setting our approach apart from existing methods. for an extended discussion on related work, please refer to appendix a. ## methodology in our self-distillation approach to unlearning, the central challenge is designing accurate soft targets that effectively guide the model toward forgetting. ideally, the outputs of the retrained model θ r , would serve as the gold standard for distillation. however, since θ r is unavailable in reality, we must approximate these targets through a principled and computationally efficient method that refines self-distillation for unlearning. unilogit. inspired by recent advances in logit adjustment for unlearning @xcite and self-distillation, we propose unilogit: a selfadjusting self-distillation method for machine unlearning. it sets the target logit so that after the softmax operation, it is equal to a uniform distribution value, while preserving the logits for all other vocabulary entries. for a model output vocabulary v , output logit function h(x; θ), parametrized by the current model's parameters θ, and a one-hot label vector t, we calculate the target logits h(x; θ): then, we calculate the soft label target distribution, where the target token label is going to be equal to a uniform probability ( for a detailed derivation, see appendix b. this design is grounded in the intuition that the current model θ and the retrained model θ r should be relatively close in both parameter space and output distributions, given that the retain and forget sets originate from the same data distribution and the forget set is significantly smaller. consequently, the non-target token logits of θ serve as a strong prior for approximating the output distribution of θ r . by explicitly setting the target token probability to a uniform value, we induce the desired unlearning effect while redistributing the lost probability mass according to this prior. we adopt a uniform probability ( 1 |v | ) for the target token as it represents a state of complete uncertainty, aligning with the goal of eliminating learned information about the forget token. this choice is also justified by prior work indicating that untrained models tend to produce nearly uniform output distributions @xcite , making it a natural approximation of an untrained state. importantly, this approach introduces no additional bias in determining the target token's probability, as we have no prior information about its true distribution post-unlearning. unilogit has two beneficial properties over un-dial: 1) it eliminates the need for a manually tuned hyperparameter γ to scale down the target logit and 2) by explicitly setting the target probability to uniform, it adaptively adjusts the reduction factor in a self-consistent manner, ensuring stability and interpretability. a crucial distinction between our approach and previous self-distillation-based unlearning methods @xcite @xcite ) is that we construct distillation targets from the current model parameters θ rather than the initial model parameters θ o . this choice is motivated by the assumption that a well-designed unlearning algorithm should progressively guide the model closer to the retrained model θ r . if this assumption holds, then at each unlearning step, the output distribution of θ should increasingly resemble that of θ r . by leveraging the latest model outputs as the basis for our self-distillation targets, we iteratively refine the approximation of θ r , leading to more accurate guidance for the unlearning process. this directly addresses the core question of our work (question 1), and we empirically validate this hypothesis in section 4.4. once we have our target distribution p, we can formulate the full training objective of unilogit: for the forget loss, we adopt reverse kldivergence (rkl) due to its advantageous properties for unlearning. unlike forward kl (fkl), which is mean-seeking, rkl is mode-seeking, thus penalizing cases where the model assigns high probability to an incorrect token for which the target distribution has assigned a low probability. this property is particularly well-suited for unlearning, as it strongly discourages the model from retaining high confidence in previously learned outputs. in appendix 4.5, we provide an ablation study demonstrating that rkl leads to superior unlearning performance compared to fkl. prior works @xcite have also observed that rkl improves evaluation metrics at the cost of reduced generation diversity. however, in the unlearning setting, this tradeoff is acceptable, as we prioritize strong unlearning performance over minor diversity loss. moreover, since our target distributions for non-target tokens remain largely consistent with the model's original outputs, with the only major change occurring in the target token's probability reduction, the overall impact on generation variance is expected to be minimal (see column flu in finally, we evaluate the different methods on an in-house e-commerce benchmark, which reflects a real-world use case scenario. the initial model employed is e-llama-3.1 8b @xcite , which has been trained extensively on public listing data from the ebay e-commerce platform. the unlearning benchmark consists of three different unlearning targets, which are comprised by different sellers from the platform. each seller has associated item listings, which need to be forgotten. the three sellers reflect three unlearning scenarios in the amount of data to be unlearned: a small, a medium and a large amount of listings. evaluation is performed by calculating rouge-recall score on text completions for forget and retain set respectively @xcite . more details on datasets and evaluation tasks used can be found in appendix c.1. settings. we closely follow the experimental setups described in the respective papers for each benchmark and cited method. all training is conducted using four a100 80gb gpus. for muse-news, we run each method for 10 epochs with λ = 1, a batch size of 32, and learning rates from the muse paper @xcite . additional hyperparameter tuning is performed to generate sweep curves, with reference ranges derived from the original method papers when available @xcite @xcite . specifically, we use a learning rate of 5e-6 for me+gd and 1e-5 for rkld+kl. please refer to for rwku, we train for 3 epochs with varying learning rates (detailed in section d.2). learning rate sweeps in the range [1e-7, 1e-5] are conducted for undial and unilogit, while for other methods, we follow hyperparameter choices from their respective papers. for our e-commerce benchmark, we train for 10 epochs for the smallest seller and 3 epochs for the medium and largest sellers. we initially apply hyperparameters as recommended in the original method papers, and perform additional tuning when necessary to ensure optimal performance for each method. ## conclusion we present unilogit, a novel method for efficient and effective machine unlearning. ting and utility preservation. our results show that unilogit achieves the most optimal pareto curves, with superior performance in retaining model utility while effectively forgetting unwanted knowledge. our method exhibits robustness to hyperparameter tuning, providing stable and consistent performance across various settings. additionally, it demonstrates computational efficiency, with faster training iterations then current methods in scenarios without a retain loss term and comparable performance when a retain loss is included. we also provide an in-depth analysis of unilogit's self-distillation targets and output distributions, highlighting its ability to generate more accurate targets compared to undial and its capacity to maintain better output distributions postunlearning. the ablation studies reinforced the importance of key design choices, such as using reverse kl divergence, adaptively calculating soft labels, and leveraging current model logits during training, all of which contribute to unilogit's superior performance. overall, unilogit demonstrates both practical robustness and theoretical advantages, making it a promising method for real-world applications.