Disentangled Knowledge Forgetting in Machine Unlearning

1 Introduction

A growing number of regulations have been introduced to grant users the right to be forgotten. A straightforward yet effective approach is to retrain from scratch, nevertheless suffering from computation overhead. Machine unlearning (MU) thus appears to remove the forgetting data (effectiveness), while preserving model performance on the remaining data (fidelity).

Existing methods address MU with either exact or approximate unlearning. The former [Bourtoule et al., 2021; Graves et al., 2021] retrains with optimized efficiency, still being costly in computation. The latter approximates the retrained model to relieve burden from retraining, where several works [Mehta et al., 2022; Foster et al., 2024] estimate the influence of forgetting data to update parameters, and others adopt knowledge distillation (KD) [Chundawat et al., 2023; Kurmanji et al., 2023] to relearn mis-forgotten knowledge. Despite achieving notable advances, they struggle to simultaneously avoid overly unlearning with low fidelity and incomplete unlearning with low effectiveness (see Fig. 1). After forgetting, model may inadvertently forget knowledge of the remaining category, or revive part of should-forgotten knowledge, that is the shared knowledge between forgetting and remembering categories. For example, the shared knowledge between tiger and cat like 'haired' and 'quadruped' lead to incomplete unlearning on tiger, since the resident feature-tocategory mappings could be coincidentally activated by utilizing shared knowledge features.

In Fig.2(a) and (b), we look into causal relationship of overly and incomplete unlearning using Structural Causal Model(SCM) [Pearl et al., 2016] respectively, including shared knowledge K, forgetting data Df , remaining data Dr, feature representations E, and prediction labels Y . Machine unlearning expects to forget Df for risk-free prediction Yr = Y - Yf . Specifically, K → Df and K → Dr represent that Dr and Df share knowledge, Df → E and Dr → E represent that E is extracted from Dr and Df , Df → Y and Dr → Y represent causal effects from images to labels, and E → Y represents prediction. In overly unlearning, MU removes K → Dr, inadvertently forgetting shared knowledge of the remaining data to induce model collapse on Dr. In incomplete unlearning, MU maintains K → Df , hindering accurately forgetting knowledge.

Since SCM implies the confusing knowledge in K as the confounder, we delve into approximate unlearning with KD to identify key path for a balanced MU, regarding MU as a two-level classification problem: (1)multi-category prediction for each remembering class; (2)binary prediction for the forgetting class and all the remembering classes. Proved by Sec 3.2, MU should be further regulated with the separated remembering loss (RL) and forgetting loss (FL), which respectively constricts knowledge of fidelity and effectiveness to aid in removing K for balanced MU. RL roughly form boundary with the remembering classes for high fidelity, whereas FL explicitly and rigidly reforms the boundary with the forgetting for high effectiveness. By contrast, traditional MU loss is coupled to emphasize knowledge about fidelity and effectiveness, failing in alleviating overly and incomplete unlearning.

Guided by the separated constraints of RL and FL, K prevents the independent optimization of these losses, requiring attribute-level semantic disentanglement for K to make it fit with their precise optimization. As identifiable semantic units, we argue the disentangled knowledge can be carefully refined via counterfactual data to enforce rigid enough boundary across semantically overlapped samples, thus separately and precisely minimizing RL and FL to guarantee high fidelity and effectiveness. Take Fig.3 to illustrate the disentanglement of the shared knowledge with counterfactual inference, where we first establish decision boundary between the remaining and forgetting data by first generating counterfactual samples near the decision boundary, and then shifting to improve rigidness of boundary by containing all forgetting data. Representing Dc as counterfactual data and E′ as the disentangled feature representation, we thus build K → Dc → E′&Y to replace K → Df → Y in Fig.2(c), effectively revealing causal effects from the disentangled knowledge to prediction in a resolvable form.

Specifically, we first establish K → Dc → Y via β-VAE to disentangle knowledge as attribute-level semantics. Since counterfactual inference provides the imagination ability to analyze confounds [Niu et al., 2021], we bravely imagine non-existent samples to form Dc based on disentangled attribute-level knowledge, which is represented as Dc → E′ . For example, based on the shared attribute knowledge

'haired' extracted from tiger and cat, we add new concept 'winged' from another known class like eagle to imagine virtual samples for learning, which could locate near the original decision boundary between tiger and cat to promote for the balanced MU. Regarding counterfactual and remembering samples as key anchors, contrastive learning is further adopted to pull positive pairs close and push negative pairs apart, thus gradually shifting boundary with high rigidness. The contributions are:

