fix pi

unlearning/LinearFiltration.py

@@ -80,49 +80,38 @@ class LinearFiltration(Strategy):
             torch.zeros_like(sums)
         )
 
+    def pi_mask(self, index, tensor):
+        mask = torch.ones(self.num_classes, dtype= torch.bool,device = tensor.device)
+        mask[index] = False
+        return tensor[:, mask]
+
     # 9
     def _compute_z(self, tensor, forget_index):
-
+      
         K = tensor.shape[0] 
+        a_pi = self.pi_mask(tensor = tensor, index = forget_index)
         
         
-        pi_a_f = torch.zeros(tensor.shape[1], device=tensor.device)
+        #pi_a_f = torch.zeros(tensor.shape[1], device=tensor.device)
 
-        t_1 = pi_a_f
+        t_1 = a_pi[forget_index] #pi_a_f
         # row vector for the forgotten class
-        a_f = tensor[forget_index, :] 
+        #a_f = tensor[forget_index, :] 
 
-        mask_a_f = torch.ones(
+        
-            a_f.shape[0],
-            dtype=torch.bool, 
-            device=tensor.device
-        )
         # We compute the target shift over features
-        t_2 = -(1.0 / (K - 1)) * a_f[mask_a_f].sum()
+        t_2 = (1.0 / (K - 1)) * t_1.sum()
 
-        mask_rows = torch.ones(K, dtype=torch.bool, device=tensor.device)
+        mask = torch.ones(self.num_classes, dtype= torch.bool,device = tensor.device)
-        mask_rows[forget_index] = False
+        mask[forget_index] = False
+        remaining_rows = a_pi[mask]
 
-        r_A = tensor[mask_rows, :]
+        #r_A = tensor[mask_rows, :]
-        t_3 = (1.0 / ((K - 1)) ** 2) * r_A.sum()
+        t_3 = (1.0 / ((K - 1)) ** 2) * remaining_rows.sum()
+
+        return t_1 - t_2 + t_3
+       
 
-        return t_1 + t_2 + t_3
-    
-    def _pi(self, a_tensor):
-        """
-        Affine transformation to normalize the logit distribution.
-        This maps the logit mean to 0 and scales based on variance 
-        to prevent logit shrinkage/expansion 'scars'.
-        """
-        # Calculate mean and std across the feature dimension
-        mean = a_tensor.mean(dim=0, keepdim=True)
-        std = a_tensor.std(dim=0, keepdim=True)
-        
-        # Avoid division by zero
-        std = torch.where(std > 1e-6, std, torch.ones_like(std))
-        
-        return (a_tensor - mean) / std
-    
 
     # Normalisation filtration
     def normalise(self, model, retain_loader, forget_loader, device, forget_index):
@@ -148,20 +137,21 @@ class LinearFiltration(Strategy):
         Z = self._compute_z(tensor=self.A, forget_index=forget_index)
         B_Z_rows = []
 
+        a_pi = self.pi_mask(index=forget_index, tensor=self.A)
+
         for i in range(self.num_classes):
             if i == forget_index:
-                # Normalise the 'erased' target vector Z
+                B_Z_rows.append(Z)
-                B_Z_rows.append(self._pi(Z.unsqueeze(0)).squeeze(0))
             else:
-                # Normalise the retained class vectors
+                # Retained classes maintain their original ideal feature directions
-                B_Z_rows.append(self._pi(self.A[i].unsqueeze(0)).squeeze(0))
+                B_Z_rows.append(a_pi[i])
         
         # 10
         # Stack back along dim=0 to match (num_classes, h_dim)
         # to get mean
         B_Z = torch.stack(B_Z_rows, dim=0)
 
-        A_inv = torch.linalg.pinv(self.A)
+        A_inv = torch.linalg.pinv(a_pi)
         # 11 
         W_Z = B_Z @ A_inv @ W