Deep Unlearning via Randomized Conditionally Independent Hessians

Task.

Given a set of input data $𝒮:{\left\{z_i\right\}}_{i=1}^n\sim 𝒟$ of size $n$, training simply identifies a hypothesis $\widehat{w}\in 𝒲$ via an iterative scheme $w_{t+1}=w_t-g\left(\widehat{w},z^{\prime }\right)$ until convergence, where $g\left(\cdot ,z^{\prime }\right)$ is a stochastic gradient of a fixed loss function. Once a model at convergence is found, machine unlearning aims to identify an update to $\widehat{w}$ through an analogous one-shot unlearning update:

for a given sample $z^{\prime }\in 𝒮$ that is to be unlearned.

Contributions.

We address several computational issues with existing approximate formulations for unlearning by taking advantage of a new statistical scheme for sufficient parameter selection. First, in order to ensure that a sample’s impact on the model predictions is minimized, we propose a measure for computing conditional independence called L-CODEC which identifies the Markov Blanket of parameters to be updated. Second, we show that the L-CODEC identified Markov Blanket enables unlearning in previously infeasible deep models, scaling to networks with hundreds of millions of parameters. Finally, we demonstrate the ability of L-CODEC to unlearn samples and entire classes on networks, from CNNs/ResNets to transformers, including face recognition and person re-identification models.

Definition 1 ($(ϵ,\delta )$– forgetting).

For all sets $𝒮$ of size $n$, with a “delete request” $z^{\prime }\in 𝒮$, an unlearning algorithm $𝒰$ is $(ϵ,\delta )$–forgetting if

In essence, for an existing model $w$, a good unlearning algorithm for request $z^{\prime }\in 𝒮$ will output a model $\widehat{w}$ close to the output of $𝒜\left(S\setminus z^{\prime }\right)$ with high probability.

Remark 1.

Definition 1 is similar to the standard definitions of differential privacy. The connection to unlearning is: if an algorithm is $(ϵ,\delta )$–forgetting for unlearning, then it is also differentially private.

If $𝒜$ is an empirical risk minimizer for the loss $f$, let

$\widehat{w}=\arg \min F(w)$ and $F(w)=\frac{1}{n}{\sum }_{i=1}^{n}f\left(w,z_i\right)$. Recall $g\left(z^{\prime }\right)$ from (1): our unlearning task essentially involves identifying the form of $g\left(z^{\prime }\right)$ for which the update in (1) is $(ϵ,\delta )$-forgetting. If an oracle provides this information, we have accomplished the unlearning task.

The difficulty, as expected, tends to depend on $f$ and $𝒜$. Recent unlearning results have identified forms of $f$ and $𝒜$ where such a $g\left(z^{\prime }\right)$ exists. The authors in [30] define $g\left(z^{\prime }\right)=\frac{1}{n-1}{H}^{\prime -1}\nabla f\left(\widehat{w},z^{\prime }\right)$ where

with additive Gaussian noise $w^{\prime }=w^{\prime }+N\left(0,{\sigma }^2\right)$ scaling as a function of $n$, $ϵ$, $\delta$, and the Lipschitz and (strong) convexity parameters of the loss $f$. We can interpret the update using (4) from the optimization perspective as a trajectory “reversal”: starting at a random initialization, the first order (stochastic gradient) trajectory of $w$ (possibly) with $z^{\prime }$ is reversed using residual second order curvature information (Hessian) at the optimal $\widehat{w}$ in (4), achieving unlearning. This is shown to satisfy Def. 1, and only incurs an additive error that scales by $O\left(\sqrt{d}/n^2\right)$ in the gap between $F\left(w^{\prime }\right)$ and the global minimizer $F\left(w^{\ast }\right)$ over the ERM $F(\widehat{w})$.

Rationale for approximate schemes.

From the reversal of $w$ optimization perspective, it is clear that there may be other choices to achieve unlearning. For a practitioner interested in unlearning, the aforementioned algorithm (as in (4)) can be directly instantiated if one has extensive computational resources. Indeed, in settings where it is not directly possible to compute the Hessian inverse necessary for $H^{\prime -1}\nabla f\left(\widehat{w},z^{\prime }\right)$, we must consider alternatives.

A potential idea.

