ready

unlearning/LinearFiltration.py

@@ -117,10 +117,6 @@ class LinearFiltration(Strategy):
     def normalise(self, model, retain_loader, forget_loader, device, forget_index):
         clf = self._get_classifier(model)
         W = clf.weight.data.clone()
-        #num_classes = W.shape[0]
-
-        # we combine the data so we can calculate the mean of prdictions
-        #full_loader = _combine_set(retain_loader, forget_loader)
         # 8
         # Computing A is the most resource intensive part of this algorithm
         # and to optimise the process, we computr it only once and re-use it
@@ -135,30 +131,42 @@ class LinearFiltration(Strategy):
 
         # 9
         Z = self._compute_z(tensor=self.A, forget_index=forget_index)
-        B_Z_rows = []
 
+        # we transpose A for column based math
+        A_t = self.A.t()
+        # pi(a) drops the forget index logit dim( K x K-1)
         a_pi = self.pi_mask(index=forget_index, tensor=self.A)
-
-        for i in range(self.num_classes):
-            if i == forget_index:
-                B_Z_rows.append(Z)
-            else:
-                # Retained classes maintain their original ideal feature directions
-                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)
+        # construct B_z
+        B_r = a_pi.clone()
+        B_r[forget_index] = Z
+        B_z = B_r.t()
 
-        A_inv = torch.linalg.pinv(a_pi)
+
+        # A inverse
+        A_inv = torch.linalg.pinv(A_t, rcond=1e-2)
         # 11 
-        W_Z = B_Z @ A_inv @ W
+        W_temp = B_z @ A_inv @ W 
+        # with one class removed, we have a  head W_temp for 19 classes.
+        # but we have loaders for 20 classes for evaluation. so ,
+        # map K-1  W back to K x K.
+        W_z = torch.zeros_like(W)
+        mask = torch.ones(self.num_classes,dtype=torch.bool, device = device)
+        mask[forget_index] = False
+        W_z[mask] = W_temp
+
 
         # 12
         clf = self._get_classifier(model)
+        # load the weights
         with torch.no_grad():
-            clf.weight.copy_(W_Z)
+            clf.weight.copy_(W_z)
+            # neutralise forget bias if exists
+            if clf.bias is not None:
+                n_bias = clf.bias.data.clone()
+                n_bias[forget_index] = -1e9
+                clf.bias.copy_(n_bias)
+
 
         return model