Rewind-to-Delete: Certified Machine Unlearning for Nonconvex Functions

In plain words

Sometimes a trained model has to forget some of its training data, for example when a person asks for their data to be removed. Retraining from scratch without that data is the honest fix, but for large models it is very expensive. The paper asks for a cheaper way that still comes with a mathematical promise: the changed model should look almost the same as one that was never shown the data.

The idea is to keep an earlier snapshot from training. To delete data, go back to that snapshot and resume training on the data that remains. Then add a little random noise to the weights. Going back further costs more computation but needs less noise. The method only needs ordinary gradient-descent training, so the training itself does not have to change. The authors test it on hospital stay data and on face images for a gender task. They report that, on the face task, it removes more of the removed users' influence than the other certified methods, and that its training is far faster than the second-order method. These are the authors' own results.

TL;DR

Rewind-to-Delete (R2D) is a certified unlearning method for nonconvex models trained with plain gradient descent, with no changes to training §1. To delete, it reloads an earlier checkpoint θ_{T−K}, runs K gradient steps on the retained data, and adds Gaussian noise whose scale σ is set by eq. (3) alg. 2 §3. Author-reported: on Lacuna-100 it lowers unlearned-set AUC more than Hessian-Free or CNS at σ=0.01 tab. 2, and on eICU it trains up to about 88× faster than Hessian-Free tab. 6. The released code does not compute the noise scale, so the certified noise cannot be reproduced from the release; see Paper vs code.

Problem

• Certified unlearning asks that the output of the unlearning step be (ε, δ)-indistinguishable from the output of training on the retained data alone §2, Def. 2.2.
• Nonconvex losses have many minima, so the classic descent-to-delete argument does not carry over. D2D is first-order but is proven only for strongly convex losses §1.1 §2.
• Existing certified methods for nonconvex losses either use second-order steps (Constrained Newton Step; Hessian-Free) or need noise injected at every training step (Langevin unlearning) §1 tab. 1. The paper's storage and training-change comparison is in tab. 1 and app. E.
• Earlier experiments mostly forget one whole class or uniformly random samples. The paper argues that real removal requests remove users, which is not i.i.d. with the training or test sets §1 §4.1.

Method

R2D has two phases. Training runs T steps of gradient descent with step size η on the average loss over the dataset D, and keeps the iterate θ_{T−K} as a checkpoint alg. 1. The released weights are θ_T with added Gaussian noise alg. 1. For the analysis, training is full-batch gradient descent §2 app. C.

Unlearning a set Z of m points gives D' = D \ Z. It loads θ_{T−K}, takes K gradient steps on the loss over D', adds fresh Gaussian noise, and serves the result alg. 2. Rewinding further (larger K) lowers the noise scale. K = T gives a noiseless full retrain §3.

If no checkpoint was saved, θ_{T−K} can be rebuilt from θ_T by K backward steps. Each backward step is a proximal point step on −f_D, which is well defined when η < 1/L §2, Lemma 2.3 alg. 3 app. A.1. Under the backward-Euler relation (eq. 1 and 2 in §2) the earlier iterate is recovered exactly from the later one, and the paper uses this to justify the backward steps §2.

The noise scale and the bound on the rewind length come from the Gaussian mechanism applied to the gap between the two trajectories §3, Thm. 3.1 app. A.2.

# Learning: T gradient steps, keep a checkpoint, release noisy weights  [alg. 1]
# Input: dataset D (n examples), initial weights, eta, T, K, sigma
def learn(data, initial_weights, eta, T, K, sigma):
    # Run T steps of gradient descent on the average training loss  [alg. 1]
    weights = initial_weights
    for step in range(1, T + 1):
        weights = weights - eta * grad_loss(weights, data)
        # Save theta_{T-K} as the rewind point  [alg. 1]
        if step == T - K:
            checkpoint = weights
    # Release theta_T with fresh Gaussian noise  [alg. 1]
    noisy = weights + gaussian_noise(sigma, weights.shape)
    return checkpoint, noisy