Our goal is to identify a form of $g\left(z^{\prime }\right)$ that approximates $H^{\prime -1}\nabla f\left(\widehat{w},z^{\prime }\right)$. Let us consider the Newton-style update suggested by (4) as a smoothing of a traditional first order gradient step. The inverse Hessian is a weighting matrix, appropriately scaling the gradients based on the second order difference between the training set mean point $F(\widehat{w})$ and at the sample of interest $f\left(\widehat{w},z^{\prime }\right)$. This smoothing can also be seen from an information perspective: the Hessian in this case corresponds to a Fisher-style information matrix, and its inverse as a conditional covariance matrix [14,16]. It is not hard to imagine that from this perspective, if there are specific set of parameters that have small gradients at $f\left(\widehat{w},z^{\prime }\right)$ or if the information matrix is zero or small, then we need not consider their effect.

Examples of this intuition in vision.

[3,11,32] and others have shown that models trained on complex tasks tend to delegate subnetworks to specific regions of the input space. That is, parameters and functions within networks tend to (or can be encouraged to) act in blocks. For example, activation maps for different filters in a trained (converged) CNN model show differences for different classes, especially for filters closer to the output layer. We formalize this observation as an assumption for samples in the training set.

Assumption 1.

For all subsets of training samples $S\subset 𝒮$, there exists a subset of trained model parameters $P^{\ast }\subset \text{Θ}$ such that

Due to the computational issues discussed above, if we could make such a simple/principled selection scheme practical, it may offer significant benefits.

Naïve, Exact Unlearning.

A number of authors have proposed methods for exact unlearning, in the case where $(ϵ=0,\delta =0)$. SVMs by [23,28], Naïve Bayes Classifiers by [5], and k-means methods by [13] have all been studied. But these algorithms do not translate to stochastic models with millions of parameters.

Approximate Unlearning.

With links to fields such as robustness and privacy, we see more developments in approximate unlearning under Definition 1. The so-called $ϵ$-certified removal by [19] puts forth similar procedures when $\delta =0$, and the model has been trained in a specific manner. [19, 22] provide updates to linear models and the last layers of networks, and [15, 16] provide updates based on linearizations that work over the full network, and follow-up work by [14] presents a scheme to unlearn under an assumption that some samples will not need to be removed.

Other recent work has taken alternative views of unlearning, which do not require/operate under probabilistic frameworks, see [4, 25]. These schemes present good guarantees in the absolute privacy setting, but they require more changes to pipelines (sharding/aggregating weaker models) and scale unsatisfactorily in large deep learning settings.

A probabilistic angle for selection.

We interpret a deep network $𝒲$ as a functional on the input space $𝒟$. This perspective is common in statistics for variable selection (e.g., LASSO), albeit used after the entire optimization procedure is performed i.e., at the optimal solution. The only difference here is that we use it at approximately optimal solutions as given by ERM minimization. Importantly, this view allows us to identify regions in $𝒲$ that contain the most information about a query sample $z^{\prime }$. We will formalize this intuition using recent results for conditional independence (CI) testing. Finding $w_P$ above should also satisfy

This CI formulation is well studied within graphical models. Many measures and hypothesis tests have been proposed to evaluate it. The coefficient of conditional dependence (CODEC) in [1], along with their algorithm for “feature ordering”, FOCI, at first seems to offer a solution to (7), and in fact, can be implemented “as is” for shallow networks. (Review of other CI tests are in the appendix.)

Using CODEC directly for Deep Unlearning is inefficient.

There are two issues: First, when applying CODEC to problems with a very large $n$ with discrete values, the cost of tie-breaking for computing nearest neighbors can become prohibitive. Second, $z^{\prime }$ is not a random variable for which we have a number of instances. We defer discussion of the second issue to Section 5, and address the first issue here.

Consider the case where a large number of elements have an equal value. With an efficient implementation using $kd$−trees, identifying the nearest-neighbor as required by CODEC would still require expanding the nodes of all elements with equal value. As an example, if we are looking for the nearest neighbor to a point at the origin and there are a large number of elements on the surface of a sphere centered at the origin, we still require checking all entries and expanding their nodes in the tree, even when we know that they are all equal for this purpose.

