GRACE:Gradient-guided Coreset Selection for LLM Unlearning

AI summary

Overview

Research area: Machine unlearning (MU) for large language models, specifically the data-selection step that decides which training examples become the forget set and which become the retain set.

Technical level: Intermediate. The paper assumes familiarity with gradient-based training, gradient sketching (Rademacher hashing), sparse approximation (orthogonal matching pursuit) and standard unlearning objectives such as gradient ascent, DPO-style preference losses and representation perturbation.

Scope: The paper proposes GRACE, a gradient-guided coreset selection method that builds compact forget and retain sets from a mixed corpus given only a few seed examples of the undesired behavior, and evaluates it across two domains, two model families and four retain-aware unlearning algorithms.

What This Paper Is About

Most LLM unlearning methods assume the forget set and the retain set are already handed to them. In real deployments, a user, legal entity or administrator can usually point to a behavior they dislike or supply a handful of examples, but cannot identify all training samples responsible for it. GRACE addresses this gap: it takes a few seed examples that trigger the undesired behavior and selects both a compact forget coreset and a compact retain coreset from a larger, mixed corpus, so that any retain-aware unlearning algorithm can then be applied. The goal is a better forget–utility tradeoff than prior gradient-based selection methods.

Key Contributions

1. Formalizing the data-selection problem for LLM unlearning. The authors frame the question: given a small set of examples that elicit undesired behavior, how do we select compact forget and retain sets from a mixed corpus so as to improve the forget–utility tradeoff of unlearning methods?

2. A forget-set selector based on gradient alignment. GRACE computes a forget direction from seed examples, ranks corpus samples by cosine similarity to that direction to form a candidate pool, then applies Non-Negative Orthogonal Matching Pursuit (NNOMP) to pick a sparse forget coreset whose non-negative weighted gradient combination approximates the forget direction.

3. A retain-set selector that is structurally orthogonal to forgetting. The method projects the forget direction out of each remaining gradient, clusters the projected gradients with K-means, and runs RetainOMP (hard projection onto each projected cluster centroid) within every cluster.

4. An empirical study across two domains, two model families and four retain-aware unlearning algorithms, including LLM-as-a-Judge evaluation with Friedman/Nemenyi statistical testing, plus ablations of the projection step, the clustering step, and a forget/retain swap experiment against RASLIK.

Main Findings

• Retrieval (RQ1, Table 1). On WMDP-Bio, GRACE's Forget Retrieval Accuracy is 82.22% with LLaMA 3.1 8B Instruct and 87.78% with Qwen 2.5 3B Instruct, versus 61.67% and 81.11% for RASLIK and 55.55% for both with the embedding baseline. The authors describe this as outperforming the strongest baseline by more than 20 percentage points with LLaMA and more than 6 points with Qwen. On MUSE, GRACE reaches 46.67% (LLaMA) and 43.33% (Qwen); the embedding baseline reaches 47.7% in both cases, about 1 point higher with LLaMA and about 4 points higher with Qwen, while RASLIK scores 28.89% and 15.56%.

• Model utility (RQ2). Averaged over the four unlearning algorithms, GRACE improves model utility by about 5–6 points across most model–dataset settings, except Qwen on MUSE, where it ties with RASLIK. Per-algorithm results (Table 2) show GRACE attaining the highest utility in 6 out of 8 algorithm–model combinations, with the largest gain of +0.16 for GradDiff on LLaMA (0.460 ± 0.042 versus 0.300 ± 0.057 for RASLIK and 0.210 ± 0.042 for Embedding). GRACE's highest reported utility is 0.525 ± 0.064 for SimNPO on LLaMA; its lowest is 0.235 ± 0.106 for GradDiff on Qwen.

• LLM-as-a-Judge on the retain side. Averaged retain scores improve across all datasets and models, with gains up to +1.33 on LLaMA with WMDP-Bio.

• Forget quality is largely invariant to selection (RQ3, Table 3). Mean Forget Quality per algorithm ranges from 0.950 to 1.000, with a maximum difference of 0.012 across selection methods within any single algorithm. The authors conclude that forget quality is driven primarily by the unlearning algorithm, while data selection affects model utility more strongly. On the judge-based forget side (Figure 4), GRACE gains up to +0.37 on MUSE with Qwen, with a small decrease only on MUSE with LLaMA.

• Statistical significance (Table 4). All six Friedman tests are highly significant (p_BH < 10⁻¹⁰ for every metric). GRACE achieves the best average rank on the three retain metrics: preservation 2.170, semantic quality 2.161, coherence/correctness 2.158 (p < 10⁻¹⁶ against both baselines). GRACE also ranks highest on forget-side response coherence (2.279, p < 10⁻¹⁶). The Embedding baseline ranks highest on answer leakage (2.156) and deviation quality (2.095).

• Ablations. Removing projection mainly affects Forget Quality: with SimNPO, FQ drops from 0.95 to 0.91 on WMDP-Bio and from 0.96 to 0.93 on MUSE, despite marginal increases in model utility; with RMU, removing projection either matches or slightly reduces performance. Removing clustering does not improve either metric in any evaluated setting. In the swap ablation with SimNPO on LLaMA 3.1 8B, utility is highest when both sets come from GRACE (0.57 on WMDP-Bio, 0.48 on MUSE), and drops to 0.55 and 0.45 when only the forget set is replaced by RASLIK's.

Methodology in Plain English

The pipeline begins with a user request identifying an undesired behavior and a few seed examples that trigger it. GRACE first computes a "forget direction" by averaging the gradients of the model's loss over those seed outputs; per-sample gradients are compressed to 65,536 dimensions using Rademacher hashing, the same technique used by RASLIK, so the comparison is isolated to the selection logic rather than the gradient approximation.