• Theoretically: We offer an insight view in MU, where causal and loss analysis guide route of disentangled knowledge to guarantee high fidelity and effectiveness.
• Methodologically: We propose disentangled knowledge forgetting, which benefits from variational autoencoder to disentangle knowledge towards independent optimization of RL and FL, and counterfactual inference to refine knowledge constrained by RL and FL.
• Empirically: Extensive experimental results demonstrate our method achieves state-of-the-art performance.

2 RelatedWork

Exact unlearning accelerates retraining by training submodels [Bourtoule et al., 2021; Yan et al., 2022] or subtracting gradients in training phases [Graves et al., 2021]. They require extensive computational cost, when facing large-scale AI systems or frequent data elimination requests.

Approximate unlearning has emerged as a research focus, focusing on reducing the influence of forgetting data. Early works directly modify model weights via theoretically verified deletion [Golatkar et al., 2020], noisy Newton update [Golatkar et al., 2020], or Fisher information matrix [Foster et al., 2024]. However, they generally suffer from locally slow convergency. Coupled with LLMs and generative model, several optimization-based methods alleviate convergency problem and address advanced challenges like hallucinations [Yuan et al., 2025], correlated knowledge persistence [Wei et al., 2025], and concept revival [George et al., 2025]. With further interest on MU deployment constraints, recent works introduce source-free Hessian estimation [Ahmed et al., 2025], certified noisy finetuning [Koloskova et al., 2025], and logits tempering [Spartalis et al., 2025] to efficiently approximate the gold standard of retraining.

3 Preliminaries

3.1 Task Definition

Let D = {xi, yi}Ni=1 be training data containing N examples, where xi is the ith sample, yi ⊆ {1, ...,K} is the corresponding class label, and K is the total number of classes. We define Df ⊆ D as subset of the training data to forget, and Dr = D \ Df as the remaining data to retain. Defining the original model trained on D as hw0 with parameters w0, we optimize model with cross-entropy loss Lce:

(1)

Regarded as the optimal unlearning model, hw∗ is the model retrained with Dr. MU aims to derive a new model hwu to forget the information in Df by updating parameters fromw0 to wu, ensuring hwu to be similar with hw∗ .

3.2 RethinkingMachine Unlearning Loss

Inspired by DKD [Zhao et al., 2022], we reformulate unlearning loss with remembering loss (RL) and forgetting loss (FL). For a forgetting sample x ∈ Df , denoting classification probability as p = [p1, p2, . . . , pf , . . . , pK] ∈ R1×K , where pi is probability to be classified as ith class and K is number of classes. Each element in p is obtained via softmax function:

where zi represents logit of the ith class. To separate the predictions as relevant and irrelevant to the forgetting class, we define binary probabilities b = [pf , p\f ] ∈ R1×2 as

Meanwhile, we define ˆ p = [ˆ p1, . . . , ˆ pf-1, ˆ pf+1, . . . , ˆ pK] ∈ R1×(K-1) to independently model the probabilities among K - 1 classes as

Unlearning loss L is defined as the KL-divergence between target probability distribution p and the grounded unlearning distribution q, which can be written as

According to Eq. 3, we define normalized probabilities ˆ pi = pi/p\f , which changes Eq. 5 as

Hence, L can be reformulated as :

KL(pf ∥ qf ) defines the discrepancy to minimize pf near 0 in forgetting class distribution, and KL(ˆ p∥ ˆ q) computes the normalized distribution gap of remaining classes weighted by 1 - pf , thus preserving irrelevant knowledge to the forgetting. Essentially, MU can be regarded as multi-task learning of forgetting and remembering [Lin et al., 2024], where loss of each task jointly constrains to learn the specific knowledge.

However, due to the coupling term (1-pf ), the updates for these two tasks during training become entangled and biased, often causing one task to dominate the model update at the expense of the other. Specifically, an over-emphasis on the remaining loss may compromise the forgetting efficacy (incomplete unlearning), whereas a dominant forgetting loss can degrade the model's classification performance on remaining classes (low fidelity). To resolve such conflict, it's essential to decouple the optimization of these two objectives, where the forgetting and remembering goals should be treated as independent constraints, allowing the optimization to explicitly balance their gradients and achieve dual optimality in both effectiveness and fidelity.

4 ProposedMethod

4.1 Overview

Built on knowledge distillation framework for unlearning, we illustrate in Fig. 4 with knowledge disentangle and refine. Regarding the original model hw0 as teacher network, we train a student network hwu to remember Dr and forget Df .