# Unlearning: rewind, take K steps on retained data, add noise  [alg. 2]
# Input: retained data D' = D minus Z, checkpoint theta_{T-K}
def unlearn(retained, checkpoint, eta, K, sigma):
    # Start from the saved iterate, not from the final weights  [alg. 2]
    weights = checkpoint
    for step in range(K):
        # Gradient steps on the loss over the retained data  [alg. 2]
        weights = weights - eta * grad_loss(weights, retained)
    # Perturb the result before it is used for inference  [alg. 2]
    return weights + gaussian_noise(sigma, weights.shape)

# Rebuild theta_{T-K} from theta_T when no checkpoint was saved  [alg. 3]
# Each backward step solves argmin_x { -f_D(x) + ||x - w||^2 / (2 eta) }
def reconstruct_checkpoint(data, final_weights, eta, K):
    # Start from the trained weights  [alg. 3]
    anchor = final_weights
    for step in range(K):
        # Approximately solve one proximal subproblem  [alg. 3] [§2]
        anchor = solve_prox(neg_loss(data), center=anchor, eta=eta)
    return anchor

Key equations, as stated in the paper:

Noise scale for certified unlearning (Thm. 3.1 assumes 0 < ε ≤ 1 and δ > 0):

\sigma = \frac{2\, m\, G \, h(K)\,\sqrt{2\log(1.25/\delta)}}{L\, n\, \epsilon}, \qquad h(K) = \bigl((1 + \tfrac{\eta L n}{n-m})^{T-K} - 1\bigr)(1+\eta L)^K eq. 3 §3

Step size condition for the same theorem: η ≤ min{1/L, n / (2(n−m)L)} §3, Thm. 3.1.

Utility bound on the average gradient norm (O(·) hides η, G and L) eq. 4:

