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
|
||||
@@ -1,47 +1,184 @@
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from .Strategy import Strategy
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
class LinearFiltration(Strategy):
|
||||
def __init__(self,target_class_index):
|
||||
super().__init__(target_class_index = target_class_index)
|
||||
def __init__(self, target_class_index):
|
||||
super().__init__(target_class_index=target_class_index)
|
||||
|
||||
def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
|
||||
model.eval()
|
||||
# Freeze internal params
|
||||
for param in model.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
device = next(model.parameters()).device
|
||||
|
||||
with torch.no_grad():
|
||||
W = model.fc.weight.data.clone()
|
||||
num_classes = W.shape[0]
|
||||
|
||||
A = self._calculate_filtration_matrix(num_classes, self.target_class_index, W.device)
|
||||
sanitized_W = torch.mm(A, W)
|
||||
model.fc.weight.copy_(sanitized_W)
|
||||
# Filter the bias (if the layer uses one)
|
||||
if model.fc.bias is not None:
|
||||
b = model.fc.bias.data.clone()
|
||||
# b is a 1D tensor of shape (num_classes),
|
||||
# so we use torch.mv (matrix-vector multiplication) or unsqueeze it
|
||||
sanitized_b = torch.mv(A, b)
|
||||
model.fc.bias.copy_(sanitized_b)
|
||||
|
||||
return model
|
||||
return self.normalise(
|
||||
model=model,
|
||||
retain_loader=retain_loader,
|
||||
forget_loader=forget_loader,
|
||||
device=device,
|
||||
forget_index=self.target_class_index
|
||||
)
|
||||
|
||||
# FIX: Added staticmethod decorator
|
||||
@staticmethod
|
||||
def get_features(model, inputs):
|
||||
# For ResNet, pass through everything up to the fc layer
|
||||
x = model.conv1(inputs)
|
||||
x = model.bn1(x)
|
||||
x = model.relu(x)
|
||||
x = model.maxpool(x)
|
||||
|
||||
x = model.layer1(x)
|
||||
x = model.layer2(x)
|
||||
x = model.layer3(x)
|
||||
x = model.layer4(x)
|
||||
|
||||
x = model.avgpool(x)
|
||||
x = torch.flatten(x, 1)
|
||||
return x
|
||||
|
||||
@staticmethod
|
||||
def _calculate_filtration_matrix(num_classes: int, forget_class: int, device: torch.device) -> torch.Tensor:
|
||||
A = torch.eye(num_classes, device=device)
|
||||
num_remaining = num_classes - 1
|
||||
|
||||
# The row of the forgotten class should average all other classes
|
||||
for j in range(num_classes):
|
||||
if j == forget_class:
|
||||
# we zero the forget class
|
||||
A[forget_class, j] = 0.0
|
||||
else:
|
||||
# and we distribute the output to the remaining
|
||||
A[forget_class, j] = 1.0 / num_remaining
|
||||
|
||||
return A
|
||||
return A
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _sums_and_counts(model, num_classes, retain_loader, forget_loader, device, forget_index, h_dim):
|
||||
model.eval()
|
||||
|
||||
sums = torch.zeros(num_classes, h_dim, device=device)
|
||||
counts = torch.zeros(num_classes, device=device)
|
||||
|
||||
# Generate values for retain
|
||||
with torch.no_grad():
|
||||
for inputs, targets in retain_loader:
|
||||
inputs = inputs.to(device)
|
||||
targets = targets.to(device)
|
||||
# FIX: Call get_features instead of model() directly
|
||||
outputs = LinearFiltration.get_features(model, inputs)
|
||||
|
||||
for j in range(num_classes):
|
||||
if j == forget_index:
|
||||
continue
|
||||
mask = (targets == j)
|
||||
|
||||
if mask.any():
|
||||
sums[j] += outputs[mask].sum(dim=0)
|
||||
counts[j] += mask.sum()
|
||||
|
||||
# Values for forget
|
||||
with torch.no_grad():
|
||||
for inputs, targets in forget_loader:
|
||||
inputs = inputs.to(device)
|
||||
targets = targets.to(device)
|
||||
# FIX: Call get_features instead of model() directly
|
||||
outputs = LinearFiltration.get_features(model, inputs)
|
||||
|
||||
mask = (targets == forget_index)
|
||||
|
||||
if mask.any():
|
||||
sums[forget_index] += outputs[mask].sum(dim=0)
|
||||
counts[forget_index] += mask.sum()
|
||||
|
||||
return sums, counts
|
||||
|
||||
@staticmethod
|
||||
def _get_means(model, num_classes, retain_loader, forget_loader, device, forget_index):
|
||||
h_dim = model.fc.in_features
|
||||
|
||||
sums, counts = LinearFiltration._sums_and_counts(
|
||||
model=model,
|
||||
num_classes=num_classes,
|
||||
retain_loader=retain_loader,
|
||||
forget_loader=forget_loader,
|
||||
device=device,
|
||||
forget_index=forget_index,
|
||||
h_dim=h_dim
|
||||
)
|
||||
A = []
|
||||
|
||||
for i in range(num_classes):
|
||||
if counts[i] > 0:
|
||||
A.append(sums[i] / counts[i])
|
||||
else:
|
||||
A.append(torch.zeros(h_dim, device=device))
|
||||
|
||||
# CORRECT: Stack along dim=0 to make it (num_classes, h_dim)
|
||||
return torch.stack(A, dim=0)
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _compute_z(tensor, forget_index):
|
||||
# Now tensor has shape (num_classes, h_dim) -> tensor.shape[0] is num_classes
|
||||
K = tensor.shape[0]
|
||||
|
||||
# pi_a0 should match the feature space dimensions (h_dim)
|
||||
pi_a0 = torch.zeros(tensor.shape[1], device=tensor.device)
|
||||
|
||||
t_1 = pi_a0
|
||||
a0 = tensor[forget_index, :] # Extracting the row vector for the forgotten class
|
||||
|
||||
mask_a0 = torch.ones(
|
||||
a0.shape[0],
|
||||
dtype=torch.bool,
|
||||
device=tensor.device
|
||||
)
|
||||
# We compute the target shift over features
|
||||
t_2 = -(1.0 / (K - 1)) * a0[mask_a0].sum()
|
||||
|
||||
mask_rows = torch.ones(K, dtype=torch.bool, device=tensor.device)
|
||||
mask_rows[forget_index] = False
|
||||
|
||||
r_A = tensor[mask_rows, :]
|
||||
t_3 = (1.0 / ((K - 1)) ** 2) * r_A.sum()
|
||||
|
||||
return t_1 + t_2 + t_3
|
||||
|
||||
|
||||
@staticmethod
|
||||
def normalise(model, retain_loader, forget_loader, device, forget_index):
|
||||
W = model.fc.weight.data.clone()
|
||||
num_classes = W.shape[0]
|
||||
|
||||
A = LinearFiltration._get_means(
|
||||
model=model,
|
||||
num_classes=num_classes,
|
||||
retain_loader=retain_loader,
|
||||
forget_loader=forget_loader,
|
||||
device=device,
|
||||
forget_index=forget_index
|
||||
)
|
||||
|
||||
Z = LinearFiltration._compute_z(tensor=A, forget_index=forget_index)
|
||||
B_Z_rows = []
|
||||
|
||||
for i in range(num_classes):
|
||||
if i == forget_index:
|
||||
B_Z_rows.append(Z)
|
||||
else:
|
||||
# Retained classes maintain their original ideal feature directions
|
||||
B_Z_rows.append(A[i])
|
||||
|
||||
# Stack back along dim=0 to match (num_classes, h_dim)
|
||||
B_Z = torch.stack(B_Z_rows, dim=0)
|
||||
|
||||
A_inv = torch.linalg.pinv(A)
|
||||
|
||||
W_Z = B_Z @ A_inv @ W
|
||||
|
||||
model.fc.weight.copy_(W_Z)
|
||||
|
||||
return model
|
||||
@@ -3,97 +3,34 @@ import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
from torch.utils.data import DataLoader
|
||||
from unlearning.Strategy import Strategy
|
||||
from .wf.WF_Net import WF_Net
|
||||
|
||||
class WeightFiltration(Strategy):
|
||||
"""
|
||||
Implements Poppi et al.'s Weight Filtering framework for linear layers.
|
||||
Uses a standard functional hook to guarantee native PyTorch autograd tracking.
|
||||
Verbatim implementation of Poppi et al.'s WF-Net framework.
|
||||
Directly filters the convolutional weights of a target layer using a learnable
|
||||
channel mask, optimizing it via weight-space regularization.
|
||||
"""
|
||||
def __init__(self, target_class_index,num_classes: int, epochs: int = 10, lr: float = 0.2, gamma: float = 10.0):
|
||||
super().__init__(target_class_index = target_class_index)
|
||||
self.num_classes = num_classes
|
||||
def __init__(self, target_class_index: int, epochs: int = 10, lr: float = 0.2, gamma: float = 10.0):
|
||||
super().__init__(target_class_index=target_class_index)
|
||||
self.epochs = epochs
|
||||
self.lr = lr
|
||||
self.gamma = gamma
|
||||
self.alpha = None
|
||||
self.hook_handle = None
|
||||
#self.alpha = None
|
||||
|
||||
|
||||
def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
|
||||
device = next(model.parameters()).device
|
||||
model.eval()
|
||||
|
||||
|
||||
|
||||
|
||||
def _optimise_filter(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader, device):
|
||||
# 1. Initialize the wrapper with your pre-trained model
|
||||
num_classes = model.fc.out_features
|
||||
wf_model = WF_Net(original_model=model, num_classes=num_classes).to(device)
|
||||
|
||||
# Locate layer4 for dynamic optimization
|
||||
target_layer = model.layer4 if hasattr(model, 'layer4') else None
|
||||
fc_layer = model.fc if hasattr(model, 'fc') and isinstance(model.fc, nn.Linear) else None
|
||||
|
||||
if target_layer is None or fc_layer is None:
|
||||
raise AttributeError("Model does not have the required layers.")
|
||||
|
||||
# Match alpha dimensions to the channels outputted by layer4
|
||||
num_features = fc_layer.weight.shape[1]
|
||||
self.alpha = nn.Parameter(torch.ones(self.num_classes, num_features, device=device) * 1.5)
|
||||
|
||||
# Freeze everything except our channel mask
|
||||
for p in model.parameters():
|
||||
p.requires_grad = False
|
||||
self.alpha.requires_grad = True
|
||||
|
||||
# Hook into layer4 dynamically to run the untraining optimization
|
||||
self.hook_handle = target_layer.register_forward_hook(self._get_hook())
|
||||
# optimise the filter to maintain accuracy on retain set
|
||||
# and decrease accuracy on forget set
|
||||
self._optimise_filter(model, forget_loader, retain_loader, device)
|
||||
|
||||
# Remove the runtime hook
|
||||
self.hook_handle.remove()
|
||||
|
||||
# Transfer the channel suppression permanently into model.fc
|
||||
with torch.no_grad():
|
||||
mask = torch.sigmoid(self.alpha[self.target_class_index]) # Shape: (num_features,)
|
||||
|
||||
# Suppress the channels ONLY for the target class row in fc
|
||||
fc_layer.weight[self.target_class_index].copy_(
|
||||
fc_layer.weight[self.target_class_index] * mask
|
||||
)
|
||||
print(f">> Baked deep channel filter into Class {self.target_class_index} weights.")
|
||||
|
||||
return model
|
||||
|
||||
def _get_hook(self):
|
||||
"""
|
||||
Filters the internal feature map channels of layer4.
|
||||
The mask scales the channels across the batch.
|
||||
"""
|
||||
def functional_hook(module, layer_input, layer_output):
|
||||
# layer_output shape: (batch, channels, height, width) -> e.g., (16, 2048, 7, 7)
|
||||
# self.alpha shape: (num_classes, channels) -> e.g., (20, 2048)
|
||||
|
||||
# Extract 1D mask for the target class: (channels,)
|
||||
mask = torch.sigmoid(self.alpha[self.target_class_index])
|
||||
|
||||
# Reshape mask to (1, channels, 1, 1) so it broadcasts over batch, height, and width
|
||||
mask = mask.view(1, -1, 1, 1)
|
||||
|
||||
# Scale the internal feature maps before they move to the next layer
|
||||
return layer_output * mask
|
||||
|
||||
return functional_hook
|
||||
|
||||
|
||||
def _optimise_filter(self, model, forget_loader, retain_loader, device):
|
||||
optimizer = optim.Adam([self.alpha], lr=self.lr)
|
||||
# 2. ONLY optimize alpha (everything else is frozen inside the wrapper)
|
||||
optimizer = optim.Adam([wf_model.alpha], lr=self.lr)
|
||||
criterion = nn.CrossEntropyLoss()
|
||||
|
||||
print(f"[{self.__class__.__name__}] Unlearning Class {self.target_class_index} with gamma={self.gamma}...")
|
||||
|
||||
# To optimise this loop we will watch improvements after each optimisation
|
||||
temp_forget_loss = None
|
||||
# this can be adjusted to optimise the best escape point
|
||||
# it is the value we set to evaluate performance improvement after each itteration.
|
||||
# if improvement is less than this, then we break itteration.
|
||||
threshold = 0.05
|
||||
|
||||
for epoch in range(self.epochs):
|
||||
forget_iter = iter(forget_loader)
|
||||
t_loss_r, t_loss_f = 0.0, 0.0
|
||||
@@ -102,6 +39,7 @@ class WeightFiltration(Strategy):
|
||||
for r_inputs, r_labels in retain_loader:
|
||||
r_inputs, r_labels = r_inputs.to(device), r_labels.to(device)
|
||||
|
||||
# Pull the matching forget batch input
|
||||
try:
|
||||
f_inputs, _ = next(forget_iter)
|
||||
except StopIteration:
|
||||
@@ -111,10 +49,19 @@ class WeightFiltration(Strategy):
|
||||
|
||||
optimizer.zero_grad()
|
||||
|
||||
# Compute Losses
|
||||
# The hook handles the weight filtering smoothly behind the scenes
|
||||
loss_r = criterion(model(r_inputs), r_labels)
|
||||
loss_f = -torch.sum((torch.ones_like(model(f_inputs)) / self.num_classes) * torch.log_softmax(model(f_inputs), dim=-1))
|
||||
# --- APPLY ALGORITHM 1 FORWARD PASS TO BOTH INPUTS ---
|
||||
# Pass the input batch AND the target unlearn class index
|
||||
outputs_r = wf_model(r_inputs, target_unlearn_class=self.target_class_index)
|
||||
outputs_f = wf_model(f_inputs, target_unlearn_class=self.target_class_index)
|
||||
|
||||
# Compute Losses using Poppi et al.'s temperature scaled entropy
|
||||
loss_r = criterion(outputs_r, r_labels)
|
||||
|
||||
temperature = 3.0
|
||||
logits_f_scaled = outputs_f / temperature
|
||||
loss_f = -torch.sum(
|
||||
(torch.ones_like(logits_f_scaled) / num_classes) * torch.log_softmax(logits_f_scaled, dim=-1)
|
||||
)
|
||||
|
||||
total_loss = loss_r + (self.gamma * loss_f)
|
||||
total_loss.backward()
|
||||
@@ -122,17 +69,56 @@ class WeightFiltration(Strategy):
|
||||
|
||||
t_loss_r += loss_r.item()
|
||||
t_loss_f += loss_f.item()
|
||||
|
||||
steps += 1
|
||||
forget_loss = t_loss_f / steps
|
||||
print(f" Epoch {epoch+1}/{self.epochs} | Retain Loss: {t_loss_r/steps:.4f} | Forget Loss: {forget_loss:.4f}")
|
||||
|
||||
print(f" Epoch {epoch+1}/{self.epochs} | Retain Loss: {t_loss_r/steps:.4f} | Forget Loss: {t_loss_f/steps:.4f}")
|
||||
|
||||
if temp_forget_loss is not None:
|
||||
return wf_model
|
||||
|
||||
|
||||
improvement = temp_forget_loss - forget_loss
|
||||
# if optimisation reaches a point of diminishing returns (improvements is less than threshold)
|
||||
# we break the loop
|
||||
if improvement < threshold:
|
||||
break
|
||||
# else we update the lasst recorded loss.
|
||||
temp_forget_loss = forget_loss
|
||||
|
||||
def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
|
||||
device = next(model.parameters()).device
|
||||
model.eval()
|
||||
|
||||
# In WF-Net, the mask targets the last major convolutional block
|
||||
# For ResNet-18, that is the final conv layer in layer4 block 1
|
||||
if hasattr(model, 'layer4') and len(model.layer4) > 1:
|
||||
target_conv = model.layer4[1].conv2
|
||||
else:
|
||||
raise AttributeError("Model architecture does not match expected ResNet-18 structure.")
|
||||
|
||||
# Store a pristine, non-grad copy of the original trained weights
|
||||
# Shape of conv2.weight: (out_channels, in_channels, kernel_size, kernel_size) -> e.g., (512, 512, 3, 3)
|
||||
original_weights = target_conv.weight.data.clone().detach()
|
||||
out_channels = original_weights.shape[0]
|
||||
|
||||
# Initialize alpha gate vector matching Poppi et al.'s initialization range
|
||||
# Shape: (out_channels,) -> acting directly as a filter-level gate
|
||||
#self.alpha = nn.Parameter(torch.ones(out_channels, device=device) * 1.5)
|
||||
|
||||
# Freeze the global model graph; only optimize our filter parameter mask
|
||||
for p in model.parameters():
|
||||
p.requires_grad = False
|
||||
#self.alpha.requires_grad = True
|
||||
|
||||
wf_model = self._optimise_filter(
|
||||
model,
|
||||
forget_loader=forget_loader,
|
||||
retain_loader=retain_loader,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# --- PERMANENT BAKING STEP ---
|
||||
# Disconnect the dynamic parameter and freeze the optimal gated state permanently into the architecture
|
||||
|
||||
with torch.no_grad():
|
||||
final_mask = torch.sigmoid(wf_model.alpha[self.target_class_index]).view(-1, 1, 1, 1)
|
||||
target_conv.weight.copy_(original_weights * final_mask)
|
||||
|
||||
# Re-enable model parameters for downstream evaluation processing
|
||||
for p in model.parameters():
|
||||
p.requires_grad = True
|
||||
|
||||
print(f">> Permanently altered {out_channels} convolutional filters in layer4 via WF-Net.")
|
||||
return model
|
||||
Reference in New Issue
Block a user