During disentanglement, we adopt β-Variational AutoEncoder(β-VAE) to extract the shared attribute knowledge of forgetting data Sf and unique attribute knowledge of remaining data Ur. Specifically, we feed forgetting samples Xf and remembering samples Xr into two learnable encoders Qϕ and Qψ to extract Sf and Ur, respectively. Afterwards, Sf is supervised via a classifier θo to predict the original output O, and Ur is used to predict the label Y via classifier θy, where both designs ensure θo and Ur capture correct semantics for knowledge disentanglement.

During refinement, we adopt schemes of counterfactual inference and contrastive learning. First, we generate counterfactual samples Xc with [Sf ,Ur]. Then, we minimize the discrepancy between Xf and Xc by optimizing the counterfactual loss. Regarding counterfactual and remembering samples as key anchors, contrastive learning is finally performed constrained by RL and FL, simultaneously pulling the embeddings of positive pairs close and pushing the embeddings of negative pairs apart.

4.2 Knowledge Disentangle Towards Independent Optimization with β-VAE

Rethinking the MU loss reveals that while FL and RL aim at distinct objectives, the shared knowledge K acts as a confounder that prevents their independent optimization. Such entanglement necessitates a shift from logit-level manipulation to attribute-level semantic disentanglement for precise optimization. With constraints of independent MU loss, we further remove confounder K by disentangling knowledge with aid of the generated counterfactual data Dc following the guidance in Fig.2(c). However, generating Dc encounters two difficulties. First, sampling Dc from real-world images is impractical, since Dc is defined as a mixture of multiple attributes with rare relations. Secondly, generating Dc with individual attributes is unreachable as well, since both class-unique attributes U and shared attributes S are inherently unobservable and entangled in sample-level semantic space. Therefore, we disentangle knowledge with Variational Auto-encoders(VAEs), which separates the unobservable and entangled semantic attributes into two groups of distinct latent variables, thus being observable as individual samples with no confounders for further counterfactual generation.

In Fig. 4, we design an encoder-decoder structure to infer the unobservable S and U with two observable constraints, i.e., the observation context O and class label Y . Specifically, we first decompose each sample X as 1) unique attributes U corresponding to the intrinsic and class-related information to predict label Y , and 2) shared attributes S with attributes that shared across different classes. Such decomposing is achieved by extract S and U from samples with encoders ϕ and ψ respectively, where ϕ encodes shared attributes as Qϕ(S|X), and ψ encodes class-unique attributes as Qψ(U|X). Then, we disentangle the joint distribution of entangled attributes as P(S, U|X) = Qϕ(S|X)Qψ(U|X) based on the Markov condition between these factors:

where P (X, S, U, O, Y ) denotes the joint distributions of all factors, P(X|S,U) denotes the reconstruction output by decoder Pθ, P(O|S) is predicted by classifier θO, and P(Y |U) is predicted by classifier θS.

To ensure disentangle is constrained by P(X, S, U, O, Y ) with D, we optimize by maximizing a weighted Evidence Lower Bound(ELBO)[Yue et al., 2021]:

where the first two terms supervise the classification accuracy of θO and θS, the third term improves the reconstruction quality by generating the original samples, and β is the weighting parameter. Essentially, ELBO optimizes the disentangling in a self-supervised manner, where ϕ is supervised by θO to focus on observation context of sample X, regardless of its belonging class. Meanwhile, ψ is supervised by θS to focus on class-unique semantics, which excludes the confounders from the shared knowledge for further counterfactual generation.

4.3 Knowledge Refine Constrained by FL and RL with Counterfactual Inference

Based on the disentanglement results Qϕ(S|X) and Qψ(U|X), we employ them to generate Dc by counterfactual inference, where we further use Dc to refine knowledge.

Counterfactual Generation. For each forgetting sample xf ∈ Df , we generate counterfactual sample xc ∈ Df by imagining a counterfactual reasoning case like "Given a tiger with attribute haired, what would it like with wings? " Such case is conducted with three typical steps, i.e., Abduction, Action, and Prediction. First, we sample the shared attributes sf ∈ Sf from disentangled S ∼ Qϕ(S|X = xf) as Abduction, which defines a haired tiger. By adding a unique attribute ur ∈ Ur from remaining data Dr which is sampled from U ∼ Qψ(U|X = xr), we then utilize Action to borrow the unique winged attribute from bird. With both sf and ur, we finally imagine a counterfactual sample xc by decoding Pθ(X|S = sf , U = Ur) as Prediction.

