strategies tested
This commit is contained in:
@@ -1,77 +1,153 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
from scipy.optimize import minimize
|
||||
from .Strategy import Strategy
|
||||
import torch.nn as nn
|
||||
from torch.utils.data import DataLoader
|
||||
from unlearning.Strategy import Strategy
|
||||
|
||||
class CertifiedRemoval(Strategy):
|
||||
"""Implements Certified Removal for machine unlearning."""
|
||||
|
||||
def __init__(self, model, data, labels, removal_bound, epsilon):
|
||||
"""
|
||||
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.model = model
|
||||
self.data = data
|
||||
self.labels = labels
|
||||
self.removal_bound = removal_bound
|
||||
self.epsilon = epsilon
|
||||
self.removal_bound = removal_bound # gamma in the paper
|
||||
self.epsilon = epsilon # Privacy budget
|
||||
self.l2_reg = l2_reg # Lambda regularization term
|
||||
|
||||
def _run(self, model: nn.Module) -> nn.Module:
|
||||
"""Runs the certified removal algorithm."""
|
||||
# 1. Linear Model Creation
|
||||
# This is a simplification for demonstration purposes. In a real implementation,
|
||||
# you'd use more sophisticated methods to learn the parameters of the
|
||||
# 'removal' model based on the example being removed.
|
||||
|
||||
def linear_model(x):
|
||||
return torch.dot(x, torch.tensor([1, 1])) # Simplified Linear Model
|
||||
|
||||
# 2. Optimization for Parameter Adjustment
|
||||
# Optimize the parameter values to minimize the loss while staying within bounds.
|
||||
original_params = torch.tensor([0.0, 0.0]) # Initial parameters for linear model
|
||||
|
||||
def objective_function(params):
|
||||
new_model = linear_model #use same function as defined above
|
||||
return torch.sum(((new_model(self.data[0]) - self.labels)**2))
|
||||
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())
|
||||
|
||||
result = minimize(objective_function, original_params, method='L-BFGS-B', bounds=[(-self.removal_bound, self.removal_bound)], options={'maxiter': 100})
|
||||
return torch.cat(all_features, dim=0), torch.cat(all_labels, dim=0)
|
||||
|
||||
if not result.success:
|
||||
print("Warning: Optimization failed!")
|
||||
print(result.message)
|
||||
return model #Return original if optimization fails
|
||||
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
|
||||
# this will be done on CPU. requires more ram so we cant afford to do it on VRAM
|
||||
# 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))
|
||||
hessian = self._compute_hessian(retain_features=retain_features, retain_features_size = N_retain)
|
||||
|
||||
# Compute the gradient of the loss with respect to the forgotten data
|
||||
# print(">> Calculating forget set gradients...")
|
||||
# num_classes = w.size(0)
|
||||
# Pass features through linear layer weights to get logits
|
||||
# logits_forget = torch.matmul(forget_features, w.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()
|
||||
|
||||
#preds_forget = torch.matmul(forget_features, w.t())
|
||||
#error = preds_forget - forget_labels_one_hot
|
||||
# error = preds_softmax - forget_labels_one_hot
|
||||
# grad_forget shape: [num_classes, 2048]
|
||||
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)
|
||||
|
||||
new_params = result.x
|
||||
# 3. New Model Creation
|
||||
# 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
|
||||
)
|
||||
# print(">> Solving Newton step via system optimization...")
|
||||
# try:
|
||||
# delta_w_t = torch.linalg.solve(Hessian, grad_forget.t())
|
||||
# delta_w = delta_w_t.t()
|
||||
# except RuntimeError:
|
||||
# print(">> Warning: Hessian matrix is singular. Falling back to pseudo-inverse.")
|
||||
# delta_w = torch.matmul(grad_forget, torch.linalg.pinv(Hessian).t())
|
||||
|
||||
new_model = lambda x: torch.dot(x, new_params)
|
||||
return new_model
|
||||
# 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}")
|
||||
|
||||
if __name__ == '__main__':
|
||||
# Example Usage - Synthetic Data for Demonstration
|
||||
np.random.seed(42) # For reproducibility
|
||||
n_samples = 100
|
||||
X = np.random.randn(n_samples, 2)
|
||||
y = (X[:, 0] + X[:, 1] > 0).astype(int)
|
||||
# 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
|
||||
|
||||
|
||||
# Create a simple linear model for demonstration
|
||||
model = nn.Linear(2, 1) # Simple linear classifier - PyTorch Version
|
||||
optimizer = torch.optim.SGD(model.parameters(), lr=0.01) # Optimizer for training the linear 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
|
||||
|
||||
|
||||
# Train a Linear Model
|
||||
for _ in range(100): #training loop
|
||||
optimizer.zero_grad()
|
||||
predictions = model(X)
|
||||
loss = torch.sum((predictions - y)**2)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
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()
|
||||
|
||||
|
||||
# Define parameters for Certified Removal
|
||||
removal_bound = 1.0
|
||||
epsilon = 0.1
|
||||
error = preds_softmax - forget_labels_one_hot
|
||||
# grad_forget shape: [num_classes, 2048]
|
||||
return torch.matmul(error.t(), forget_features) / forget_features.size(0)
|
||||
|
||||
|
||||
# Create the CertifiedRemoval object with the trained model, data and labels
|
||||
certified_removal_obj = CertifiedRemoval(model, X, y, removal_bound, epsilon)
|
||||
|
||||
# Run Certified Removal
|
||||
new_model = certified_removal_obj.apply(model)
|
||||
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
|
||||
Reference in New Issue
Block a user