123 lines
5.6 KiB
Python
123 lines
5.6 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import math
|
|
from torch.utils.data import DataLoader
|
|
from unlearning.Strategy import Strategy
|
|
|
|
class CertifiedRemovalFacebook(Strategy):
|
|
"""
|
|
Implements Certified Removal (Guo et al.) mapped for Multi-Class models
|
|
by executing a single-class One-vs-Rest (OvR) block-removal update step.
|
|
Math matches the facebookresearch/certified-removal reference repository.
|
|
"""
|
|
def __init__(self, target_class_index: int, removal_bound: float, epsilon: float, l2_reg: float = 0.1):
|
|
super().__init__(target_class_index=target_class_index)
|
|
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 = []
|
|
|
|
with torch.no_grad():
|
|
for inputs, _ in loader:
|
|
inputs = inputs.to(device)
|
|
# Pass through frozen backbone to get the 2048-dimensional embedding
|
|
features = backbone(inputs)
|
|
all_features.append(features.cpu())
|
|
|
|
return torch.cat(all_features, dim=0)
|
|
|
|
def _fb_lr_grad(self, w, X, y, lam):
|
|
"""
|
|
Replicates exact lr_grad calculation from Facebook's codebase.
|
|
Note: The resulting gradient has a flipped sign due to the structure of (z - 1).
|
|
"""
|
|
# X.mv(w) computes raw linear margins
|
|
z = torch.sigmoid(y * X.mv(w))
|
|
# Gradient formula: X^T * ((z - 1) * y) + lambda * N * w
|
|
return X.t().mv((z - 1) * y) + lam * X.size(0) * w
|
|
|
|
def _fb_lr_hessian_inv(self, w, X, y, lam, device, batch_size=50000):
|
|
"""
|
|
Replicates exact lr_hessian_inv calculation from Facebook's codebase.
|
|
Scales the L2 regularization matrix explicitly by dataset row count (N * lambda * I).
|
|
"""
|
|
z = torch.sigmoid(X.mv(w).mul_(y))
|
|
D = z * (1 - z) # Element-wise variance vector
|
|
|
|
H = None
|
|
num_batch = int(math.ceil(X.size(0) / batch_size))
|
|
for i in range(num_batch):
|
|
lower = i * batch_size
|
|
upper = min((i + 1) * batch_size, X.size(0))
|
|
X_i = X[lower:upper]
|
|
|
|
# Stepwise feature weighting via element-wise variance columns
|
|
if H is None:
|
|
H = X_i.t().mm(D[lower:upper].unsqueeze(1) * X_i)
|
|
else:
|
|
H += X_i.t().mm(D[lower:upper].unsqueeze(1) * X_i)
|
|
|
|
# Scale identity buffer by dataset split size: lambda * N_retain
|
|
reg_matrix = lam * X.size(0) * torch.eye(X.size(1), device=device).float()
|
|
return torch.linalg.inv(H + reg_matrix)
|
|
|
|
def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
|
|
"""
|
|
Applies Certified Removal strictly to the target class parameters
|
|
belonging to the final fully connected layer (model.fc).
|
|
"""
|
|
device = next(model.parameters()).device
|
|
k = self.target_class_index
|
|
|
|
# Isolate final layer and extract raw deep embeddings using frozen backbone
|
|
linear_head = model.fc
|
|
model.fc = nn.Identity()
|
|
|
|
print(">> Extracting deep features from model backbone...")
|
|
X_retain = self._get_features(model, retain_loader, device).to(device)
|
|
X_forget = self._get_features(model, forget_loader, device).to(device)
|
|
|
|
# Restore the classification head back
|
|
model.fc = linear_head
|
|
|
|
# Extract current model weight row for the target class channel
|
|
w_k = model.fc.weight.data[k].clone().to(device)
|
|
|
|
# Create One-vs-Rest binary target indicator arrays (+1.0 / -1.0)
|
|
# Retain dataset instances are negative labels (-1.0) for the target class channel
|
|
y_retain_binary = torch.full((X_retain.size(0),), -1.0, device=device)
|
|
# Forget dataset instances are positive labels (+1.0) for the target class channel
|
|
y_forget_binary = torch.full((X_forget.size(0),), 1.0, device=device)
|
|
|
|
# Compute Inverse Hessian (on Retain Data) and Gradient (on Forget Data)
|
|
H_inv = self._fb_lr_hessian_inv(w_k, X_retain, y_retain_binary, self.l2_reg, device)
|
|
grad_forget = self._fb_lr_grad(w_k, X_forget, y_forget_binary, self.l2_reg)
|
|
|
|
# 5. Compute the Weight Update Step Vector (Delta)
|
|
multiplier = 0.5
|
|
delta_w_k = torch.mv(H_inv, grad_forget) * multiplier
|
|
|
|
# Verify Theoretical Removal Bound Criteria
|
|
norm_delta = torch.norm(delta_w_k).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}")
|
|
|
|
# Apply Update (Using '+' since Facebook's grad calculation yields a negative sign output)
|
|
new_w_k = w_k + delta_w_k
|
|
|
|
# Calibrate and Inject Perturbation Noise (Objective Perturbation Verification)
|
|
sigma = 2.0 / (self.l2_reg * self.epsilon)
|
|
noise = torch.randn_like(new_w_k, device=device) * (sigma / X_retain.size(0))
|
|
new_w_k = new_w_k + noise
|
|
|
|
# Commit updated weight vector row back into model head parameters in-place
|
|
model.fc.weight.data[k] = new_w_k
|
|
|
|
print(">> Certified Removal process completed successfully.")
|
|
return model |