import torch import torch.nn as nn from torch.utils.data import DataLoader, RandomSampler from torch.autograd import grad from unlearning.Strategy import Strategy from sets.Data import * # Single-Batch Certified Unlearning for DNNs class CertifiedUnlearning(Strategy): """ Implements Certified Unlearning for non-convex DNNs (Zhang et al.). Uses a modified, stabilized stochastic Newton step using Taylor-expansion HVP estimation across the entire parameter space, capped with calibrated noise. """ def __init__( self, target_class_index: int, l2_reg: float = 0.0005, gamma: float = 0.01, scale: float = 50000.0, s1: int = 2, s2: int = 350, std: float = 0.001, unlearn_bs: int = 2 ): super().__init__(target_class_index) self.l2_reg = l2_reg self.gamma = gamma self.scale = scale self.s1 = s1 self.s2 = s2 self.std = std self.unlearn_bs = unlearn_bs def get_params(self, model: nn.Module, named): """ Safely collects named parameter tuples while skipping InceptionV3 auxiliary layers and tracking gradients. """ inner_model = getattr(model, "model", model) # Check if the current architecture is an Inception variant is_inception = inner_model.__class__.__name__.lower() == "inception3" params_list = [] for name, p in inner_model.named_parameters(): if p.requires_grad: # Discard the disconnected auxiliary training branch weights if is_inception and "AuxLogits" in name: continue # CRITICAL: Append as a tuple so it can be unpacked as (name, param) params_list.append((name, p)) return params_list if named else [e[1] for e in params_list] def _compute_loss_gradient(self, model, loader, device: torch.device): model.eval() criterion = nn.CrossEntropyLoss(reduction='sum') params = self.get_params(model, False) grad_accumulator = [torch.zeros_like(p, device=device) for p in params] total_samples = 0 with torch.set_grad_enabled(True): for data, targets in loader: total_samples += targets.shape[0] data, targets = data.to(device), targets.to(device) outputs = model(data) loss = criterion(outputs, targets) mini_grads = grad(loss, params, retain_graph=False) for i in range(len(grad_accumulator)): grad_accumulator[i] += mini_grads[i] # Data mean conversion for i in range(len(grad_accumulator)): grad_accumulator[i] /= total_samples # regularisation gradient for i, param in enumerate(params): grad_accumulator[i] += 2 * self.l2_reg * param.detach() return grad_accumulator def _hvp(self, loss, params, v): first_grads = grad(loss, params, retain_graph=True, create_graph=True) elemwise_products = sum(torch.sum(g_elem * v_elem) for g_elem, v_elem in zip(first_grads, v)) return grad(elemwise_products, params, create_graph=False) def _stochastic_newton_update(self, g, dataset, model, device): model.eval() criterion = nn.CrossEntropyLoss() params = self.get_params(model, False) h_res = [torch.zeros_like(p, device=device) for p in g] total_steps = self.s1 * self.s2 step_interval = max(1, total_steps // 100) global_step = 0 current_pct = 0 # Create DataLoader outside or use optimal sampling sampler = RandomSampler(dataset, replacement=True, num_samples=self.unlearn_bs * self.s2 * self.s1) res_loader = DataLoader(dataset, batch_size=self.unlearn_bs, sampler=sampler) res_iter = iter(res_loader) for _ in range(self.s1): h_estimate = [p.clone() for p in g] # hesian estimation for _ in range(self.s2): global_step += 1 try: data, target = next(res_iter) except StopIteration: res_iter = iter(res_loader) data, target = next(res_iter) data, target = data.to(device), target.to(device) # forward outputs = model(data) loss = criterion(outputs, target) l2_reg_term = sum(p.pow(2).sum() for p in params) loss += (self.l2_reg + self.gamma) * l2_reg_term h_s = self._hvp(loss, params, h_estimate) # OPTIMIZATION 4: Avoid deprecated .data, use detach() and in-place ops with torch.no_grad(): for k in range(len(params)): h_estimate[k].copy_(h_estimate[k] + g[k] - (h_s[k] / self.scale)) # feed back on status if global_step % step_interval == 0 and current_pct < 100: current_pct += 1 print(f"\rProgress: {current_pct}% done", end="", flush=True) with torch.no_grad(): for k in range(len(params)): h_res[k] += h_estimate[k] / self.scale return [p / self.s1 for p in h_res] def _certify(self, model, device, delta, full_certification): certification = "full " if full_certification else "partial" print(f"Performing {certification} certification") delta_idx = 0 # named_parameters to monitor layer positions for name, param in self.get_params(model, True): if param.requires_grad: noise = self.std * torch.randn(param.data.size(), device=device) if full_certification: param.data.add_(delta[delta_idx] + noise) else: # option for applying certification only to last layers # deprecated if "layer4" in name or "fc" in name: param.data.add_(delta[delta_idx] + noise) else: # Keep early low-level vision filters entirely pristine pass # Move to the next calculated Hessian vector block only after a valid update step delta_idx += 1 return model def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module: device = next(model.parameters()).device print(">> Calculating stable base gradients over the Forget set...") g = self._compute_loss_gradient(model, forget_loader, device) print(">> Estimating non-convex inverse Hessian trajectories via Taylor series...") dataset = retain_loader.dataset delta = self._stochastic_newton_update(g, dataset, model, device) print(">> Applying parameter removal adjustments (-delta)...") model = self._certify( model= model, device = device, delta = delta, full_certification = True ) print(">> Certified Unlearning process completed successfully.") return model # overriden function def _split_data(self, dataset): # Certified unlearning does require both forget and retain sets # split horizontaly. one class to forget and the rest to retain return get_unlearning_loaders( dataset=dataset, forget_class_idx=self.target_class_index, batch_size = 32 )