added pi

unlearning/LinearFiltration.py

@@ -108,6 +108,21 @@ class LinearFiltration(Strategy):
 
         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):
@@ -135,10 +150,11 @@ class LinearFiltration(Strategy):
 
         for i in range(self.num_classes):
             if i == forget_index:
-                B_Z_rows.append(Z)
+                # Normalise the 'erased' target vector Z
+                B_Z_rows.append(self._pi(Z.unsqueeze(0)).squeeze(0))
             else:
-                # Retained classes maintain their original ideal feature directions
-                B_Z_rows.append(self.A[i])
+                # Normalise the retained class vectors
+                B_Z_rows.append(self._pi(self.A[i].unsqueeze(0)).squeeze(0))
         
         # 10
         # Stack back along dim=0 to match (num_classes, h_dim)