359 lines
16 KiB
Python
359 lines
16 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
import os
|
|
from scipy.spatial.distance import jensenshannon
|
|
from torch.utils.data import DataLoader
|
|
from sklearn.linear_model import LogisticRegression
|
|
from sklearn.metrics import accuracy_score
|
|
from architectures.Model import Model
|
|
|
|
class UnlearningAttack:
|
|
def __init__(self, arch, class_size):
|
|
"""
|
|
Initializes the robust verification suite with universal architecture metadata.
|
|
Matches Section 5.5 of the thesis text by implementing distinct
|
|
Parameter-Space and Logit-Space (Look-alike) attack pipelines uniformly.
|
|
"""
|
|
self.arch = arch
|
|
self.class_size = class_size
|
|
|
|
self.hook = None
|
|
self.model = None
|
|
self._hook_features = []
|
|
self.criterion = nn.CrossEntropyLoss(reduction='none')
|
|
self.collecting = False
|
|
|
|
|
|
def calculate_js_dist(self, model1, model2, loader, device, target_class):
|
|
"""Calculates Jensen-Shannon Distance between output probability distributions."""
|
|
model1.eval(); model2.eval()
|
|
probs1, probs2 = [], []
|
|
with torch.no_grad():
|
|
for data, _ in loader:
|
|
data = data.to(device)
|
|
# Handle WF_Module specific gate signal if needed
|
|
if model1.__class__.__name__ == "WF_Module":
|
|
gate = torch.full((data.size(0),), target_class, device=device)
|
|
p1 = torch.softmax(model1(data, target_class_indices=gate), dim=1)
|
|
else:
|
|
p1 = torch.softmax(model1(data), dim=1)
|
|
p2 = torch.softmax(model2(data), dim=1)
|
|
probs1.extend(p1.cpu().numpy()); probs2.extend(p2.cpu().numpy())
|
|
|
|
# JS Distance is the square root of JS Divergence
|
|
return np.mean(jensenshannon(np.array(probs1), np.array(probs2), axis=1))
|
|
|
|
def calculate_a_dist(self, latent1, latent2):
|
|
"""Calculates formal A-Distance: 2 * (1 - 2 * epsilon)."""
|
|
combined = np.vstack([latent1, latent2])
|
|
mean = np.mean(combined, axis=0)
|
|
std = np.std(combined, axis=0) + 1e-8
|
|
l1 = (latent1 - mean) / std
|
|
l2 = (latent2 - mean) / std
|
|
|
|
# 2. Use the same balanced split and regularization (C=0.01)
|
|
# as the look-alike method to ensure stability.
|
|
X = np.vstack([l1, l2])
|
|
y = np.concatenate([np.ones(len(l1)), np.zeros(len(l2))])
|
|
|
|
# Shuffle and split
|
|
idx = np.arange(len(X)); np.random.shuffle(idx)
|
|
X, y = X[idx], y[idx]
|
|
split = int(len(X) * 0.7)
|
|
|
|
clf = LogisticRegression(solver='liblinear').fit(X[:split], y[:split])
|
|
epsilon = 1.0 - accuracy_score(y[split:], clf.predict(X[split:]))
|
|
|
|
return 2.0 * np.abs(0.5 - epsilon)
|
|
|
|
|
|
|
|
def _hook_fn(self, module, input, output):
|
|
if not self.collecting:
|
|
return
|
|
flattened_embeddings = torch.flatten(output, 1)
|
|
self._hook_features.append(flattened_embeddings.detach())
|
|
|
|
def register_model_hook(self, inner_model):
|
|
if hasattr(inner_model, "original_model"):
|
|
core_model = inner_model.original_model
|
|
else:
|
|
core_model = inner_model
|
|
|
|
if hasattr(core_model, 'avgpool'):
|
|
pool_layer = core_model.avgpool
|
|
else:
|
|
pool_layer = None
|
|
for name, module in core_model.named_modules():
|
|
if 'pool' in name:
|
|
pool_layer = module
|
|
break
|
|
|
|
if pool_layer is None:
|
|
raise AttributeError("The target model architecture lacks an 'avgpool' layer block.")
|
|
|
|
self.hook = pool_layer.register_forward_hook(self._hook_fn)
|
|
|
|
def shutdown_hook(self):
|
|
if hasattr(self, 'hook') and self.hook:
|
|
self.hook.remove()
|
|
self.hook = None
|
|
|
|
def _extract_attack_features(self, target_model, loader, device, target_class):
|
|
'''target_model.eval()
|
|
all_probs = []
|
|
all_entropies = []
|
|
all_losses = []
|
|
self._hook_features = []
|
|
self.collecting = True
|
|
with torch.no_grad():
|
|
for data, targets in loader:
|
|
data, targets = data.to(device), targets.to(device)
|
|
if target_model.__class__.__name__ == "WF_Module":
|
|
gate_signals = torch.full(
|
|
(data.size(0),),
|
|
target_class,
|
|
dtype=torch.long,
|
|
device=data.device
|
|
)
|
|
outputs = target_model(data, target_class_indices=gate_signals)
|
|
else:
|
|
outputs = target_model(data)
|
|
|
|
probs = torch.softmax(outputs, dim=1)
|
|
all_probs.extend(probs.cpu().numpy())
|
|
|
|
entropy = -torch.sum(probs * torch.log(probs + 1e-10), dim=1)
|
|
all_entropies.extend(entropy.cpu().numpy())
|
|
|
|
loss = self.criterion(outputs, targets)
|
|
all_losses.extend(loss.cpu().numpy())
|
|
|
|
self.collecting = False
|
|
|
|
X_probs = np.array(all_probs)
|
|
X_entropy = np.array(all_entropies).reshape(-1, 1)
|
|
X_loss = np.array(all_losses).reshape(-1, 1)
|
|
X_features = np.hstack([X_probs, X_entropy, X_loss])
|
|
|
|
if self._hook_features:
|
|
compiled_latent = torch.cat(self._hook_features, dim=0).cpu().numpy()
|
|
else:
|
|
compiled_latent = np.zeros((len(X_features), 512))
|
|
|
|
return X_features, compiled_latent'''
|
|
target_model.eval()
|
|
all_probs, all_entropies, all_losses = [], [], []
|
|
|
|
with torch.no_grad():
|
|
for data, targets in loader:
|
|
data, targets = data.to(device), targets.to(device)
|
|
if target_model.__class__.__name__ == "WF_Module":
|
|
gate = torch.full((data.size(0),), target_class, dtype=torch.long, device=device)
|
|
outputs = target_model(data, target_class_indices=gate)
|
|
else:
|
|
outputs = target_model(data)
|
|
|
|
probs = torch.softmax(outputs, dim=1)
|
|
all_probs.extend(probs.cpu().numpy())
|
|
|
|
# Entropy: -sum(p * log(p))
|
|
log_probs = torch.log(probs + 1e-10)
|
|
entropy = -torch.sum(probs * log_probs, dim=1)
|
|
all_entropies.extend(entropy.cpu().numpy())
|
|
|
|
loss = self.criterion(outputs, targets)
|
|
all_losses.extend(loss.cpu().numpy())
|
|
|
|
# Combine output-based features only
|
|
return np.hstack([
|
|
np.array(all_probs),
|
|
np.array(all_entropies).reshape(-1, 1),
|
|
np.array(all_losses).reshape(-1, 1)
|
|
])
|
|
|
|
def run_parameter_space_mia(self, unlearned_model, shadow_model, forget_loader, retain_test_loader, device, index):
|
|
'''X_shadow_mem, _ = self._extract_attack_features(shadow_model, forget_loader, device, index)
|
|
X_shadow_non, _ = self._extract_attack_features(shadow_model, retain_test_loader, device, index)
|
|
|
|
min_train = min(len(X_shadow_mem), len(X_shadow_non))
|
|
X_train = np.vstack([X_shadow_mem[:min_train], X_shadow_non[:min_train]])
|
|
y_train = np.concatenate([np.ones(min_train), np.zeros(min_train)])
|
|
|
|
attack_classifier = LogisticRegression(max_iter=1000)
|
|
attack_classifier.fit(X_train, y_train)
|
|
|
|
X_eval_mem, latent_features = self._extract_attack_features(unlearned_model, forget_loader, device, index)
|
|
X_eval_non, retain_latent = self._extract_attack_features(unlearned_model, retain_test_loader, device, index)
|
|
|
|
min_test = min(len(X_eval_mem), len(X_eval_non))
|
|
X_test = np.vstack([X_eval_mem[:min_test], X_eval_non[:min_test]])
|
|
y_test = np.concatenate([np.ones(min_test), np.zeros(min_test)])
|
|
|
|
predictions = attack_classifier.predict(X_test)
|
|
mia_accuracy = accuracy_score(y_test, predictions)
|
|
|
|
#clean_centroid = np.mean(retain_latent, axis=0)
|
|
#forget_distances = np.linalg.norm(latent_features - clean_centroid, axis=1)
|
|
|
|
return mia_accuracy, np.mean(forget_distances)'''
|
|
X_shadow_mem = self._extract_attack_features(shadow_model, forget_loader, device, index)
|
|
X_shadow_non = self._extract_attack_features(shadow_model, retain_test_loader, device, index)
|
|
|
|
# Train MIA Classifier
|
|
min_train = min(len(X_shadow_mem), len(X_shadow_non))
|
|
X_train = np.vstack([X_shadow_mem[:min_train], X_shadow_non[:min_train]])
|
|
y_train = np.concatenate([np.ones(min_train), np.zeros(min_train)])
|
|
|
|
attack_classifier = LogisticRegression(max_iter=1000)
|
|
attack_classifier.fit(X_train, y_train)
|
|
|
|
# Evaluate MIA
|
|
X_eval_mem = self._extract_attack_features(unlearned_model, forget_loader, device, index)
|
|
X_eval_non = self._extract_attack_features(unlearned_model, retain_test_loader, device, index)
|
|
|
|
min_test = min(len(X_eval_mem), len(X_eval_non))
|
|
X_test = np.vstack([X_eval_mem[:min_test], X_eval_non[:min_test]])
|
|
y_test = np.concatenate([np.ones(min_test), np.zeros(min_test)])
|
|
|
|
predictions = attack_classifier.predict(X_test)
|
|
mia_accuracy = accuracy_score(y_test, predictions)
|
|
|
|
# Note: Latent distance is removed as it's not a black-box metric
|
|
return mia_accuracy, 0.0
|
|
|
|
def _comput_adversarial_accuracy(self, filtered, naive, axis=-1):
|
|
|
|
# Z-Score Normalisation
|
|
filtered = (filtered - np.mean(filtered, axis=-1, keepdims=True)) / (np.std(filtered, axis=-1, keepdims=True) + 1e-8)
|
|
naive = (naive - np.mean(naive, axis=-1, keepdims=True)) / (np.std(naive, axis=-1, keepdims=True) + 1e-8)
|
|
|
|
# shuffle indices
|
|
num_images = len(filtered)
|
|
image_indices = np.arange(num_images)
|
|
np.random.shuffle(image_indices)
|
|
|
|
# split to train and test set
|
|
split = int(num_images * 0.7)
|
|
train_img_idx, test_img_idx = image_indices[:split], image_indices[split:]
|
|
|
|
# create a balanced test and train set
|
|
data_train = np.vstack([filtered[train_img_idx], naive[train_img_idx]])
|
|
data_test = np.vstack([filtered[test_img_idx], naive[test_img_idx]])
|
|
|
|
# labels for attcker (1 from unlearned and 0s to retrained)
|
|
# we do this because every output retrained gives us is a result of unseen
|
|
# but unlearned has seen these data.
|
|
label_train = np.concatenate([np.ones(len(train_img_idx)), np.zeros(len(train_img_idx))])
|
|
# test set
|
|
label_test = np.concatenate([np.ones(len(test_img_idx)), np.zeros(len(test_img_idx))])
|
|
|
|
|
|
adversary = LogisticRegression(max_iter=1000)
|
|
adversary.fit(data_train, label_train)
|
|
# evaluate similarity of outputs
|
|
return accuracy_score(label_test, adversary.predict(data_test))
|
|
|
|
def run_logit_space_lookalike_mia(self, filtered_model, naive_retrained, forget_loader, device, target_class):
|
|
filtered_model.eval()
|
|
naive_retrained.eval()
|
|
|
|
filtered_logits = []
|
|
naive_logits = []
|
|
|
|
with torch.no_grad():
|
|
for data, _ in forget_loader:
|
|
data = data.to(device)
|
|
|
|
if filtered_model.__class__.__name__ == "WF_Module":
|
|
gate_signals = torch.full((data.size(0),), target_class, dtype=torch.long, device=data.device)
|
|
out_filtered = filtered_model(data, target_class_indices=gate_signals)
|
|
else:
|
|
out_filtered = filtered_model(data)
|
|
|
|
out_naive = naive_retrained(data)
|
|
|
|
filtered_logits.append(out_filtered)
|
|
naive_logits.append(out_naive)
|
|
|
|
# Concatenate everything
|
|
filtered = torch.cat(filtered_logits, dim=0).cpu().numpy()
|
|
naive = torch.cat(naive_logits, dim=0).cpu().numpy()
|
|
|
|
# evaluate similarity of outputs
|
|
lookalike_accuracy = self._comput_adversarial_accuracy(filtered=filtered, naive=naive)
|
|
# so that the metric is between 0 and 1.
|
|
return 2.0 * np.abs(lookalike_accuracy - 0.5)
|
|
|
|
def run_complete_evaluation(self, framework_name, target_class, forget_loader, retain_test_loader, unlearned_instance, base_shadow_instance, device):
|
|
"""Orchestrates specific pipeline routing cleanly using cached constructor parameters."""
|
|
target_dir = os.path.join("reports", framework_name)
|
|
os.makedirs(target_dir, exist_ok=True)
|
|
current_log_file = os.path.join(target_dir, "attack_values.csv")
|
|
|
|
if not os.path.exists(current_log_file):
|
|
with open(current_log_file, "w") as f:
|
|
f.write("target_class, parameter_mia_accuracy, latent_distance_tell, lookalike_accuracy, A-Dist, JS-Dist\n")
|
|
|
|
self.register_model_hook(unlearned_instance.model)
|
|
|
|
# 1. Parameter-Space MIA and Latent Distance
|
|
parameter_mia_acc, latent_dist = self.run_parameter_space_mia(
|
|
unlearned_model=unlearned_instance.model,
|
|
shadow_model=base_shadow_instance.model,
|
|
forget_loader=forget_loader,
|
|
retain_test_loader=retain_test_loader,
|
|
device=device,
|
|
index=target_class
|
|
)
|
|
|
|
# we load a retrained model to evaluate look_alike tests
|
|
ghost_checkpoint_path = f"trained_models/class_{target_class}_retrained.pth"
|
|
|
|
if os.path.exists(ghost_checkpoint_path):
|
|
# Safe clean wrapper boot utilizing internal instance state properties
|
|
ghost_model_instance = Model.create(arch=self.arch, device=device, size=self.class_size)
|
|
state_dict = torch.load(ghost_checkpoint_path, map_location=device, weights_only=True)
|
|
ghost_model_instance.model.load_state_dict(state_dict)
|
|
reference_model_torch = ghost_model_instance.model
|
|
else:
|
|
raise FileNotFoundError(
|
|
f"Retrained weights not found at: {ghost_checkpoint_path}. \nDid you forget to save models or are they saved with a different path?"
|
|
)
|
|
|
|
lookalike_acc = self.run_logit_space_lookalike_mia(
|
|
filtered_model=unlearned_instance.model,
|
|
naive_retrained=reference_model_torch,
|
|
forget_loader=forget_loader,
|
|
device=device,
|
|
target_class=target_class
|
|
)
|
|
|
|
# 1. Calculate JS-Dist (Logit-space probability comparison)
|
|
js_dist = self.calculate_js_dist(unlearned_instance.model, reference_model_torch, forget_loader, device, target_class)
|
|
|
|
# 2. Extract latent features for A-Dist
|
|
# We need features from both Unlearned and Retrained model
|
|
#_, unlearned_latent = self._extract_attack_features(unlearned_instance.model, forget_loader, device, target_class)
|
|
#_, retrained_latent = self._extract_attack_features(reference_model_torch, forget_loader, device, target_class)
|
|
# Extract features (now just one returned object)
|
|
unlearned_features = self._extract_attack_features(unlearned_instance.model, forget_loader, device, target_class)
|
|
retrained_features = self._extract_attack_features(reference_model_torch, forget_loader, device, target_class)
|
|
|
|
# Calculate A-Dist using these features
|
|
a_dist = self.calculate_a_dist(unlearned_features, retrained_features)
|
|
# 3. Calculate A-Dist (Replacing latent_distance)
|
|
#a_dist = self.calculate_a_dist(unlearned_latent, retrained_latent)
|
|
|
|
|
|
print(f"[{framework_name}] Class {target_class} | Parameter MIA: {parameter_mia_acc:.4f} | Latent Dist: {latent_dist:.4f} | Lookalike: {lookalike_acc:.4f}" )
|
|
|
|
with open(current_log_file, "a") as f:
|
|
f.write(f"{target_class},{parameter_mia_acc:.6f},{latent_dist:.6f},{lookalike_acc:.6f}, {a_dist:.6f}, {js_dist:.6f}\n")
|
|
|
|
return {
|
|
"parameter_mia_accuracy": parameter_mia_acc,
|
|
"latent_distance": latent_dist,
|
|
"lookalike_accuracy": lookalike_acc
|
|
} |