Tinsae/Finetuning
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

unlearning done

Browse Source
This commit is contained in:
Tinsae
2026-06-27 20:38:17 +02:00
parent c4fdc034b2
commit 0680a920ff
11 changed files with 307 additions and 740 deletions
Download Patch File Download Diff File Expand all files Collapse all files

98
Tune_new.py
View File

@@ -9,11 +9,9 @@ import Util
from sets.Data import *
from sets.IdentitySubset import IdentitySubset
from architectures.Model import Model, Architecture
from unlearning.CertifiedRemoval import CertifiedRemoval
from unlearning.CertifiedUnlearning import CertifiedUnlearning
from unlearning.LinearFiltration import LinearFiltration
from unlearning.WeightFiltration import WeightFiltration
from unlearning.WF import WeightF
# Global Hyperparameters
@@ -140,40 +138,12 @@ def run_unlearning_and_strategy_eval(env_dict, forget_class_idx, strategy, evalu
train_data = env_dict["train_data"]
test_data = env_dict["test_data"]
# testing valuse * *
#---------------------------------------------------------------------------
# S1 50 5 5 5 5 5
# S2 1000 200 1000 500 200 300
# BS 5 5 5 5 5 5
# scale 2000 500 8000 5000 10000 8000
# std 0.00001 0.00001 0.00001 0.00001 0.00001 0.00001
# Initialize the strategy hyperparameters matching standard settings
# increase s2, decrease scale ---sweet spot
'''certified_removal = CertifiedRemoval(
target_class_index=forget_class_idx,
s1=4,
s2=350, # 350 best
unlearn_bs=5,
scale=6000.0, # 6000 was good
std=0.00001
)'''
'''certified_removal = CertifiedUnlearning(
target_class_index=0,
l2_reg=0.0005,
gamma=0.1,
scale=7000.0,
s1=2,
s2=350,
std=1e-5,
unlearn_bs=2
)'''
# Segment specific unlearning loaders using class index boundaries
forget_train_loader, retain_train_loader = get_unlearning_loaders(
retain_train_loader , forget_train_loader= get_unlearning_loaders(
dataset=train_data, forget_class_idx=forget_class_idx, batch_size=BATCH_SIZE
)
forget_test_loader, retain_test_loader = get_unlearning_loaders(
retain_test_loader, forget_test_loader = get_unlearning_loaders(
dataset=test_data, forget_class_idx=forget_class_idx, batch_size=BATCH_SIZE
)
@@ -189,9 +159,16 @@ def run_unlearning_and_strategy_eval(env_dict, forget_class_idx, strategy, evalu
print("fine tunned model loaded into evaluation sandbox")
# Execute strategic parameter unlearning step
strategy.apply(reloaded.model, forget_train_loader, retain_train_loader)
unlearned = strategy.apply(reloaded.model, train_data)
strategy_in_use = strategy.__class__.__name__
if isinstance(unlearned,nn.Module):
reloaded.model = unlearned
else:
reloaded = unlearned
# Define validation tracking steps dynamically
evaluation_domains = [
{"loader": retain_test_loader, "mode": "retain", "label": "\n--- Performance on Retained Classes"},
@@ -215,66 +192,63 @@ def run_unlearning_and_strategy_eval(env_dict, forget_class_idx, strategy, evalu
# entry
if __name__ == "__main__":
try:
# Run Data Infrastructure and Architecture Builder
runtime_environment = prepare_data_and_model_environment()
# Baseline Evaluation
finetuning = False
# switch finetuning for tests on strategies only
run_finetuning_or_baseline_eval(runtime_environment, run_training=finetuning)
run_finetuning_or_baseline_eval(runtime_environment, run_training = finetuning)
finetuning = True
# Unlearning Iterations
for i in range(0, 1):
# strategies
#
#certified_removal = CertifiedRemoval(
# target_class_index=i,
# s1=4,
# s2=350, # 350 best
# unlearn_bs=5,
# scale=6000.0, # 6000 was good
# std=0.00009
# )
certified_unlearning = CertifiedUnlearning(
target_class_index=i,
target_class_index=0,
l2_reg=0.000002,
gamma=0.1,
scale= 20000,# 16400.0, # took ages to reach this sweet spot
scale= 16400.0,# 16400.0, # took ages to reach this sweet spot
s1=2,
s2=300,
std=0.00001,
unlearn_bs=16
unlearn_bs=8
)
# works perfectly
linear_filtration = LinearFiltration(
target_class_index=i
target_class_index=0
)
weight_filtration = WeightF( #WeightFiltration(
target_class_index=i,
epochs=3,
lr=0.05,
gamma=5
weight_filtration = WeightFiltration(
target_class_index=0,
epochs=6,
lr=150.0,
gamma=0.001
)
strategies = [
# certified_unlearning,
certified_unlearning,
weight_filtration,
# linear_filtration
linear_filtration
]
# Unlearning Iteration
for i in range(0, CLASS_SIZE):
print(f"\n>>> Executing Unlearning Framework for Target Identity Index: {i} <<<")
for strategy in strategies:
# update target class to be unlearned
strategy.set_target_class(i)
print(f"Unlearning class {i} with {strategy.strategy_name}")
# forget
run_unlearning_and_strategy_eval(
runtime_environment,
forget_class_idx=i,
strategy=strategy,
evaluate= not finetuning
evaluate = not finetuning
)
except KeyboardInterrupt:
print("program interrupted. Exit!")

3
Util.py
View File

@@ -46,3 +46,6 @@ def log_metric(log_file, execution_time: float):
"""Appends the execution time to this strategy's specific file."""
with open(log_file, "a") as f:
f.write(f"{execution_time:.6f}\n")

5
architectures/Model.py
View File

@@ -7,7 +7,7 @@ import time
import numpy as np
from sklearn.metrics import classification_report
from pathlib import Path
from unlearning.Strategy import Strategy
#from unlearning.Strategy import Strategy
import copy
from torch.optim.lr_scheduler import CosineAnnealingLR
@@ -84,7 +84,7 @@ class Model(ABC):
print(f'Model loaded from {file_path}')
def unlearn(self, strategy: Strategy, forget_loader, retain_loader):
def unlearn(self, strategy: 'Strategy', forget_loader, retain_loader):
""" Executes a targeted unlearning strategy and profiles efficiency """
print(f"Executing: {strategy.__class__.__name__}...")
@@ -103,6 +103,7 @@ class Model(ABC):
Evaluates the model, prints terminal reports, and routes metrics to
a file logger based on the current context mode.
"""
self.model.eval()
all_preds, all_labels = [], []
print(f"\nEvaluating Domain: [{mode}]...")

59
sets/Data.py
View File

@@ -1,5 +1,5 @@
from torchvision import datasets, transforms
from torch.utils.data import Dataset, DataLoader, Subset
from torch.utils.data import Dataset, DataLoader, Subset, ConcatDataset
import torch
import numpy as np
import os
@@ -181,4 +181,59 @@ def get_unlearning_loaders(dataset: Dataset, forget_class_idx: int, batch_size:
print(f"[Data Split] Local Class {forget_class_idx}: {len(forget_subset)} samples | Remaining Classes: {len(retain_subset)} samples.")
return forget_loader, retain_loader
return retain_loader, forget_loader
def vertical_split(dataset, batch_size,num_classes):
"""
Executes a class-wise vertical split.
Divides the samples of every single identity class exactly in half:
50% of each class goes to the Retain Set, 50% goes to the Forget Set.
"""
# 1. Group dataset indices by their respective ground-truth classes
class_to_indices = {c: [] for c in range(num_classes)}
print(" [Vertical Split] Tracking class indices across the combined dataset...")
for idx in range(len(dataset)):
# Extract the label cleanly from the underlying dataset structure
_, label = dataset[idx]
if label in class_to_indices:
class_to_indices[label].append(idx)
retain_indices = []
forget_indices = []
# 2. Slice each class identity vertically (exactly 50/50)
for c, indices in class_to_indices.items():
if len(indices) < 2:
print(f" Warning: Class {c} has fewer than 2 samples. Cannot split vertically.")
retain_indices.extend(indices)
continue
# Deterministic shuffle per class to ensure honest distribution before splitting
np.random.shuffle(indices)
mid = len(indices) // 2
forget_indices.extend(indices[:mid]) # First half assigned to unlearning
retain_indices.extend(indices[mid:]) # Second half assigned to retention
print(f" Vertical split complete: Retain Index Size = {len(retain_indices)} | Forget Index Size = {len(forget_indices)}")
# 3. Construct lightweight PyTorch Subsets using our sliced index maps
retain_subset = Subset(dataset, retain_indices)
forget_subset = Subset(dataset, forget_indices)
# 4. Return pristine, shuffled DataLoaders mirroring your environment's batch specifications
retain_loader = DataLoader(retain_subset, batch_size=batch_size, shuffle=True)
forget_loader = DataLoader(forget_subset, batch_size=batch_size, shuffle=True)
return retain_loader, forget_loader
def _combine_set(loader_one, loader_two):
full_train_dataset = ConcatDataset([loader_one.dataset, loader_two.dataset])
return DataLoader(
full_train_dataset,
batch_size=loader_one.batch_size,
shuffle=True
)

214
unlearning/CertifiedRemoval.py
View File

@@ -1,214 +0,0 @@
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, RandomSampler
from torch.autograd import grad
from unlearning.Strategy import Strategy
class CertifiedRemoval(Strategy):
"""
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, 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 _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
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)
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:
device = next(model.parameters()).device
num_forget = len(forget_loader.dataset)
num_retain = len(retain_loader.dataset)
scaling_ratio = num_forget / num_retain
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)
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
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

123
unlearning/Certified_facebook.py
View File

@@ -1,123 +0,0 @@
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

125
unlearning/LastK_Certified.py
View File

@@ -1,125 +0,0 @@
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

135
unlearning/LinearFiltration.py
View File

@@ -2,6 +2,7 @@ import torch
import torch.nn as nn
from .Strategy import Strategy
from torch.utils.data import DataLoader
from sets.Data import get_unlearning_loaders, _combine_set
class LinearFiltration(Strategy):
def __init__(self, target_class_index):
@@ -23,40 +24,8 @@ class LinearFiltration(Strategy):
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
for j in range(num_classes):
if j == forget_class:
A[forget_class, j] = 0.0
else:
A[forget_class, j] = 1.0 / num_remaining
return A
@staticmethod
def _sums_and_counts(model, num_classes, retain_loader, forget_loader, device, forget_index, h_dim):
def _sums_and_counts(self, model, num_classes, loader, device, forget_index, h_dim):
model.eval()
sums = torch.zeros(num_classes, h_dim, device=device)
@@ -64,11 +33,11 @@ class LinearFiltration(Strategy):
# Generate values for retain
with torch.no_grad():
for inputs, targets in retain_loader:
for inputs, targets in loader:
inputs = inputs.to(device)
targets = targets.to(device)
# FIX: Call get_features instead of model() directly
outputs = LinearFiltration.get_features(model, inputs)
# predictions
outputs = model(inputs)
for j in range(num_classes):
if j == forget_index:
@@ -79,65 +48,54 @@ class LinearFiltration(Strategy):
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(
#
def _get_means(self,model, num_classes, loader, device, forget_index):
h_dim = model.fc.out_features
# all predictions
sums, counts = self._sums_and_counts(
model=model,
num_classes=num_classes,
retain_loader=retain_loader,
forget_loader=forget_loader,
loader=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))
#A = []
# CORRECT: Stack along dim=0 to make it (num_classes, h_dim)
return torch.stack(A, dim=0)
counts_safe = counts.unsqueeze(1)
A = torch.where(
counts_safe > 0,
sums / counts_safe,
torch.zeros_like(sums)
)
# 6
return A
@staticmethod
def _compute_z(tensor, forget_index):
# Now tensor has shape (num_classes, h_dim) -> tensor.shape[0] is num_classes
# 9
def _compute_z(self, tensor, forget_index):
K = tensor.shape[0]
# pi_a0 should match the feature space dimensions (h_dim)
pi_a0 = torch.zeros(tensor.shape[1], device=tensor.device)
# pi_a_forget should match the feature space dimensions (h_dim)
pi_a_f = torch.zeros(tensor.shape[1], device=tensor.device)
t_1 = pi_a0
a0 = tensor[forget_index, :] # Extracting the row vector for the forgotten class
t_1 = pi_a_f
# Extracting the row vector for the forgotten class
a_f = tensor[forget_index, :]
mask_a0 = torch.ones(
a0.shape[0],
mask_a_f = torch.ones(
a_f.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()
t_2 = -(1.0 / (K - 1)) * a_f[mask_a_f].sum()
mask_rows = torch.ones(K, dtype=torch.bool, device=tensor.device)
mask_rows[forget_index] = False
@@ -148,21 +106,23 @@ class LinearFiltration(Strategy):
return t_1 + t_2 + t_3
@staticmethod
def normalise(model, retain_loader, forget_loader, device, forget_index):
# Normalisation filtration
def normalise(self, model, retain_loader, forget_loader, device, forget_index):
W = model.fc.weight.data.clone()
num_classes = W.shape[0]
A = LinearFiltration._get_means(
# we combine the data so we can calculate the mean of prdictions
full_loader = _combine_set(retain_loader, forget_loader)
# 8
A = self._get_means(
model=model,
num_classes=num_classes,
retain_loader=retain_loader,
forget_loader=forget_loader,
loader=full_loader,
device=device,
forget_index=forget_index
)
Z = LinearFiltration._compute_z(tensor=A, forget_index=forget_index)
# 9
Z = self._compute_z(tensor=A, forget_index=forget_index)
B_Z_rows = []
for i in range(num_classes):
@@ -172,13 +132,24 @@ class LinearFiltration(Strategy):
# Retained classes maintain their original ideal feature directions
B_Z_rows.append(A[i])
# 10
# Stack back along dim=0 to match (num_classes, h_dim)
# to get mean
B_Z = torch.stack(B_Z_rows, dim=0)
A_inv = torch.linalg.pinv(A)
# 11
W_Z = B_Z @ A_inv @ W
# 12
model.fc.weight.copy_(W_Z)
return model
# overriden function
def _split_data(self, dataset):
return get_unlearning_loaders(
dataset=dataset,
forget_class_idx=self.target_class_index,
batch_size = 32
)

18
unlearning/Strategy.py
View File

@@ -16,13 +16,21 @@ class Strategy:
self.log_file = Path(f"reports/{self.strategy_name}/metrics.txt")
Util._initialize_log_file(log_file= self.log_file)
def apply(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
def set_target_class(self, target_class_index: int):
"""Dynamically switch the unlearning target without retraining."""
self.target_class_index = target_class_index
def apply(self, model: nn.Module, dataset) -> nn.Module:
"""
Wraps the unlearning execution with automated timing and strategy-specific logging.
DO NOT override this method in subclasses. Override _run instead.
"""
start_time = time.perf_counter()
retain_loader, forget_loader = self._split_data(dataset)
# Execute core unlearning logic
processed_model = self._run(model, forget_loader, retain_loader)
@@ -42,3 +50,11 @@ class Strategy:
def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
"""Subclasses implement their core unlearning logic here."""
raise NotImplementedError
'''
different strategies split data in to different partitions differently.
So a strategy will implement its own and since this part is startegy specific.
not all should compute it the same.
'''
def _split_data(self,dataset):
pass

167
unlearning/WeightFiltration.py
View File

@@ -1,126 +1,135 @@
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.utils.data import DataLoader, ConcatDataset, Subset
from unlearning.Strategy import Strategy
from .wf.WF_Net import WF_Net
import numpy as np
from sklearn.metrics import classification_report
from architectures.WFNet import WF_Net_Model
from sets.Data import vertical_split
class WeightFiltration(Strategy):
"""
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: int, epochs: int = 10, lr: float = 0.2, gamma: float = 10.0):
def __init__(self,
target_class_index: int,
num_classes: int = 20,
epochs: int = 6,
lr: float = 100.0,
gamma: float = 0.01,
):
super().__init__(target_class_index=target_class_index)
self.epochs = epochs
self.lr = lr
self.gamma = gamma
#self.alpha = None
self.num_classes = num_classes
self.wf_model = None
self.lambda_1 = 25
def _optimise_filter(self, model: nn.Module, retain_loader: DataLoader, forget_loader: DataLoader, device) -> nn.Module:
# new WF_Model instance
wf_model = WF_Net_Model(
device=device,
size=self.num_classes,
original_model=model,
target_class_index=self.target_class_index
)
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)
# 2. ONLY optimize alpha (everything else is frozen inside the wrapper)
optimizer = optim.Adam([wf_model.alpha], lr=self.lr)
# a WF_net module to be trained (unlearned) to generate alpha
wf_net = wf_model.get()
optimizer = optim.SGD([wf_net.alpha], lr=self.lr)
criterion = nn.CrossEntropyLoss()
for epoch in range(self.epochs):
forget_iter = iter(forget_loader)
t_loss_r, t_loss_f = 0.0, 0.0
steps = 0
for r_inputs, r_labels in retain_loader:
# forget and retain
for (r_inputs, r_labels), (f_inputs, f_labels) in zip(retain_loader, forget_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:
forget_iter = iter(forget_loader)
f_inputs, _ = next(forget_iter)
f_inputs = f_inputs.to(device)
f_inputs, f_labels = f_inputs.to(device), f_labels.to(device)
optimizer.zero_grad()
# --- 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)
# retain data paired with randomly selected rows of alpha to compute the retaining loss
random_rows = []
for label in r_labels:
allowed = [i for i in range(self.num_classes) if i != label.item()]
random_rows.append(np.random.choice(allowed))
# Compute Losses using Poppi et al.'s temperature scaled entropy
gate_signals_r = torch.tensor(random_rows, dtype=torch.long, device=device)
outputs_r = wf_net(r_inputs, target_class_indices=gate_signals_r)
loss_r = criterion(outputs_r, r_labels)
temperature = 3.0
logits_f_scaled = outputs_f / temperature
# Forget set is paired with corresponding labels as row selectors for alpha
# and used to compute unlearning loss
outputs_f = wf_net(f_inputs, target_class_indices=f_labels)
# Compute uniform target entropy per-sample, then average over the batch
log_probs_f = torch.log_softmax(logits_f_scaled, dim=-1)
uniform_target = torch.ones_like(logits_f_scaled) / num_classes
loss_f = -torch.sum(uniform_target * log_probs_f, dim=-1).mean()
loss_f = 0.0
classes_in_batch = 0
total_loss = loss_r + (self.gamma * loss_f)
# every image of class c will unlearn over the same row of alpha_l (poppi et al page 5)
for c in range(self.num_classes):
class_mask = (f_labels == c)
if not class_mask.any():
continue
labels_c = f_labels[class_mask]
# Slice the existing outputs instead of recalculating a forward pass
outputs_f_c = outputs_f[class_mask]
loss_f_ce = criterion(outputs_f_c, labels_c)
# Poppi et al. suggest employing reciprocal of the forget loss
# to avoid shortcomings of negative gradient approach
loss_f += 1.0 / (loss_f_ce + 1e-6)
classes_in_batch += 1
# Average forget loss by number of distinct classes seen in this batch
if classes_in_batch > 0:
loss_f = loss_f / classes_in_batch
# Regilarisation penality
loss_reg = torch.sum(1.0 - torch.sigmoid(wf_net.alpha))
# back propagation
total_loss = loss_r + (self.lambda_1 * loss_f) + (self.gamma * loss_reg)
total_loss.backward()
optimizer.step()
t_loss_r += loss_r.item()
t_loss_f += loss_f.item()
t_loss_f += loss_f.item() if classes_in_batch > 0 else 0.0
steps += 1
print(f" Epoch {epoch+1}/{self.epochs} | Retain Loss: {t_loss_r/steps:.4f} | Forget Loss: {t_loss_f/steps:.4f}")
return wf_model
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.")
if self.wf_model is None:
print(">> Initializing and compiling global WF-Net matrix (Run Once for all classes)...")
# 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(
self.wf_model = self._optimise_filter(
model,
forget_loader=forget_loader,
retain_loader=retain_loader,
device=device,
forget_loader=forget_loader,
device=device
)
else:
print(f">> Gating matrix loaded. Switching layout to target class index: {self.target_class_index}")
self.wf_model.target_class_index = self.target_class_index
return self.wf_model
def _split_data(self, dataset):
return vertical_split(
dataset= dataset,
batch_size=32,
num_classes=self.num_classes
)
# --- PERMANENT BAKING STEP ---
with torch.no_grad():
# Grab the alpha mask vector for the forgotten class and cast to 4D tensor shape
final_mask = torch.sigmoid(wf_model.alpha[self.target_class_index]).view(-1, 1, 1, 1)
# Apply filter masking permanently back onto the base layer
target_conv.weight.copy_(original_weights * final_mask)
# Unfreeze architecture parameters for evaluations downstream
for p in model.parameters():
p.requires_grad = True
print(f">> Permanently altered {out_channels} convolutional filters in layer4 via WF-Net.")
return model

38
unlearning/wf/WF_Net.py
View File

@@ -35,50 +35,50 @@ class WF_Net(nn.Module):
#self.alpha = nn.Parameter(torch.ones(num_classes, out_channels) * 1.5)
self.alpha = nn.Parameter(torch.ones(num_classes, out_channels))
def forward(self, x: torch.Tensor, target_unlearn_class: int) -> torch.Tensor:
"""
Implements Algorithm 1: General forward step of a WF model
Inputs:
x: Input tensor (Xin)
target_unlearn_class: The class index we are actively filtering out (Yunl)
"""
def forward(self, x: torch.Tensor, target_class_indices: torch.Tensor) -> torch.Tensor:
# 1. Run through early sequence of layers undisturbed
x = self.maxpool(self.relu(self.bn1(self.conv1(x))))
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
# Run layer4 block 0 and block 1 conv1 normally
# Run layer4 block 0 normally
x = self.layer4[0](x)
# -------------------------------------------------------------
# HERE IT IS: Save the structural skip connection (identity)
# BEFORE modifying features via block 1's convolutions
# -------------------------------------------------------------
identity = x
# Now enter layer4 block 1
x = self.layer4[1].conv1(x)
x = self.layer4[1].bn1(x)
x = self.layer4[1].relu(x)
# 2. CORE WF-NET MATH: w_hat_l <- alpha_l[Yunl] ⊙ w_l
# Extract 1D vector for target class and reshape to (out_channels, 1, 1, 1) for 4D convolution broadcasting
mask = torch.sigmoid(self.alpha[target_unlearn_class]).view(-1, 1, 1, 1)
w_hat = self.original_w * mask
# [Your Step 1 Masking Math happens right here...]
batch_alpha = self.alpha[target_class_indices]
mask = torch.sigmoid(batch_alpha).view(x.size(0), -1, 1, 1)
# 3. Pass gated weights straight to functional forward pass: l(Xi, w_hat_l)
# Run the functional convolution
x = F.conv2d(
x,
weight=w_hat,
weight=self.original_w,
bias=self.target_conv.bias,
stride=self.target_conv.stride,
padding=self.target_conv.padding
)
# Apply your WF-Net channel mask
x = x * mask
x = self.layer4[1].bn2(x)
# Handle residual shortcut skip connection manually since we opened up block 1
# In ResNet-18 layer4, block 1 has no downsample shortcut layer; it's a direct identity add
# -------------------------------------------------------------
# HERE IT IS USED: Add the pristine identity back to the gated output
# -------------------------------------------------------------
x = self.layer4[1].relu(x + identity)
# 4. Final Classification Head Sequence
# Final Classification Head Sequence
x = self.avgpool(x)
x = torch.flatten(x, 1)
y_out = self.fc(x)