214 lines
9.7 KiB
Python
214 lines
9.7 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.data import DataLoader, RandomSampler
|
|
from torch.autograd import grad
|
|
from unlearning.Strategy import Strategy
|
|
|
|
class CertifiedRemoval(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 = 1000.0,
|
|
s1: int = 10, s2: int = 1000, 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 _compute_loss_gradient(self, model, loader, device: torch.device):
|
|
model.eval()
|
|
criterion = nn.CrossEntropyLoss(reduction='sum')
|
|
params = [p for p in model.parameters() if p.requires_grad]
|
|
grad_accumulator = [torch.zeros_like(p).cpu() for p in params]
|
|
total_samples = 0
|
|
|
|
for data, targets in loader:
|
|
total_samples += targets.shape[0]
|
|
data, targets = data.to(device), targets.to(device)
|
|
outputs = model(data)
|
|
|
|
mini_grads = list(grad(criterion(outputs, targets), params))
|
|
for i in range(len(grad_accumulator)):
|
|
grad_accumulator[i] += mini_grads[i].cpu().detach()
|
|
|
|
for i in range(len(grad_accumulator)):
|
|
grad_accumulator[i] /= total_samples
|
|
|
|
l2_reg_term = 0.0
|
|
for param in model.parameters():
|
|
l2_reg_term += torch.norm(param, p=2)
|
|
|
|
reg_grads = list(grad(self.l2_reg * l2_reg_term, params))
|
|
for i in range(len(grad_accumulator)):
|
|
grad_accumulator[i] += reg_grads[i].cpu().detach()
|
|
|
|
return [p.to(device) for p in grad_accumulator]'''
|
|
def _compute_loss_gradient(self, model, loader, device: torch.device):
|
|
model.eval()
|
|
# Use reduction='sum' matching the original framework
|
|
criterion = nn.CrossEntropyLoss(reduction='sum')
|
|
params = [p for p in model.parameters() if p.requires_grad]
|
|
grad_accumulator = [torch.zeros_like(p).cpu() for p in params]
|
|
total_samples = 0
|
|
|
|
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)
|
|
|
|
# Incorporate L2 weight regularization directly inside the backprop graph
|
|
# to keep scaling bounded and aligned with the data volume
|
|
l2_reg_term = 0.0
|
|
for param in model.parameters():
|
|
if param.requires_grad:
|
|
l2_reg_term += torch.norm(param, p=2)
|
|
|
|
total_loss = loss + (self.l2_reg * l2_reg_term)
|
|
|
|
mini_grads = list(grad(total_loss, params, retain_graph=False))
|
|
for i in range(len(grad_accumulator)):
|
|
grad_accumulator[i] += mini_grads[i].cpu().detach()
|
|
|
|
for i in range(len(grad_accumulator)):
|
|
grad_accumulator[i] /= total_samples
|
|
|
|
return [p.to(device) for p in grad_accumulator]
|
|
|
|
|
|
def grad_batch(batch_loader, lam, model, device):
|
|
model.eval()
|
|
criterion = nn.CrossEntropyLoss(reduction='sum')
|
|
params = [p for p in model.parameters() if p.requires_grad]
|
|
grad_batch = [torch.zeros_like(p).cpu() for p in params]
|
|
num = 0
|
|
for batch_idx, (data, targets) in enumerate(batch_loader):
|
|
num += targets.shape[0]
|
|
data, targets = data.to(device), targets.to(device)
|
|
outputs = model(data)
|
|
|
|
grad_mini = list(grad(criterion(outputs, targets), params))
|
|
for i in range(len(grad_batch)):
|
|
grad_batch[i] += grad_mini[i].cpu().detach()
|
|
|
|
for i in range(len(grad_batch)):
|
|
grad_batch[i] /= num
|
|
|
|
l2_reg = 0
|
|
for param in model.parameters():
|
|
l2_reg += torch.norm(param, p=2)
|
|
grad_reg = list(grad(lam * l2_reg, params))
|
|
for i in range(len(grad_batch)):
|
|
grad_batch[i] += grad_reg[i].cpu().detach()
|
|
return [p.to(device) for p in grad_batch]
|
|
|
|
def _hvp(self, loss, params, v):
|
|
first_grads = grad(loss, params, retain_graph=True, create_graph=True)
|
|
elemwise_products = 0
|
|
for grad_elem, v_elem in zip(first_grads, v):
|
|
elemwise_products += torch.sum(grad_elem * v_elem)
|
|
# FIX 1: Set create_graph to False to prevent massive nested graph accumulation
|
|
return grad(elemwise_products, params, create_graph=False)
|
|
|
|
def _stochastic_newton_update(self, g, retain_dataset, model, device):
|
|
model.eval()
|
|
criterion = nn.CrossEntropyLoss()
|
|
params = [p for p in model.parameters() if p.requires_grad]
|
|
h_res = [torch.zeros_like(p) for p in g]
|
|
|
|
for _ in range(self.s1):
|
|
h_estimate = [p.clone() for p in g]
|
|
sampler = RandomSampler(retain_dataset, replacement=True, num_samples=self.unlearn_bs * self.s2)
|
|
res_loader = DataLoader(retain_dataset, batch_size=self.unlearn_bs, sampler=sampler)
|
|
res_iter = iter(res_loader)
|
|
|
|
for j in range(self.s2):
|
|
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)
|
|
outputs = model(data)
|
|
|
|
loss = criterion(outputs, target)
|
|
l2_reg_term = 0.0
|
|
for param in model.parameters():
|
|
l2_reg_term += torch.norm(param, p=2)
|
|
loss += (self.l2_reg + self.gamma) * l2_reg_term
|
|
|
|
h_s = self._hvp(loss, params, h_estimate)
|
|
|
|
with torch.no_grad():
|
|
for k in range(len(params)):
|
|
# FIX 2: Added .detach() to decouple history strings across iterative update blocks
|
|
#h_estimate[k] = (h_estimate[k] + g[k] - h_s[k] / self.scale).detach()
|
|
next_estimate = h_estimate[k].data + g[k].data - (h_s[k].data / self.scale)
|
|
h_estimate[k] = next_estimate.clone()
|
|
del h_s, loss, outputs
|
|
|
|
for k in range(len(params)):
|
|
h_res[k] = h_res[k] + h_estimate[k] / self.scale
|
|
|
|
return [p / self.s1 for p in h_res]
|
|
|
|
'''def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
|
|
device = next(model.parameters()).device
|
|
|
|
num_forget = len(forget_loader.dataset)
|
|
num_retain = len(retain_loader.dataset)
|
|
scaling_ratio = num_forget / num_retain
|
|
|
|
print(">> Calculating base gradients over target FORGET set...")
|
|
# FIX 3: Base gradients MUST be evaluated from forget_loader to drop target class distributions
|
|
g = self._compute_loss_gradient(model, forget_loader, device)
|
|
|
|
print(">> Estimating non-convex inverse Hessian trajectories via Taylor series...")
|
|
retain_dataset = retain_loader.dataset
|
|
delta = self._stochastic_newton_update(g, retain_dataset, model, device)
|
|
|
|
print(">> Applying stabilized parameter adjustments and randomized certification noise...")
|
|
with torch.no_grad():
|
|
for i, param in enumerate(model.parameters()):
|
|
if param.requires_grad:
|
|
noise = self.std * torch.randn(param.data.size(), device=device)
|
|
#param.data.add_(-delta[i] + noise)
|
|
param.data.add_(scaling_ratio * delta[i] + noise)
|
|
|
|
print(">> Certified Unlearning process completed successfully across the complete landscape.")
|
|
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 RETAIN set...")
|
|
# To match the author's snippet perfectly, g MUST be computed on the retain data.
|
|
# If this loader is too large for your VRAM, use a smaller batch size (e.g. 16 or 32)
|
|
# in your main training script when creating retain_loader.
|
|
g = self._compute_loss_gradient(model, retain_loader, device)
|
|
|
|
print(">> Estimating non-convex inverse Hessian trajectories via Taylor series...")
|
|
retain_dataset = retain_loader.dataset
|
|
delta = self._stochastic_newton_update(g, retain_dataset, model, device)
|
|
|
|
print(">> Applying parameter removal adjustments (-delta)...")
|
|
with torch.no_grad():
|
|
for i, param in enumerate(model.parameters()):
|
|
if param.requires_grad:
|
|
noise = self.std * torch.randn(param.data.size(), device=device)
|
|
|
|
# MATCHING THE SNIPPET: Subtract delta exactly as the authors do
|
|
# This removes the influence trace of the omitted data.
|
|
param.data.add_(-delta[i] + noise)
|
|
|
|
print(">> Certified Unlearning process completed successfully.")
|
|
return model |