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.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): accuracy_score = self._comput_adversarial_accuracy(latent1, latent2, axis=0) epsilon = 1 - accuracy_score return 2.0 * np.abs(0.5 - epsilon) def _extract_attack_features(self, target_model, loader, device, target_class): 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) # 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 def _comput_adversarial_accuracy(self, filtered, naive, axis=0): # Z-Score Normalisation filtered = (filtered - np.mean(filtered, axis = axis, keepdims=True)) / (np.std(filtered, axis = axis, keepdims = True) + 1e-8) naive = (naive - np.mean(naive, axis = axis, keepdims = True)) / (np.std(naive, axis = axis, 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, axis = -1) # 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") # 1. Parameter-Space MIA and Latent Distance parameter_mia_acc = 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 ) # Calculate JS-Dist js_dist = self.calculate_js_dist(unlearned_instance.model, reference_model_torch, forget_loader, device, target_class) # Extract features 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) print(f"[{framework_name}] Class {target_class} | Parameter MIA: {parameter_mia_acc:.4f} Lookalike: {lookalike_acc:.4f}" ) with open(current_log_file, "a") as f: f.write(f"{target_class},{parameter_mia_acc:.6f},{0.00000},{lookalike_acc:.6f}, {a_dist:.6f}, {js_dist:.6f}\n") return { "parameter_mia_accuracy": parameter_mia_acc, "latent_distance": a_dist, "lookalike_accuracy": lookalike_acc }