\frac{1}{T}\Big(\sum_{t=0}^{T-K-1}\|\nabla f_{D'}(\theta_t)\|^2 + \sum_{t=0}^{K-1}\|\nabla f_{D'}(\theta''_t)\|^2 + \mathbb{E}\|\nabla f_{D'}(\tilde\theta'')\|^2\Big) \le O\Big(\sigma^2 + \frac{n}{T(n-m)} + \frac{(T-K-1)\,m}{T(n-m)}\Big) eq. 4

Rewind length that reaches a given ε for a fixed σ (Cor. 3.2) eq. 5:

K(\epsilon) \le \frac{1}{\log\!\big(1+\frac{\eta L n}{n-m}\big)} \log\!\Big(\big(1+\tfrac{\eta L n}{n-m}\big)^{T} - \frac{\sigma L n \epsilon}{2mG\sqrt{2\log(1.25/\delta)}}\Big) eq. 5 §3

PL condition with parameter μ > 0 (Def. 3.3) eq. 6:

\tfrac{1}{2}\|\nabla f(x)\|^2 \ge \mu\,\bigl(f(x) - f^*\bigr) eq. 6 §3

Population-risk bound for PL losses (Cor. 3.4, unnumbered display in §3):

\mathbb{E}[F(\tilde\theta'')] - F^* \le L\sqrt{d}\,\sigma + \frac{2G^2}{(n-m)\mu} + \frac{L}{2\mu}(1-\eta\mu)^K\Big[\big(1-\tfrac{\eta\mu(n-m)}{n}\big)^{T-K} \big(f_{D'}(\theta_0) - f^*_{D'}\big) + \frac{Gm(G+L\eta)}{\mu(n-m)}\Big] §3, Cor. 3.4

Hyperparameters

Values from the paper unless marked otherwise. Released scripts are cited where they state a value.

Results

Author-reported. Setup: eICU (MLP, binary length-of-stay task, about 1% of the data removed by user) and Lacuna-100 (ResNet-18, binary gender task, about 2% removed, from VGGFace2 with MAAD-Face labels) §4.1 app. B.1. Baselines are Constrained Newton Step (CNS) and Hessian-Free (HF) at σ = 0.01, and Finetune, Fisher Forgetting and SCRUB applied to the R2D original model §4.1 tab. 2.

Model AUC on the unlearned set after unlearning, σ = 0.01 tab. 2. The parenthesised value is the change from the same algorithm's original model, in AUC points:

Membership inference AUC, σ = 0.01 for the certified rows; lower means less leakage tab. 3. The paper calls MIA an imperfect measure in this setting app. B.1.1:

Training and unlearning wall-clock time on an RTX A6000 tab. 6 tab. 7 app. B.2:

The paper's "up to 88 times faster during training" matches 5.65 h against 230.76 s, about 88× §1 tab. 6. The paper says the Hessian-Free run uses fewer training iterations than R2D, so this is not an equal-compute comparison app. B.2.

Checkpoint reconstruction by Algorithm 3, Euclidean distance to the saved checkpoint tab. 8: eICU 0.065 at 10.3% rewind, 0.292 at 23.1%, 0.998 at 35.9%, 2.924 at 48.7%. Lacuna-100 1.174 at 7.4%, 4.464 at 44.4%. The paper finds reconstruction faithful for K < T/2 §2 app. B.3.

LLM preliminary experiment: Mistral-7B-Instruct-v0.2 with LoRA on Alpaca, math tasks as the unlearned set, BLEU on those tasks tab. 11 app. B.4. A lower BLEU is read as stronger unlearning, since the text says both R2D variants beat Finetuning:

Paper vs code

Official repository: github.com/siqiaomu/r2d, MIT licence, commit 8e2534c93dd3aa5bf9ccba5a1d2063832d65d4fa; the paper links it §4 app. B.1.

• Model selection: the paper selects the checkpoint with the lowest validation loss app. B.1. The code selects by validation error main.py.
• MLP depth: the paper says three hidden layers §4.1 app. B.1. The code's MLP has two hidden layers, of 256 and 64 units models.py. The paper's count of 136,386 parameters tab. 4 equals this two-layer MLP with 467 inputs (by my count: 467·256+256 + 256·64+64 + 64·2+2). The 467 is the code's default input size.
• Checkpoint interval: the paper says every 10 epochs §4.1. The code saves every 5 epochs main.py.
• Lipschitz estimate: the paper's formula puts the parameter difference over the gradient difference app. B.1. The code divides the gradient difference by the parameter difference main.py, which is the usual definition of the smoothness constant. The text's ratio is its inverse.
• Lipschitz and gradient bound for Lacuna-100: the released notebook reads G and L from the eICU training log. For Lacuna-100 both lines are commented out proxpoint.ipynb proxpoint.ipynb, and the Lacuna-100 run script does not request the Lipschitz estimate checkpointbash_lacuna.sh.
• Reconstruction (Algorithm 3): the paper solves each proximal subproblem by gradient descent or Newton's method §2. The code does one pass of mini-batch SGD per backward step, minimising the negated mini-batch loss plus the proximal term, with the centre fixed for that pass checkpoint_reconstruct.py. The inner learning rate defaults to 0.01 and the mini-batch size to 64 checkpoint_reconstruct.py checkpoint_reconstruct.py. The step size is passed as η = lr = 0.01 checkpoint_reconstruct.py.
• Noise on unlearned models: Algorithm 2 adds Gaussian noise after the K steps alg. 2. No released script adds this noise in an unlearning run. The only σ = 0.01 perturbation is in the Lipschitz estimate main.py. The noise helpers and the analytic calibration are in r2d.py r2d.py r2d.py, but h_function is never called in the released notebooks; it is only imported proxpoint.ipynb. The σ values in Tables 2 and 3 cannot be reproduced from the release.
• Forget-set choice: the paper says a random subset §4.1. For Lacuna-100 the code draws the same number of IDs from each class datasets.py. For eICU it draws IDs uniformly datasets.py.
• Unlearning runs: the eICU scripts resume from a checkpoint and train for 78 minus the checkpoint epoch, on the retained data, with the same mini-batch SGD and no added noise checkpointbash_eicu.sh.
• LLM rewinds: the text pairs the first checkpoint with two more epochs and the second with one app. B.4. The script trains both checkpoints for two epochs lora.py, so the 33% row is not produced by this script.
• Epoch counts: the paper uses 78 epochs for eICU tab. 4. The released training command in the README uses 80 with model selection README.md. The unlearning scripts use 78 checkpointbash_eicu.sh.

The learning rate (0.01), batch sizes (2048 and 512) and epoch counts (78 and 270) in the run scripts match Table 4 checkpointbash_eicu.sh checkpointbash_lacuna.sh tab. 4.

Pitfalls

• ε is never reported for the σ = 0.01 rows. The paper fixes σ for utility and then computes the privacy level from eq. (3) §3 app. B.1. To get ε you need L, G, n, m, K, T, η and δ. Lacuna-100 has no stated L or G (see Paper vs code).
• Theorem 3.1 covers 0 < ε ≤ 1 only §3. For ε > 1 the paper switches to the analytic Gaussian calibration of Balle and Wang app. B.1. That calibration is in r2d.py r2d.py, but nothing in the release calls it for R2D.
• G is an estimate, not a bound. The theorem assumes every per-example gradient is bounded by G everywhere §3. The paper estimates G as the largest mini-batch gradient norm seen in training app. B.1. Mini-batch norms are not per-example bounds, and the authors list this among the sources of imprecision app. C.
• L is an estimate too. It is the largest ratio over 100 noisy copies of the trained weights, with σ = 0.01 app. B.1. Its accuracy depends on the random samples.
• The step-size condition η ≤ min{1/L, n / (2(n−m)L)} §3 can be checked for eICU (η = 0.01 and L = 0.2065 satisfy η < 1/L, tab. 4) but not for Lacuna-100 from the paper.
• Noise shrinks slowly with rewinding. σ falls only slowly as K grows, so small σ needs K close to T, which means close to a full retrain app. C §3 fig. 1.
• Noise scales with n and m. σ shrinks as n grows and grows with m §3, eq. 3. For a fixed σ, more deletions weaken the guarantee. Sequential requests weaken it further, so the paper's advice is to re-run the algorithm on the new retained set §3.
• Dimension enters the population bound. Cor. 3.4 has a term L√d σ §3. For Lacuna-100, d = 11,160,258 tab. 4, so √d ≈ 3,341 (computed here). The bound is loose unless σ is small.
• Reconstruction degrades past half the trajectory. The paper says reconstruction is faithful only for K < T/2 §2. Table 8 shows the error rising from 0.065 to 2.924 on eICU tab. 8.
• Mini-batches are not full batch. The guarantees assume full-batch gradient descent §2 §3. The experiments use batch 2048 and 512 tab. 4, which the authors say slows convergence and makes the certificate approximate app. C.
• SmeLU β is not stated in the paper. The code's default is 0.5, while its docstring says 2 models.py models.py.
• Table captions overstate the metric. The caption of Table 2 says "relative decrease", but the numbers are absolute differences from each algorithm's own original model (for HF on eICU, 0.7614 to 0.7587 is −0.0027) tab. 2. The non-certified rows use the R2D original model, while the certified baselines have different originals §4.2.
• Small differences and one seed. The LLM result is one run, and the BLEU gap between R2D 33% and Finetuning is 0.0026 tab. 11. The paper gives ± values only for MIA, from cross-validation folds app. B.5. The released scripts use seed 1 checkpointbash_eicu.sh. The paper also reuses one Gaussian noise vector across σ values under the same seed, so table rows are correlated app. B.1.
• Small mismatch between tables. The main-text HF row for Lacuna-100 gives 0.4950 and 0.8460 tab. 3; the appendix gives 0.4949 and 0.8461 tab. 14.
• Default arguments are not the experiment settings. main.py defaults to lr 0.1, batch 256 and 100 epochs main.py. The scripts override these checkpointbash_eicu.sh.
• Storage. The extra storage is small: two model snapshots, about 8 MB each for eICU and about 45 MB each for Lacuna-100 app. B.2.
• Data access and licences. eICU needs credentialed access app. B.1. Lacuna-100 is built from VGGFace2 with MAAD-Face labels app. B.1. The licences of these datasets were not checked here; check them before commercial reuse.

Lineage

• Builds on: Descent-to-Delete (D2D), which the paper extends from strongly convex to nonconvex losses §1.1 §2. Also the certified removal line in Certified Data Removal (the Newton-step family the paper cites as Guo et al.).
• Compared against: Constrained Newton Step §4.1; Hessian-Free unlearning §4.1; Langevin unlearning, which the paper does not run §4.1; Finetune, Fisher Forgetting [golatkar1] and SCRUB [kurmanji2023towards] §4.1.
• Superseded by: not known from this wiki.