PaperlayerView exact revision on arXiv

Source-linked AI summary

UNDIAL: Self-Distillation with Adjusted Logits for Robust Unlearning in Large Language Models

Yijiang River Dong, Hongzhou Lin, Mikhail Belkin, Ramon Huerta, Ivan Vulić

arXiv:2402.10052v2cs.CLcs.AI

TL;DR

LLMs can retain and expose sensitive information, while existing direct-tuning unlearning methods can collapse through over-unlearning. UNDIAL uses self-distillation toward adjusted logits; the authors report favorable unlearning–usefulness trade-offs and robustness across experiments, while noting limits in model coverage and focused-token selection.

  • Problem

    LLMs can expose memorized private information, and direct-tuning methods such as GA and NPO face over-unlearning and model collapse.

    Source: §1 · p. 1, §3 · p. 3

  • Method

    UNDIAL self-distills toward a target distribution formed by reducing the logit of the token to be unlearned.

    Source: §3.1 · p. 4

  • Results

    UNDIAL achieves a favorable usefulness–unlearning Pareto frontier and the authors report robust performance across benchmarks, hyperparameters, forget-set sizes, and sequential requests.

    Source: §4.4 · pp. 5–6, §6 · p. 8

  • Takeaways & Limitations

    The authors conclude that UNDIAL balances reducing memorization with preserving language generation and understanding capabilities.

    Source: §6 · p. 8

    The study evaluates selected models, including GPT-Neo and LLaMA-2 7B, and could be extended to other and larger LLMs.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Mitigating the retention of sensitive or private information in large language models is essential for enhancing privacy and safety. Existing unlearning methods, like Gradient Ascent and Negative Preference Optimization, directly tune models to remove unwanted information. However, these methods often become unstable because they fine-tune by maximizing cross-entropy loss, which is the opposite of traditional loss minimization in learning. This reversal creates instability, especially on larger datasets, as the model struggles to balance unlearning with maintaining language capacity, leading to over-unlearning. In this paper, we introduce UnDIAL (Unlearning via Self-Distillation on Adjusted Logits), a novel and robust unlearning method. Our approach leverages self-distillation to adjust logits and selectively reduce the influence of targeted tokens. This technique ensures smooth convergence and avoids catastrophic forgetting, even in challenging unlearning tasks with large datasets and sequential unlearning requests. Extensive experiments show that UnDIAL can achieve both robustness in unlearning and scalability while maintaining stable training dynamics and resilience to hyperparameter tuning.

1 Introduction

LLMs can expose memorized private information, while data-protection rules create demand for removing specific information. The introduction presents UNDIAL as a direct-tuning method intended to make that unlearning more robust as data and requests scale.

1 Introduction

  • LLMs can expose memorized personal information, and extraction attacks can heighten data-security concerns.The cited examples include names and addresses as personally identifiable information.Source: §1 · p. 1
  • The GDPR and CCPA’s “right to be forgotten” motivates techniques that remove specific information from LLMs to prevent disclosure.Source: §1 · p. 1
  • Existing unlearning methods can reduce general language usefulness on larger corpora, limiting their real-world use and ability to accommodate continual updates.Source: §1 · p. 1
  • UNDIAL uses self-distillation to adjust the target-token logit and train the model toward the resulting distribution.The adjustment is based on the original model’s outputs, as illustrated in Figure 1.Source: §1 · p. 2, Figure 1 · p. 2
  • The authors report that UNDIAL is robust across hyperparameter settings, forget-set sizes, and multiple unlearning requests.Source: §1 · p. 2

2 Background and Related Work

The background describes how LLM memorization can expose private or copyrighted material, then surveys post-training unlearning approaches. It distinguishes direct tuning from auxiliary-model methods and frames evaluation around forgetting while preserving unrelated knowledge.

2.1 Memorization in Large Language Models

  • LLM memorization can expose private or copyrighted material and increase susceptibility to extraction or membership-inference attacks.Source: §2.1 · p. 2
  • As LLMs grow, the fraction of training data that can be extracted has been reported to increase linearly; deduplication and private training have shown positive effects on reducing memorization.Source: §2.1 · p. 2

2.2 Unlearning in Large Language Models

  • Machine unlearning aims to remove memorized content after training because retraining LLMs on their massive datasets for each removal is impractical.Source: §2.2 · p. 2
  • Direct-tuning methods include Gradient Ascent, which maximizes loss on forgotten tokens, and NPO, which diverges more slowly; alternatives train models to answer “I don’t know” or predict random labels.Source: §2.2 · p. 2
  • Auxiliary-model approaches use contrastive decoding or parameter merging to suppress forget-set generations, while the paper focuses on removing knowledge from the base LLM via fine-tuning.Figure 2 evaluates direct-tuning dynamics using average KL divergence from MUSE reference models on Forget and Retain sets.Source: §2.2 · pp. 2–3, Figure 2 · p. 3