For the forget set, every sample in the corpus is scored by cosine similarity with the forget direction, and the top-scoring samples form a candidate pool (400 candidates for MUSE, 800 for WMDP-Bio, four times the target forget size). Non-Negative Orthogonal Matching Pursuit then greedily picks a small subset whose gradients, combined with non-negative weights, approximate the forget direction. The non-negativity constraint and the rule that only positively aligned candidates are eligible prevent the selection of corrective counter-forces, yielding a compact and non-redundant set.

For the retain set, the forget direction is projected out of each remaining gradient so only the component independent of the undesired behavior survives. K-means runs on these projected gradients (10 clusters for MUSE, 20 for WMDP-Bio) so that overrepresented data types do not dominate. One pass of RetainOMP is run per cluster, targeting the projected cluster centroid, and 10 samples per cluster are kept, giving retain coresets of 100 (MUSE) and 200 (WMDP-Bio) samples that match the forget-set sizes.

Evaluation uses two datasets. The heterogeneous one combines the 100 Harry Potter forget samples from MUSE Books with Dolly-15k, which contains roughly 15k instruction–response pairs. The domain-specific one combines roughly 1.2k WMDP-Bio biosecurity multiple-choice items (converted to instruction–response pairs, with 200 randomly selected as the forget set) with 20k AlpaCare-MedInstruct samples. Held-out test splits of 200 samples are used for evaluation. Ten seed samples are used for MUSE and 20 for WMDP-Bio. Models are LLaMA 3.1 8B Instruct and Qwen 2.5 3B Instruct, fine-tuned with LoRA for 10 epochs at batch size 32 and learning rate 1×10⁻⁴. Four retain-aware unlearning algorithms are applied for 200 steps at batch size 8: GradDiff, SimNPO + L_r, NPO + L_r and RMU + L_r. All experiments run on a single 48GB A100 with an AMD EPYC 7282 processor. Baselines are embedding-based retrieval with BAAI/bge-large-en-v1.5 and RASLIK. Metrics are Forget Retrieval Accuracy, Forget Quality (inverted harmonic mean of ROUGE-L and conditional probability) and Model Utility (harmonic mean of ROUGE-L, conditional probability and cosine similarity), supplemented by a 1–10 LLM-as-a-Judge protocol, inter-judge agreement measured with Krippendorff's alpha, and Friedman tests with Nemenyi post-hoc comparisons under Benjamini–Hochberg correction.

Why This Matters

Research impact. The paper shifts attention from unlearning algorithms to the data-selection step that precedes them, showing that selector choice mainly moves model utility rather than forget quality when forget quality is already saturated. It also provides a reusable framing—forget direction, orthogonal projection, clustered OMP—that connects coreset selection research to unlearning.

Real-world applications:

• Content removal requests, such as deleting copyrighted book content from a deployed model when only a handful of example passages are available.
• Safety and biosecurity compliance, where an operator wants to suppress hazardous-domain knowledge identified by a few triggering prompts.
• Regulatory compliance with the GDPR's Right to be Forgotten, where an organization must act on a request without a documented map of its training data.
• Domain or identity-level unlearning in production, such as removing an entire concept or a person's information while keeping general assistant quality intact.

Industry relevance. Production teams rarely hold a curated gold forget set, but they do have logs of bad outputs. GRACE's starting point—seed examples of undesired behavior—matches that operational reality, and its reliance on gradient signals means the selected coresets are compact enough to feed existing retain-aware unlearning pipelines directly. The reported gains in utility translate into fewer post-unlearning regressions in general assistant quality.

Future Directions

• Heterogeneous requests. GRACE assumes the seeds induce a reasonably coherent forget direction; handling undesired behaviors that span unrelated phenomena likely requires decomposing a request into multiple homogeneous sub-requests, and this is untested.

• Deployment-realistic evaluation. Current tests rely on controlled settings with known ground-truth forget samples to measure retrieval accuracy; evaluation where the responsible training examples are genuinely unknown or inaccessible remains open.

• Broader benchmarking. The authors call for additional domains, larger models, multilingual settings and more diverse unlearning requests beyond the two domains, two model families and four algorithms studied here.

• Scalability of the selection step. GRACE needs per-sample gradient representations and NNOMP computation; Rademacher hashing and running times comparable to prior gradient-based selection mitigate this, but the authors state scalability still depends on efficient implementations of coreset selection.

Target Audience

Researchers and practitioners working on LLM unlearning, data selection and coreset methods, and machine-learning engineers responsible for privacy, copyright or safety-driven model edits in production. It is also relevant to readers interested in gradient-based data attribution who want a concrete application of orthogonal matching pursuit and gradient projection outside standard training-data curation. Readers need comfort with gradient notation and sparse approximation to follow the method section; the experimental results are readable without it.

Authors’ abstract

Machine Unlearning methods for Large Language Models typically assume pre-specified forget and retain sets. In realistic settings, however, requests may provide only a few examples of undesired behavior, requiring forget and retain sets to be inferred from heterogeneous corpora. We study this data-selection problem and propose GRACE , a gradient-guided coreset selection method that constructs both forget and retain sets for LLM unlearning. GRACE first computes a forget direction from seed examples that elicit the undesired behavior, then selects a compact forget coreset whose gradients approximate this direction using non-negative orthogonal matching pursuit. To preserve model utility, it selects retain examples after projecting out the forget direction and applying clustered orthogonal matching pursuit in the remaining gradient space. Across two target domains, two model families, and four unlearning algorithms, GRACE improves model utility while maintaining comparable forget quality, with particularly consistent gains over prior gradient-based selection methods.