127 lines
5.6 KiB
Python
127 lines
5.6 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.data import DataLoader
|
|
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.
|
|
"""
|
|
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 _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 = []
|
|
|
|
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())
|
|
|
|
return torch.cat(all_features, dim=0), torch.cat(all_labels, dim=0)
|
|
|
|
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).
|
|
"""
|
|
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)
|
|
|
|
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)
|
|
|
|
forget_labels_one_hot = torch.nn.functional.one_hot(forget_labels, num_classes=num_classes).float()
|
|
|
|
|
|
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 |