Interestingly, this problem has a relatively elegant solution. We introduce a randomized version of CODEC, L-CODEC. For variables $A$, $B$, $C$:

where $\tilde{B}=B+N\left(0,{\sigma }^2\right)$, and similarly for $\tilde{C}$, $\tilde{A}$. This additive noise can simply be scaled to the inverse of the largest distance between any points in the set. By requiring this noise to be smaller than any distance between items in the set, the ranking will remain the same between unique discrete values, and will be perturbed slightly for equal ones. In expectation, this will still lead to the true dependence measure. The noise addition is consistent with the Randomization criterion for conditional independence – for random variables $A$, $B$, $C$ in Borel spaces, $A\perp B\mid CiffA\stackrel{\text{a.s.}}{=}h(B,U)$ for some measurable function $h$ and uniform random variable $U\sim \text{Uniform}(0,1)$ which is independent of $(B,C)$ as in [26].

Remark 2.

An altered version of this setup also gives us a form of explainability, where we can apply sensitivity analysis to each input feature or pixel and estimate its effect on the output via a similarly randomized version of the Chatterjee rank coefficient $T(A,B)$, proposed by [8].

4.1. Efficient Subset Selection that is also Sufficient for Predictive Purposes

The above test is good for (7) if we know which subset $P\in 𝒫(\text{Θ})$ to test. Recent work by [36] proposes a selection procedure using an iterative scheme to slowly build the sufficient set, adding elements which maximally increase the information explained in the outcome of interest. While it is efficient (polynomial in size), we must know the maximal degree. A priori, we may have no knowledge of what this size is, and for parameter subsets it may be very high.

When using L-CODEC, we can use a more straightforward Markov Blanket identification procedure adapted from [1]. FOCI more directly selects which variables are valuable for explaining $z^{\prime }$, and in fact, is proven to identify the sufficient set (Markov Blanket) with a reasonable number of samples. Briefly, in our L-FOCI, the sufficient set is built incrementally with successive calls to L-CODEC, moving the most “dependent” feature from the independent set to the sufficient set. See appendix for details.

Computational Gains.

A direct observation is that now we are doing sampling, which adds a linear computational load. However, directly updating all parameters requires $O\left(d^3\right)$ computation due to matrix inversion, while this procedure requires $O\left(md+dm\log m+p^3\right)$, for the forward passes, FOCI algorithm, and subsequent subsetted matrix inversion. For any reasonable setting, we have $p\ll d$, and so this clearly offers significant practical advantages.

5.1. Theoretical Analysis

By definition, any neural network as described above is actually a Markov Chain: we know that the output of a layer is conditionally independent of the penultimate one given the previous one, and clearly a change in one layer will propagate forward through the rest of the network. However, when trained for a task with a large number of samples, the influence or “memory” of the network with respect to a specific sample may not be clear. While the output of the layers may follow a Markov Chain, the parameters in the layers themselves do not, and their influence on a sample through the forward pass may be highly dependent or correlated. Practically, we would hope that unlearning samples at convergence does not cause too much damage to the model’s performance on the rest of the input samples. Following traditional unlearning analysis, we can bound the residual gradient norm to relieve this tension.

L-CODEC Evaluation.

To assess speedup gained in the discrete setting when running L-CODEC, we construct the Markov Blanket for specific attributes provided as side information with the CelebA dataset. Fig. 3 shows the wall-clock times for Markov Blanket Selection via FOCI and L-FOCI for each attribute.

Markov Blanket Identification.

We replicate the experimental setup in Sec 5.3 of [36], where a high dimensional distribution over a ground truth graph is generated, and feature mappings are used to reduce the dimension and map to a latent space. Table 1 summarizes subset identification efficacy and runtime. Replacing conditional mutual information (CMI) with L-CODEC, we see a clear improvement in both runtime and Markov Blanket identification over the raw data, and comparable results in the latent feature space. Using L-FOCI directly in the feature space, we identify an additional spurious feature not part of the Markov Blanket, but runtime is significantly faster.

Spurious Feature Regularization.

