certified
This commit is contained in:
@@ -1,127 +1,214 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.data import DataLoader
|
||||
from torch.utils.data import DataLoader, RandomSampler
|
||||
from torch.autograd import grad
|
||||
from unlearning.Strategy import Strategy
|
||||
|
||||
class CertifiedRemoval(Strategy):
|
||||
"""
|
||||
Implements Certified Removal (Guo et al.) adapted for deep architectures
|
||||
like ResNet50 by isolating and updating the final classification layer.
|
||||
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, removal_bound: float, epsilon: float, l2_reg: float = 0.1):
|
||||
super().__init__()
|
||||
self.removal_bound = removal_bound # gamma in the paper
|
||||
self.epsilon = epsilon # Privacy budget
|
||||
self.l2_reg = l2_reg # Lambda regularization term
|
||||
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 _get_features(self, backbone: nn.Module, loader: DataLoader, device: torch.device):
|
||||
"""Passes data through the frozen ResNet backbone to extract embedding features."""
|
||||
backbone.eval()
|
||||
all_features = []
|
||||
all_labels = []
|
||||
'''
|
||||
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
|
||||
|
||||
with torch.no_grad():
|
||||
for inputs, labels in loader:
|
||||
inputs = inputs.to(device)
|
||||
# Pass through backbone to get the 2048-dimensional feature vector
|
||||
features = backbone(inputs)
|
||||
all_features.append(features.cpu())
|
||||
all_labels.append(labels.cpu())
|
||||
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)
|
||||
|
||||
return torch.cat(all_features, dim=0), torch.cat(all_labels, dim=0)
|
||||
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:
|
||||
"""
|
||||
Entry point expected by your Model.unlearn() architecture interface.
|
||||
Applies Certified Removal strictly to the final linear layer (model.fc).
|
||||
"""
|
||||
'''def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
|
||||
device = next(model.parameters()).device
|
||||
|
||||
# Isolate the final NN (Fully connected) layer from the model
|
||||
linear_head = model.fc
|
||||
# Temporarily turn the fc layer into a identity pass-through
|
||||
model.fc = nn.Identity()
|
||||
|
||||
print(">> Extracting deep features from model backbone...")
|
||||
retain_features, retain_labels = self._get_features(model, retain_loader, device)
|
||||
forget_features, forget_labels = self._get_features(model, forget_loader, device)
|
||||
|
||||
# Restore the linear head back
|
||||
model.fc = linear_head
|
||||
|
||||
# Extract weights from the classification layer
|
||||
# w shape: [num_classes, 2048]
|
||||
w = model.fc.weight.data.clone().cpu()
|
||||
|
||||
# Compute the Exact Hessian Matrix over the remaining (retained) features
|
||||
# Formula: H = (X^T * X) / N + lambda * I
|
||||
N_retain = retain_features.size(0)
|
||||
hessian = self._compute_hessian(retain_features=retain_features, retain_features_size = N_retain)
|
||||
|
||||
grad_forget = self._compute_loss_gradient(
|
||||
forget_labels=forget_labels,
|
||||
forget_features=forget_features,
|
||||
model_weights=w)
|
||||
#torch.matmul(error.t(), forget_features) / forget_features.size(0)
|
||||
|
||||
# Compute the Newton step update via solving: H * Delta_W^T = Grad_forget^T
|
||||
delta_w = self._compute_newton_step(
|
||||
tensor = hessian,
|
||||
gradient= grad_forget
|
||||
)
|
||||
# Apply the Certified Removal update rule: W_new = W + Delta_W
|
||||
new_w = w + delta_w
|
||||
# Calibrate noise based on your epsilon budget
|
||||
# (Guo et al. use a perturbation based on the regularization lambda and epsilon)
|
||||
sigma = 2.0 / (self.l2_reg * self.epsilon)
|
||||
noise = torch.randn_like(new_w) * (sigma / N_retain)
|
||||
new_w = new_w + noise
|
||||
|
||||
# Theoretical Guarantee verification
|
||||
norm_delta = torch.norm(delta_w).item()
|
||||
if norm_delta > self.removal_bound:
|
||||
print(f"!! Warning: Removal budget exceeded! Norm: {norm_delta:.4f} > Bound: {self.removal_bound}")
|
||||
else:
|
||||
print(f">> Certificate valid. Norm: {norm_delta:.4f} <= Bound: {self.removal_bound}")
|
||||
|
||||
# Push updated parameters back into the model instance in-place
|
||||
model.fc.weight.data = new_w.to(device)
|
||||
num_forget = len(forget_loader.dataset)
|
||||
num_retain = len(retain_loader.dataset)
|
||||
scaling_ratio = num_forget / num_retain
|
||||
|
||||
print(">> Certified Removal process completed successfully.")
|
||||
return model
|
||||
|
||||
|
||||
# computing the hessian matrix
|
||||
def _compute_hessian(self, retain_features, retain_features_size):
|
||||
print(">> Computing exact Hessian matrix...")
|
||||
# N_retain = retain_features.size(0)
|
||||
X_T_X = torch.matmul(retain_features.t(), retain_features)
|
||||
reg_matrix = self.l2_reg * torch.eye(retain_features.size(1))
|
||||
return (X_T_X / retain_features_size) + reg_matrix
|
||||
|
||||
|
||||
def _compute_loss_gradient(self, forget_features, forget_labels, model_weights):
|
||||
print(">> Calculating forget set gradients...")
|
||||
num_classes = model_weights.size(0)
|
||||
# Pass features through linear layer weights to get logits
|
||||
logits_forget = torch.matmul(forget_features, model_weights.t())
|
||||
# Apply softmax to get true class probabilities
|
||||
preds_softmax = torch.softmax(logits_forget, dim=1)
|
||||
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)
|
||||
|
||||
forget_labels_one_hot = torch.nn.functional.one_hot(forget_labels, num_classes=num_classes).float()
|
||||
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
|
||||
|
||||
error = preds_softmax - forget_labels_one_hot
|
||||
# grad_forget shape: [num_classes, 2048]
|
||||
return torch.matmul(error.t(), forget_features) / forget_features.size(0)
|
||||
|
||||
|
||||
def _compute_newton_step(self,tensor, gradient):
|
||||
print(">> Solving Newton step via system optimization...")
|
||||
try:
|
||||
delta_w_t = torch.linalg.solve(tensor, gradient.t())
|
||||
delta_w = delta_w_t.t()
|
||||
except RuntimeError:
|
||||
print(">> Warning: Hessian matrix is singular. Falling back to pseudo-inverse.")
|
||||
delta_w = torch.matmul(gradient, torch.linalg.pinv(tensor).t())
|
||||
return delta_w
|
||||
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
|
||||
Reference in New Issue
Block a user