3 Methodology

The methodology motivates UNDIAL by the collapse and over-unlearning of direct-tuning baselines, then defines a self-distillation target that reduces the target token’s logit. It also introduces a focused variant that applies the loss only to selected nouns and entities.

3 Methodology

  • GA and NPO can over-unlearn until model capacity collapses; NPO slows divergence but does not prevent long-term collapse in practice.Source: §3 · p. 3
  • On MUSE, both GA and NPO collapse, with NPO diverging more slowly; NPO also plateaus before briefly unlearning and then collapsing, making it brittle to tuning and stopping time.Source: §3 · p. 3

3.1 UNDIAL: Method Description

  • UNDIAL reduces the target token’s logit by γ, forms an adjusted distribution, and minimizes cross-entropy so the model approximates that fixed target.When the forgotten token is the highest-logit token, the adjustment encourages generation of the second-highest token.Source: §3.1 · p. 4
  • A defined target distribution gives UNDIAL a convergence objective, and the authors report smooth, stable training that permits flexible stopping and avoids degradation.Source: §3 · p. 3, §3.1 · p. 4

3.2 Variant: Focused UNDIAL (FUNDIAL)

  • FUNDIAL applies its loss only to selected key tokens, treating nouns and named entities as targets; spaCy extracts these tokens in practice.The authors describe this noun-and-entity selection as a simplifying assumption and leave more dynamic selection for future work.Source: §3.2 · p. 4

4 Case Study One: Extraction Data

The evaluation uses extractable GPT-Neo training examples and measures both memorization and language usefulness. UNDIAL maintains a favorable unlearning–generation trade-off, including at low memorization levels, and its focused variant further improves that trade-off.

4.1 Dataset and Model

  • The forget set contains 15,000 extractable 200-token sequences from the Pile, previously shown to be memorized by GPT-Neo-family models.Its smaller size relative to MUSE permits extensive ablation studies.Source: §4.1 · p. 4

4.2 Unlearning Metrics

  • Memorization Accuracy measures exact-token recall and Extraction Likelihood measures n-gram overlap; both metrics approximate context evaluation at 40-token intervals.The interval reduces the computational cost of evaluating every context length.Source: §4.2 · p. 5

4.3 ‘Model Usefulness’ Metrics

  • Usefulness is evaluated through six QA-style NLU accuracy benchmarks and WikiText-103 generation measured by MAUVE, repetition, and perplexity.Generation uses 5,000 samples and a 32-token context prompt.Source: §4.3 · p. 5

4.4 Results and Discussion

  • UNDIAL occupies the favorable EL–MAUVE trade-off with state-of-the-art language performance and high unlearning accuracy; NPO matches only with careful tuning and loses robustness as settings change.The figure varies unlearning strength to show the trade-off between extraction likelihood and language capacity.Source: Figure 3 · p. 5, §4.4 · p. 5
  • At EL around 0.1 (about 50% less memorization), NLU remains within 5% of baseline, while GA, TA, DP, and CD sharply reduce MAUVE; NPO and UNDIAL retain high MAUVE.NPO and UNDIAL also show less degradation in perplexity and Rep3 than the auxiliary-model methods.Source: §4.4 · pp. 5–6, Table 1 · p. 6
  • Focused UNDIAL targets entity names or nouns rather than functional words; FUNDIAL outperforms standard UNDIAL on the EL–MAUVE trade-off.The figure also compares stable training dynamics across variants at γ=30.Source: Figure 4 · p. 6, §4.4 · pp. 6–7
  • Below EL 0.05, NPO’s MAUVE falls below 0.03, while UNDIAL continues producing high-quality outputs.The NPO outputs include repetitive or unnatural sentences.Source: §4.4 · p. 6

5 Case Study Two: MUSE Benchmark

On MUSE, UNDIAL balances forgetting targeted BBC News knowledge with preserving utility, and the reported results show robustness across hyperparameters, larger forget sets, and sequential requests.

Evaluation setup

  • The MUSE evaluation uses BBC News Forget and Retain sets, measuring forgotten knowledge memorization and retained knowledge utility with Q&A accuracy.The setup uses LLaMA-2 7B with LoRA rank 8.Source: §5 · p. 7, Figure 5 · p. 7

UNDIAL Achieves a Better Pareto Frontier

  • UNDIAL achieves a superior Pareto frontier to direct-tuning and auxiliary-model baselines, including variants that add retain-set objectives.It reaches state-of-the-art performance without using the Retain set.Source: Figure 5 · p. 7, §5 · p. 7