With the generated Dc, we refine knowledge with dual objectives of remembering and forgetting. Remembering establishes a new causal route K → Dc → Y to resolve the confounders for Dr. Since multiple counterfactual samples xc ∈ Dc and forgetting sample xf ∈ Df shares attributes S, we design counterfactual loss to keep S in forgetting sample:

where hw0 (·) and hwu (·) denotes the output of original and unlearned model, respectively.

By leveraging the decoupled Sf and Ur in counterfactual samples, minimizing Lf aligns Uf with the manifold of Ur while keeping Sf invariant. Such mechanism implicitly substitutes Df with Dc, thereby effectively eliminating the confounder. The optimization objective of Lf implies:

Optimizing Lf is equivalent to minimizing (FL). This drives the forgetting class probability (pf ) to zero, which enforces FL constraint to guarantee the effectiveness.

Contrastive Learning. By establishing K → Dc → E′ , we refine knowledge to form its disentangled feature representation E′ . Inspired by the impressive performance of contrastive learning [Chen et al., 2020] to boost representation capability, we utilize Dc to refine E′ by comparing both forgetting and remembering samples. Specifically, we firstly take (xc, xf) as positive pairs, where Px(xf) = {xp|xp ∈ Dc} and Nx(xf) = {xn|xn ∈ Dr, yn ̸= yf} are formed as negative pairs. We then pull close the positives to encourage Uf being varying to the random anchors Ur, meanwhile we push away negatives to enforce Uf shifting from its original representation through contrastive loss:

where τ is a scalar temperature parameter, and z = hwu (x) denotes the extracted feature representation.

Minimizing the contrastive loss Lc preserves the intrinsic topology of the representations for the remaining data. Such structural stability implies that the normalized prediction distribution over the remaining classes is maintained. This alignment effectively minimizes the KL divergence between the unlearning model and the original model on the remaining data, thereby explicitly satisfying RL constraint and guaranteeing high fidelity on the remaining data.

Considering the causal route of K → Dc → E′&Y , the overall loss for knowledge refinement is L = Lf + λLc, where λ is designed to balance remembering and forgetting.

5 Experiments

5.1 Experimental Setup

Unlearning Settings. We compare performance with several baselineMUmethods in three scenarios, i.e., single-class forgetting, subclass forgetting and multi-class forgetting.

Datasets andModels. Following [Chundawat et al., 2023; Chen et al., 2023], we evaluate with three image classification datasets: CIFAR-10, CIFAR-20, CIFAR-100 [Krizhevsky et al., 2009]. ResNet-18 [He et al., 2016] and ResNet-34 [He et al., 2016] are used as backbone network.

Baseline. The adopted baselines are: Retrain, Finetune, Negative Gradient [Golatkar et al., 2020], Random Labels [Golatkar et al., 2020], BadT [Chundawat et al., 2023], Boundary [Chen et al., 2023], SSD [Foster et al., 2024], L1-Sparse [Jia et al., 2023], ERM-KTP [Lin et al., 2023], SCRUB [Kurmanji et al., 2023], GDR-GMA [Lin et al., 2024], SalUn [Fan et al., 2024].

Metrics. Following [Chundawat et al., 2023; Chen et al., 2023], we use 6 metrics to evaluate unlearning: (1) AccDr , the classification accuracy on remaining data; (2) AccDf , the classification accuracy on forgetting data; (3) Accval, the classification accuracy on validation dataset; (4) Membership inference attack (MIA), examining whether the forgetting samples are in the training set; (5) Average Gap (Avg.Gap), performance gap comparing with Retrain; (6) RTE offers the consuming time of unlearning. Note that the better performance of an MU method corresponds to the smaller performance gap with Retrain.

5.2 Experimental Evaluation

Comparison experiments for single-class forgetting. Following settings of [Chundawat et al., 2023; Chen et al., 2023], Table 1 presents the comparison results for singleclass forgetting on CIFAR-10 and CIFAR-100 dataset. As the most vital metrics, our method achieves the lowest performance gaps against Retrain, reaching 1.32 on CIFAR-10 and 1.46 on CIFAR-100 respectively. Some methods may be the strongest when considering only a single metric, but this comes at the cost of sacrificing the other metrics. For example, Salun suffers from low fidelity by reaching 3.90% lower in AccDr than Retrain, while Finetune fails in privacy scenario proved by gaining 15.84% MIA gap against Retrain. These results demonstrate that our idea of attributelevel semantic disentanglement enables precise unlearning, maintaining high fidelity without compromising effectiveness across most application scenarios.

