Source-linked AI summary
GRACE:Gradient-guided Coreset Selection for LLM Unlearning
Praveen Bushipaka, Andrea D'Angelo, Lucia Passaro, Tommaso Cucinotta
TL;DR
LLM unlearning often lacks complete forget and retain sets when requests provide only sparse undesired-behavior examples. GRACE constructs compact coresets by matching gradients for forgetting and projecting out the forget direction for retention, improving utility in most evaluated settings while maintaining comparable forget quality. The method assumes that seed examples induce a reasonably coherent forget direction.
Problem
LLM unlearning commonly assumes pre-specified forget and retain sets, although realistic requests may provide only sparse behavior examples and require inference from heterogeneous corpora.
Method
GRACE computes a forget direction from undesired outputs, selects a compact aligned forget coreset with non-negative orthogonal matching pursuit, and selects retain examples after projection and clustered matching pursuit.
Results
Across two domains, two model families, and four unlearning algorithms, GRACE improves model utility in most evaluated settings while maintaining comparable forget quality.
Takeaways & Limitations
GRACE provides a stronger forget-utility tradeoff than prior selection baselines for constructing LLM-unlearning coresets from sparse undesired-behavior examples.
Takeaways & Limitations
GRACE assumes that seed examples induce a reasonably coherent forget direction and may be less effective for heterogeneous undesired behavior spanning unrelated phenomena.
Abstract
from arXiv · showhide
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.
1 Introduction
Realistic LLM unlearning must infer compact forget and retain sets from heterogeneous corpora because requests often provide only sparse behavior examples. GRACE addresses this selection problem with gradient-guided coresets and improves utility while maintaining forget quality.
- Motivation: LLM unlearning methods commonly assume that forget and retain sets are already specified, although realistic requests may identify only undesired behavior or a few examples.Training data may be undocumented or only partially accessible, requiring both sets to be constructed from heterogeneous corpora.
- Research question: The central problem is selecting compact forget and retain sets from a mixed corpus to improve the forget–utility tradeoff.The selection must begin from examples that elicit the undesired behavior.
- Gap: RASLIK ranks samples independently by gradient cosine similarity, but does not explicitly optimize aggregate gradient alignment or preserve useful directions while excluding the forget direction.This can select redundant similar samples and does not ensure that dissimilar samples support model utility.
- GRACE: GRACE computes a forget direction, selects a compact forget coreset with non-negative orthogonal matching pursuit, then selects retain examples after removing that direction from gradient space.The resulting coresets are passed to retain-aware unlearning algorithms.
- Results: Across two target domains, two model families, and four unlearning algorithms, GRACE improves downstream model utility in most settings while maintaining forget quality.The method shows consistent gains over RASLIK.
2 Related Work
Prior work spans LLM unlearning, general data selection, and methods for constructing forget or retain sets. GRACE targets the less-explored setting where both sets must be selected from heterogeneous data given only sparse undesired-behavior examples.
- LLM Unlearning: LLM unlearning methods include input-based, fine-tuning-based, auxiliary, and editing-based approaches, with retain-aware methods using retain losses to preserve non-target knowledge.The studied retain-aware methods are GradDiff, NPO, SimNPO, and RMU.
- Data selection in LLMs: Coreset selection seeks a small representative subset that preserves essential dataset patterns, while several LLM methods address efficient data selection or reduction.Some prominent methods rely on external models, and TAGCOS provides a clustering- and OMP-based foundation.
- Data selection in LLM Unlearning: Existing unlearning data-selection approaches either assume forget and retain sets are available or depend on sufficient data structure to identify relevant neighbors.Smaller forget sets can suffice but may require more unlearning steps.
- Prior selection baseline: RASLIK selects both sets from a dataset using gradient cosine similarity to undesired-behavior samples, assigning closest samples to forgetting and antipodal samples to retention.It is identified as the only LLM-unlearning method in this setting among the cited prior work.
- Positioning: GRACE differs from prior gradient-guided and coreset methods by jointly constructing forget and retain sets for the distinct objectives of LLM unlearning.Cosine similarity may not retrieve the best samples, while generic gradient and coreset methods do not distinguish forgetting from retention.
3 Background
LLM unlearning seeks an updated model that approximates retraining on retained data without the cost of retraining from scratch. Because full training and retain-set coverage is difficult, GRACE evaluates compact coresets using proxy metrics.
- Gradient representation: Rademacher hashing maps large per-sample gradients to compact randomized sketches that support approximate gradient-similarity computations.The sketches project gradients onto random vectors with entries in {−1,+1}.
- Machine Unlearning: Machine unlearning removes the influence of a forget set so the updated model behaves as if those samples had not been used during training.The ideal target is a model retrained from scratch on the retain set.
- Gold model: A gold model is obtained by applying the training algorithm from scratch to the retain set, but this ideal solution is generally infeasible after each request.Unlearning methods instead approximate this gold model while avoiding retraining cost.
- Formal objective: A machine-unlearning method takes the trained model, forget set, and retain set and outputs an unlearned model as close as possible to the gold model.The retain set is defined as the training set excluding the forget set.
- Evaluation context: Modern LLM settings rarely provide full coverage of the training and retain sets, so unlearning evaluations use proxy metrics for forget quality and utility retention.GRACE therefore focuses on selecting compact forget and retain coresets.
4 GRACE
GRACE uses first-order gradient information to build two coresets from sparse undesired outputs: a forget coreset aligned with the forget direction and a retain coreset separated from it. Projection, clustering, and orthogonal matching pursuit organize representative retention data.
- Forget coreset: GRACE averages hashed loss gradients from observed undesired outputs to form a more stable forget direction and ranks training samples by cosine similarity to it.The highest-scoring samples form the candidate forget pool.
- Forget coreset: Non-Negative Orthogonal Matching Pursuit selects a sparse forget subset whose non-negative weighted gradient combination approximates the forget direction.The non-negativity constraint prevents corrective counter-forces, and negatively aligned candidates are excluded.
- Forget coreset: The resulting forget coreset is compact and non-redundant, with aggregate gradients closely matching the undesired behavior.If a forget set is externally provided, the forget-direction construction step can be omitted.
- Retain coreset: For retention, GRACE projects the forget direction out of remaining gradients, isolating components independent of the undesired behavior.This projection defines the gradient space used for subsequent retain selection.
- Retain coreset: GRACE clusters projected gradients to prevent overrepresented data types from dominating, then runs RetainOMP independently within each cluster.Both clustering and selection are orthogonal to the forget direction, structurally excluding forget-aligned content.
- Scope of projection: Hard projection is most useful when the forget direction is coherent, such as in class-level or identity-level unlearning.Class-specific gradients are described as structurally distinct from those of other classes, reducing entanglement with retained data.
5 Experimental Setup
The experiments test whether GRACE improves forget-sample retrieval and post-unlearning utility while maintaining forget quality across heterogeneous and domain-specific datasets, multiple models, and unlearning algorithms.
- The evaluation addresses three questions: retrieval of ground-truth forget samples, post-unlearning model utility, and retained forget quality.
- The study uses MUSE Books combined with Dolly-15k as a heterogeneous dataset and WMDP-Bio combined with AlpaCare-MedInstruct as a domain-specific dataset.The datasets contain approximately 15k heterogeneous instruction-response pairs and 20k medical instruction-response samples, respectively.
- Forget retrieval begins with 10 MUSE or 20 WMDP-Bio seed samples, which are combined with retrieved samples to form the final forget set.
- GRACE selects retain coresets by clustering gradients into 10 MUSE or 20 WMDP-Bio clusters and choosing 10 samples per cluster.The resulting retain coresets contain 100 and 200 samples, matching the respective forget-set sizes.
- The study evaluates LLaMA 3.1 8B and Qwen 2.5 3B using GradDiff, SimNPO + Lr, NPO + Lr, and RMU + Lr.All models use the same fine-tuning setup across MUSE and WMDP-Bio, with 10 training epochs, batch size 32, and learning rate 1 × 10−4.
- FRA measures correctly retrieved ground-truth forget samples, while FQ and MUT jointly assess forgetting efficacy and utility preservation.MUT uses held-out test splits and combines ROUGE-L, conditional probability, and cosine similarity through a harmonic mean.
6 Results
Across quantitative, qualitative, and statistical evaluations, GRACE generally improves model utility over selection baselines while preserving comparable forget quality. Its retain-side advantage is especially consistent across models, datasets, and unlearning algorithms.
- 6.1 RQ1: GRACE outperforms the strongest WMDP-Bio retrieval baseline by more than 20 points with LLaMA and more than 6 points with Qwen.
- 6.2 RQ2: GRACE improves quantitative model utility by about 5–6 points across most model–dataset settings, tying RASLIK only for Qwen on MUSE.Utility is averaged across GradDiff, SimNPO, NPO, and RMU.
- 6.2 RQ2: GRACE achieves the highest utility in 6 of 8 algorithm–model combinations, with the largest gain reaching +0.16 for GradDiff on LLaMA.
- 6.2 RQ2: LLM-as-a-Judge retain metrics improve across all datasets and models, with gains up to +1.33 on LLaMA with WMDP-Bio.
- 6.3 RQ3: Forget quality varies by at most 0.012 across selection methods within each unlearning algorithm, indicating limited selector impact in the evaluated setup.
- 6.3 RQ3: Qualitative forget-side scores favor GRACE in most settings, with gains up to +0.37 on MUSE with Qwen and a small decrease only on MUSE with LLaMA.
- 6.4 Statistical Analysis: All six Friedman tests are significant at pBH < 10^-10, while GRACE has the highest average ranks on all three retain metrics.GRACE significantly outperforms both Embedding and RASLIK on retain metrics; Embedding ranks highest on answer leakage and deviation quality.
7 Ablation studies on GRACE
Ablations show that projection and clustering contribute to GRACE’s overall forget–utility tradeoff, while forget quality is largely saturated and utility gains are not attributable solely to retain-set construction.
- Projection ablation: Removing projection decreases SimNPO FQ from 0.95 to 0.91 on WMDP-Bio and from 0.96 to 0.93 on MUSE.These decreases occur despite marginal increases in MUT.
- Clustering ablation: Removing clustering does not improve either metric, and GRACE matches or outperforms the no-clustering variant in every evaluated setting.The largest differences occur in MUT for RMU.
- Forget Quality is largely saturated across the experiments, consistent with prior evidence that small forget coresets can achieve high FQ.The swap ablation therefore tests whether GRACE’s utility gains arise solely from retain-set construction.
- Forget–retain swap: Swap experiments show similarly high FQ across configurations, while utility gains cannot be attributed solely to retain-set construction.The comparison fixes SimNPO on LLaMA 3.1 8B while interchanging GRACE and RASLIK selections.
8 Conclusion
GRACE constructs compact forget and retain coresets from limited undesired-behavior examples rather than assuming those sets are given. Across diverse domains, models, and unlearning algorithms, it improves utility while maintaining comparable forget quality.
- GRACE constructs compact forget and retain coresets from a small set of examples or outputs exhibiting undesired behavior.It approximates the forget direction with non-negative Orthogonal Matching Pursuit and selects retain examples after projecting out that direction.
- Across two domains, two model families, and four unlearning algorithms, GRACE improves model utility in most settings while maintaining comparable forget quality.The resulting tradeoff is stronger than prior selection baselines.
- Despite saturated post-unlearning forget quality, GRACE consistently improves forget-set retrieval.
- Ablations show that removing projection or clustering does not improve the overall forget–utility tradeoff.Forget–retain swaps further show that utility gains are not attributable solely to retain-set construction.
9 Limitations
GRACE is limited by assumptions about coherent forget directions, controlled evaluation conditions, incomplete benchmark coverage, and the computational cost of per-sample gradient representations.
- GRACE assumes seed examples induce a coherent forget direction, limiting effectiveness for heterogeneous behaviors spanning unrelated phenomena.Such requests may need decomposition into more homogeneous sub-requests.
- The controlled evaluation uses ground-truth forget samples for retrieval measurement, only partially reflecting deployments where responsible training examples may be unknown or inaccessible.
- Evaluation across two domains, two model families, and four algorithms is not a comprehensive benchmark for LLM-unlearning data selection.The authors call for broader testing across domains, model sizes, languages, and request types.
- Per-sample gradient representations cost more than embedding-based retrieval signals, so scalability depends on efficient coreset-selection implementations, especially NNOMP.Rademacher hashing reduces this cost, and reported running times are comparable to prior gradient-based selection.
A Retrieval Methods
The retrieval pipeline reduces gradient dimensionality, selects a forget coreset aligned with seed-induced behavior, and constructs a retain coreset from gradients orthogonal to that forget direction before applying retain-aware unlearning objectives.
- Gradient dimension reduction: Gradient representations are projected with random Rademacher vectors, permuted, and normalized to reduce them efficiently to k dimensions.GRACE and RASLIK use k = 65,536 for comparison.
- Forget selection: GRACE averages seed gradients to obtain forget direction gf, retrieves a candidate pool, and applies Non-negative OMP to select a forget coreset.Seed counts are 10 for MUSE and 20 for WMDP-Bio; the candidate pool is four times the forget-set size.
- Retain selection: For retain selection, GRACE excludes forget candidates, projects gf out of each gradient, clusters the remainder, and applies OMP within each cluster.It uses K = 10 for MUSE and K = 20 for WMDP-Bio, with |Dr| equal to |Df|.
- Retrieval cost: Embedding retrieval is cheapest because embeddings are pre-stored, whereas RASLIK and GRACE require gradient computation; their running times are comparable.
- Unlearning algorithms: The retain-aware methods studied are GradDiff, NPO, SimNPO, and RMU, using retain objectives to preserve non-target knowledge and stabilize training.The retain loss commonly uses cross-entropy or KL-divergence regularization.
- Evaluation metrics: Forget Quality combines ROUGE-L and conditional probability, while Model Utility combines ROUGE-L, conditional probability, and cosine similarity through harmonic means.All evaluation scores range from 0 to 1, with higher values better.
B.4 LLM-as-a-Judge
GRACE’s evaluations combine LLM-as-a-Judge scores with agreement analysis and additional sensitivity studies. These results support its utility-preserving performance and examine retrieval and clustering choices.
- Inter-Judge Agreement: The three judge models show strong agreement on five of six evaluation dimensions, especially for test-set metrics.Agreement is measured with Krippendorff’s α across DeepSeek V4-Flash, GLM 5.1, and Kimi k2.6.
- Inter-Judge Agreement: Judges consistently favor GRACE over RASLIK on five of six dimensions despite disagreement over whether meaningless repetition represents valid deviation.The disagreement concerns the Deviation Quality metric, whose interpretation differs across judges.
- LLM-as-a-Judge Evaluation: GRACE preserves model utility better than baselines, with higher average Probability and Cosine Similarity scores except for ROUGE-L.The comparison averages utility across the evaluated unlearning algorithms.
- Ablations and Sensitivity: A four-times forget candidate pool gives the highest retrieval accuracy and is selected as an efficiency–performance trade-off because larger pools increase NNOMP computation.Retrieval accuracy remains stable across the evaluated pool sizes without consistent improvement from larger pools.
C.3 Additional Results with LLaMA 3.2 1B
Additional experiments with LLaMA 3.2 1B Instruct reproduce GRACE’s earlier trends. GRACE outperforms the comparison methods in forget retrieval and achieves the highest utility in most algorithm–model combinations.
- Additional Results: With LLaMA 3.2 1B Instruct, GRACE outperforms Embedding and RASLIK in Forget Retrieval Accuracy.The additional experiments are reported in Table 15.
- Additional Results: GRACE consistently balances Forget Quality and Model Utility after unlearning on LLaMA 3.2 1B Instruct.The corresponding results are reported in Table 16.
- Additional Results: GRACE achieves the highest model utility in 5 out of 8 algorithm–model combinations.This result is part of the additional consistency study with LLaMA 3.2 1B Instruct.
- Evaluation: The utility evaluation compares selection mechanisms using quantitative metrics averaged over unlearning algorithms.Figure 5 describes the quantitative model-utility comparison, while Table 8 covers Forget Quality and Model Utility across models and datasets.