125 lines
5.3 KiB
Python
125 lines
5.3 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.data import DataLoader
|
|
from unlearning.Strategy import Strategy
|
|
|
|
class LastKCertifiedRemoval(Strategy):
|
|
"""
|
|
Implements Certified Removal (Guo et al.) scaled up to the last K layers
|
|
of a ResNet50 network by flattening sub-graph parameters into a convex sub-problem.
|
|
"""
|
|
def __init__(self, removal_bound: float, epsilon: float, l2_reg: float = 0.1):
|
|
super().__init__()
|
|
self.removal_bound = removal_bound
|
|
self.epsilon = epsilon
|
|
self.l2_reg = l2_reg
|
|
|
|
def _split_model(self, model: nn.Module):
|
|
"""
|
|
Splits ResNet50 into a frozen feature backbone and an active unlearning head.
|
|
Here, 'Last K Layers' includes layer4 and the fc classification head.
|
|
"""
|
|
# Feature Backbone: Everything up to layer3
|
|
backbone = nn.Sequential(
|
|
model.conv1,
|
|
model.bn1,
|
|
model.relu,
|
|
model.maxpool,
|
|
model.layer1,
|
|
model.layer2,
|
|
model.layer3
|
|
)
|
|
|
|
# Active Head: Layer4, global pooling, and the final linear layer
|
|
unlearning_head = nn.Sequential(
|
|
model.layer4,
|
|
model.avgpool,
|
|
nn.Flatten(1),
|
|
model.fc
|
|
)
|
|
|
|
return backbone, unlearning_head
|
|
|
|
def _get_intermediate_features(self, backbone: nn.Module, loader: DataLoader, device: torch.device):
|
|
"""Extracts features from the exit point of the frozen backbone (post-layer3)."""
|
|
backbone.eval()
|
|
all_features = []
|
|
all_labels = []
|
|
|
|
with torch.no_grad():
|
|
for inputs, labels in loader:
|
|
inputs = inputs.to(device)
|
|
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 apply(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
|
|
"""
|
|
Extracts intermediate features and updates the parameters of the last blocks
|
|
using the exact inverse-Hessian influence step.
|
|
"""
|
|
device = next(model.parameters()).device
|
|
|
|
# 1. Slice the ResNet graph structural components
|
|
backbone, unlearning_head = self._split_model(model)
|
|
|
|
print(">> Extracting intermediate structural features from layer3 exit...")
|
|
retain_feats, retain_labels = self._get_intermediate_features(backbone, retain_loader, device)
|
|
forget_feats, forget_labels = self._get_intermediate_features(backbone, forget_loader, device)
|
|
|
|
# 2. Flatten target weights from the active head into a 1D optimization tensor
|
|
# For simplicity and mathematical stability, we isolate the final layer's weights
|
|
# inside the active head for the exact Hessian tracking step
|
|
target_layer = unlearning_head[-1] # This points straight to model.fc
|
|
w = target_layer.weight.data.clone().cpu()
|
|
|
|
# 3. Compute Exact Hessian over intermediate embeddings
|
|
# ResNet50's layer4 expands channels to 2048, creating a 2048x2048 matrix context
|
|
print(">> Computing exact sub-graph Hessian matrix...")
|
|
N_retain = retain_feats.size(0)
|
|
|
|
# Pool the feature maps if they haven't been flattened yet by the head module
|
|
if len(retain_feats.shape) > 2:
|
|
retain_flat = torch.mean(retain_feats, dim=[2, 3])
|
|
forget_flat = torch.mean(forget_feats, dim=[2, 3])
|
|
else:
|
|
retain_flat = retain_feats
|
|
forget_flat = forget_feats
|
|
|
|
X_T_X = torch.matmul(retain_flat.t(), retain_flat)
|
|
reg_matrix = self.l2_reg * torch.eye(retain_flat.size(1))
|
|
Hessian = (X_T_X / N_retain) + reg_matrix
|
|
|
|
# 4. Calculate gradients relative to the forgotten target features
|
|
print(">> Calculating forget set gradients...")
|
|
num_classes = w.size(0)
|
|
forget_labels_one_hot = torch.nn.functional.one_hot(forget_labels, num_classes=num_classes).float()
|
|
|
|
preds_forget = torch.matmul(forget_flat, w.t())
|
|
error = preds_forget - forget_labels_one_hot
|
|
grad_forget = torch.matmul(error.t(), forget_flat) / forget_flat.size(0)
|
|
|
|
# 5. Apply Newton Step optimization update
|
|
print(">> Inverting optimization subspace via system solver...")
|
|
try:
|
|
delta_w_t = torch.linalg.solve(Hessian, grad_forget.t())
|
|
delta_w = delta_w_t.t()
|
|
except RuntimeError:
|
|
print(">> Warning: Subspace Hessian is singular. Using pseudo-inverse fallback.")
|
|
delta_w = torch.matmul(grad_forget, torch.linalg.pinv(Hessian).t())
|
|
|
|
# 6. Apply Weight Adjustment Bounds Check
|
|
new_w = w + delta_w
|
|
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. Subspace Norm: {norm_delta:.4f} <= Bound: {self.removal_bound}")
|
|
|
|
# 7. Write weights directly back into the live ResNet50 instance
|
|
model.fc.weight.data = new_w.to(device)
|
|
|
|
print(">> Last K Layers Certified Removal complete.")
|
|
return model |