Comparison experiments for subclass forgetting. Following settings of [Chundawat et al., 2023; Foster et al., 2024], we compare performance of subclass forgetting in Table 2, where the experiment of forgetting one sub-class per hyper-class on CIFAR-20 dataset is challenging, due to significant semantic overlap in the embedding space. While most methods suffer from severe performance degradation, our method achieves close performance with Retrain, achieving the smallest average gap of 2.19 against Retrain. Such phenomenon ensures our method is capable to balance forgetting and remembering, even with high semantic interference.

Comparison experiments for multi-class forgetting. Table 3 presents the results of comparison experiments by forgetting multiple classes on CIFAR-100 dataset, i.e., 1, 5, 10 and 20 classes. while baseline performance degrades significantly as the number of forgetting classes gets larger, our method demonstrates superior scalability with the lowest Avg.Gap. In the highly challenging 20-class scenario, our method gains 95.80% AccDr , surpassing 92.50% achieved by SCRUB. Such phenomenon proves that our method with attribute-level knowledge disentanglement effectively resolves the confound of the shared knowledge, especially facing complex scenarios of forgetting multiple classes.

5.3 Ablation Study

Different classes to be forgotten. In Table 4, we present ablation results on CIFAR-10, where the forgetting class Df is varied across 10 independent runs. Our method performs consistently well across all runs, regardless of the semantic overlap of the chosen class. Such insensitivity experiment highlights the robustness and applicability of our method.

Impact of λ. Fig. 5(a) shows the ablation results with different λ values, where λ is the weight of the overall loss. We can observe that a larger λ reduces the Avg.Gap, meanwhile removing it (λ = 0) increases the Avg.Gap by 0.59, demonstrating its effectiveness in establishing a rigid decision boundary. However, an excessive value of λ = 2 results in a significant performance drop with the rising Avg. Gap of 3.19, due to the over-fitting with the contrastive constraints. The optimal configuration is empirically determined to be λ = 1, which achieves the minimum Avg. Gap comparing with Retrain.

Impact of β. In Fig. 5(b), we compare different values of β, where β is the penalty factor in β-VAE loss. A lower β fails to fully remove the confounding knowledge, resulting in an incomplete unlearning with a high AccDf of 5.82%. Conversely, an excessive value of defining β = 8 overcompresses features for discriminative representation, causing AccDr to greatly drop. Therefore, we defined β = 6 to achieve the minimum Avg.Gap.

5.4 Qualitative Analysis

Visualization of Decision Space. As shown in Figure 6, we visualize the decision space by projecting with t-SNE, where each dot denotes a sample and the corresponding color indicating its belonging classes. Before unlearning in (a), all classes are learned as clusters in decision space, with clear boundary separating each other. When removing "deer" during training, its corresponding samples are misclassified in (b) achieved by Retrain. In Fig.6(c), the forgetting data has been predicted as the nearest classes after boundary shifting, even the cluster does not entirely spread out like Retrain.

Visualization of Forgetting and Remembering. Following [Chen et al., 2023] to test capability of privacy preserving, we apply machine unlearning on face recognition task [Cao et al., 2018]. As shown in the supplementary material, our method outperforms all other methods and achieves the smallest average performance gap against Retrain model. To further illustrate in Fig.7, we employ Grad-CAM [Selvaraju et al., 2017] to visualize model's attention areas. The original model focuses on facial regions in (a), while the Retrain model and our method effectively adjust the recognition focus on areas out of face, thus successfully forgetting individual's privacy. Meanwhile, both methods maintain strong attention on relevant facial features of the remembering cases in (b), demonstrating high fidelity achieved by our method comparable to the Retrain model.

6 Conclusion

In this paper, we propose disentangled knowledge forgetting to address overly and incomplete unlearning. Extensive experimental results demonstrate that our method achieves state-of-the-art performance.

Acknowledgments

This work was supported in part by the National Key R&D Program of China under Grant 2023YFC3006501, the Natural Science Foundation of Jiangsu Province of China under Grants BK20242050 and BK20251504, the National Natural Science Foundation of China under Grant 6250074646, the High Performance Computing Platform, Hohai University, the Fundamental Research Funds for the Central Universities under Grant B250201042, and the Open Research Project of State Key Laboratory for Novel Software Technology, Nanjing University, under Grant KFKT2025B04.

Contribution Statement

Yuhang Xia and Cheng Zhen contributed equally to this work. Yirui Wu is the corresponding author.