This Markov Blanket $(MB)$ identification scheme can be used to address spurious feature effects on traditional NN models. A straightforward approach would be to directly add a loss term for each potentially important feature over which we would like to regularize, $𝓛(\theta )+\sum _{S\in 𝒮}{R}_S(\theta )$. However, with a large number of outside factors $S$, this can adversely effect training. We instead use L-FOCI to identify the set of minimal factors that, when conditioned, make the rest conditionally independent. Then it is only necessary to include regularizers over $S\in MB(Y)$.

We evaluate a simple attribute image classification setting using the CelebA dataset. We run L-FOCI over the attributes as in our L-CODEC evaluation, and regularize using a Gradient Reversal Layer for a simple accuracy term over those attributes. Results in Fig. 4 clearly show that selection with FOCI provides the best result, maintaining high overall accuracy but also preserving high accuracy on sets of samples with/without correlated attributes.

7.1. Compare to Full Hessian Computation

For simple regressors, we can compute the full Hessian and compare results generated by a traditional unlearning update, our L-FOCI update, and a random selection update. To reduce variance and show the best possible random selection, we run our L-FOCI and randomly choose a set of the same size for each random selection. Fig. 5 (left) shows validation and residual accuracies for 1000 random removals from MNIST (average over 10 runs).

7.2. Removal in NLP models

We now scrub samples from transformer based models using LEDGAR [34], a multilabel corpus of legal provisions in contracts. We use the prototypical subset which contains 110156 provisions pertaining to 13 most commonly used labels based on frequency. Our model is a fine-tuned DistilBERT [29] and uses the $[CLS]$ token as an input to the classification head. Table. 7 shows results of scrubbing the provisions from two different classes; Governing Laws and Terminations which have the highest/lowest support in the test set. As expected with increasing $ϵ$, i.e., lower privacy guarantees, we can support more number of removals based on the Micro F1 score of the overall model. The Micro F1 scores, for the removed class fall off rapidly, while the change in overall scores is more gradual.

7.3. Removal from Pretrained Models

The above settings show settings where a sample from one specific source may be removed. A more direct application of unlearning is completely removing samples from a specific class; a compelling use case is face recognition.

We utilize the VGGFace dataset and model, pretrained from the original work in [21, 27]. The model uses a total of approximately 1 million images to predict the identity of 2622 celebrities in the dataset. Using a reconstructed subset of 100 images from each person, we first fine-tune the model on this subset for 5 epochs, and use the resultant models as estimates of the Hessian. In this setting, the VGGFace model is very large, including a linear layer of size 25088 × 4096. Selecting even a few slices from this layer results in a Hessian matrix unable to fit in typical memory. For this reason, we run a “cheap” version of L-FOCI: we select only one slice that results in the largest conditional dependence on the output loss.

Fig. 7 (left) show results for scrubbing consecutive images from one individual in the dataset for a strong privacy guarantee of $ϵ={10}^{-5}$. As the number of samples scrubbed increases, the performance on that class drops faster than on the residual set, exactly as desired.

7.4. Removal from Person re-identification model

As a natural extension to our experiments on face recognition, we evaluate unlearning of deep neural networks trained for person re-identification. Here, the task is to associate the images pertaining to a particular individual but collected in diverse camera settings, both belonging to the same camera or from multiple cameras. In our experiments, we use the Market-1501 dataset [38] and a Resnet50 architecture which was trained for the task. We unlearn samples belonging to a particular person, one at a time, and check the performance of the model. Experimental results are in agreement with results reported for the transformer model as well as the VGGFace model. With very small values of $ϵ$ i.e. 0.0005 the number of supported removals is limited to less than 10 depending on the person id being removed. However, with a larger value of $ϵ$, e.g., 0.1, all potential samples can be removed without a noticeable degradation in model performance in terms of mAP scores. In Fig. 8, we clearly see that after scrubbing a model for a particular person, its predictions for that particular individual become meaningless whereas the predictions on other classes are still possible with confidence, as desired. Additional experiments with different datasets, model architectures and other ablations for deep unlearning for person re-identification models are presented in the appendix.

Social Impact.

Indiscriminate use of personal data in training large AI models is ethically questionable and sometimes illegal. We need mechanisms to ensure that AI models operate within boundaries specified by society and legal guardrails. As opt-out laws get implemented, compliance on the service-provider end will entail costs. While our contributions cannot guarantee perfect forgetting, with additional validation they can become a part of a suite of methods for unlearning.