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 }