added A-Dist and JS-Dist

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,