A-dist

eval/UnlearningAttack.py

@@ -46,8 +46,16 @@ class UnlearningAttack:
 
     def calculate_a_dist(self, latent1, latent2):
         """Calculates formal A-Distance: 2 * (1 - 2 * epsilon)."""
-        X = np.vstack([latent1, latent2])
+        combined = np.vstack([latent1, latent2])
-        y = np.concatenate([np.ones(len(latent1)), np.zeros(len(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)
@@ -162,31 +170,7 @@ class UnlearningAttack:
         
         return mia_accuracy, np.mean(forget_distances)
     
-    def run_logit_space_lookalike_mia(self, filtered_model, naive_retrained, forget_loader, device, target_class):
+    def _comput_adversarial_accuracy(self, filtered, naive, axis=-1):
-        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()
 
         # Z-Score Normalisation
         filtered = (filtered - np.mean(filtered, axis=-1, keepdims=True)) / (np.std(filtered, axis=-1, keepdims=True) + 1e-8)
@@ -216,7 +200,36 @@ class UnlearningAttack:
         adversary = LogisticRegression(max_iter=1000)
         adversary.fit(data_train, label_train)
         # evaluate similarity of outputs
-        lookalike_accuracy = accuracy_score(label_test, adversary.predict(data_test))
+        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)