diff --git a/eval/UnlearningAttack.py b/eval/UnlearningAttack.py index 5a381a3..96743d4 100644 --- a/eval/UnlearningAttack.py +++ b/eval/UnlearningAttack.py @@ -2,6 +2,7 @@ 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 @@ -23,6 +24,43 @@ class UnlearningAttack: 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).""" + X = np.vstack([latent1, latent2]) + y = np.concatenate([np.ones(len(latent1)), np.zeros(len(latent2))]) + + # 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 @@ -190,7 +228,7 @@ class UnlearningAttack: 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\n") + f.write("target_class, parameter_mia_accuracy, latent_distance_tell, lookalike_accuracy, A-Dist, JS-Dist\n") self.register_model_hook(unlearned_instance.model) @@ -226,10 +264,22 @@ class UnlearningAttack: target_class=target_class ) - print(f"[{framework_name}] Class {target_class} | Parameter MIA: {parameter_mia_acc:.4f} | Latent Dist: {latent_dist:.4f} | Lookalike: {lookalike_acc:.4f}") + # 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) + + # 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}\n") + 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,