UNDIAL Shows Robust Training Dynamics

  • UNDIAL maintains stable training dynamics, unlike GA and NPO, which exhibit over-unlearning and catastrophic forgetting.The paper attributes this stability to defining the target distribution during unlearning.Source: §5 · pp. 7–8

UNDIAL is More Robust to Different Hyperparameter Setups

  • Across learning rates 1e-4 and 1e-5 and unlearning strengths γ = 2, 4, 8, and 16, UNDIAL converges without collapse.Higher unlearning strength can hurt usefulness; γ values 2, 4, and 8 yield similar forget distances, while larger values lead to higher retain distances.Source: §5 · p. 7, Figure 6 · p. 8

Unlearning with UNDIAL is More Scalable and Sustainable

  • With larger forget sets, UNDIAL’s usefulness remains reasonably well maintained while GA and NPO scores drop sharply.For sequential requests, UNDIAL maintains utility above 0.4 with minimal degradation, while baselines’ utility drops close to zero.Source: Figure 7 · p. 8, §5 · p. 8

6 Conclusion

UNDIAL uses self-distillation to balance reducing memorization with preserving language capabilities. Experiments report state-of-the-art unlearning performance and robustness across several settings.

Conclusion

  • UNDIAL is a self-distillation method intended to balance reducing memorization with preserving language generation and understanding.The conclusion describes it as a direct-tuning unlearning method.Source: §6 · p. 8
  • Experiments on Extraction Data and MUSE report state-of-the-art unlearning performance.The paper also reports resilience across varying hyperparameters, forget-set sizes, and sequential requests.Source: §6 · p. 8
  • The authors present UNDIAL as a promising next step for real-world applications requiring LLM unlearning.They report that it prevents model collapse and scales efficiently.Source: §6 · p. 8

Limitations

The study’s model coverage and the focused FUNDIAL variant’s token-selection assumption delimit the scope of its conclusions.

Model scope

  • The study focuses on GPT-Neo and LLaMA-2 7B, leaving extension to other and larger LLMs for future work.The authors cite prior benchmark use and available computational resources and budget as reasons for this selection.Source: Limitations · p. 9

Focused FUNDIAL

  • Focused FUNDIAL assumes nouns and named entities are the target tokens, which may not always identify sensitive information accurately.The authors suggest more sophisticated token selection, such as automatic detection, as future work.Source: Limitations · p. 9

Ethical Consideration

The authors frame LLM privacy improvements as potentially supporting trust and safety in applications.

Ethical Consideration

  • The authors suggest that mitigating private-data concerns could foster greater trust and safety in LLM applications.They also express hope that such efforts would increase societal trust in models’ controlled usefulness.Source: Ethical Consideration · p. 9

A.1 Hyperparameter Selection

Experiments used GPT-Neo models at three parameter scales, with batch and optimizer settings specified for the runs.

A.1 Hyperparameter Selection

  • The experiments tested models with 125M and 1.3B parameters for extracting memorization data (Carlini et al., 2021).Source: §A.1 · p. 13
  • The 125M model used batch size 64; the 1.3B model used mini-batches of 16 with gradient accumulation of 4.Source: §A.1 · p. 13
  • All experiments used the AdamW optimizer.Source: §A.1 · p. 13

A.2 (A Brief Summary of) Baseline Models

This appendix outlines baseline unlearning methods, reports UnDIAL’s scaling tests, and provides implementation and generation examples. The reported scaling result covers larger GPT-Neo models and both full and LoRA fine-tuning.

A.2 (A Brief Summary of) Baseline Models

  • Gradient Ascent decreases memorized-token generation probability by maximizing log-likelihood loss on memorized data, reversing typical loss minimization.Source: §A.2 · p. 13

Negative Preference Optimization (NPO)

A.4 Implementation of UNDIAL

  • The implementation modifies the Hugging Face Trainer with a compute-loss function; Figure 8 presents Python code for UnDIAL.Source: §A.4 · p. 13, Figure 8 · p. 14

A.3 Scaling GPT-Neo on Unlearning Extraction Data

  • UnDIAL is reported robust from 1.3B to 2.7B parameters and across full fine-tuning and LoRA on the Extraction dataset.Table 2 reports EL3 and MAUVE deltas and NLUa, the average over six NLU tasks; lower is better except for MAUVE and NLUa.Source: §A.3 · p. 13, Table 2 · p. 14

A.5.1 Example Generation on CC News

  • The CC News examples include contrastive-decoding outputs at coefficients 0.25 and 0.5, with MAUVE values 0.172 and 0.333, respectively.Source: §A.5.1 · p. 16

A.5.2 Example Generation on Forget Set

Related papers

Loading 2402.10052v2…