wf_net

unlearning/WeightFiltration.py

@@ -59,20 +59,21 @@ class WeightFiltration(Strategy):
                 
                 temperature = 3.0
                 logits_f_scaled = outputs_f / temperature
-                loss_f = -torch.sum(
-                    (torch.ones_like(logits_f_scaled) / num_classes) * torch.log_softmax(logits_f_scaled, dim=-1)
-                )
                 
+                # Compute uniform target entropy per-sample, then average over the batch
+                log_probs_f = torch.log_softmax(logits_f_scaled, dim=-1)
+                uniform_target = torch.ones_like(logits_f_scaled) / num_classes
+                loss_f = -torch.sum(uniform_target * log_probs_f, dim=-1).mean()
+                                
                 total_loss = loss_r + (self.gamma * loss_f)
                 total_loss.backward()
                 optimizer.step()
-                
+                                
                 t_loss_r += loss_r.item()
                 t_loss_f += loss_f.item()
                 steps += 1
-                
+                            
             print(f"  Epoch {epoch+1}/{self.epochs} | Retain Loss: {t_loss_r/steps:.4f} | Forget Loss: {t_loss_f/steps:.4f}")
-
         return wf_model
             
 
@@ -110,15 +111,16 @@ class WeightFiltration(Strategy):
         )
 
         # --- PERMANENT BAKING STEP ---
-        # Disconnect the dynamic parameter and freeze the optimal gated state permanently into the architecture
-    
         with torch.no_grad():
+            # Grab the alpha mask vector for the forgotten class and cast to 4D tensor shape
             final_mask = torch.sigmoid(wf_model.alpha[self.target_class_index]).view(-1, 1, 1, 1)
-            target_conv.weight.copy_(original_weights * final_mask)
             
-            # Re-enable model parameters for downstream evaluation processing
+            # Apply filter masking permanently back onto the base layer
+            target_conv.weight.copy_(original_weights * final_mask)
+                        
+            # Unfreeze architecture parameters for evaluations downstream
             for p in model.parameters():
                 p.requires_grad = True
-                
-        print(f">> Permanently altered {out_channels} convolutional filters in layer4 via WF-Net.")
-        return model
+                        
+            print(f">> Permanently altered {out_channels} convolutional filters in layer4 via WF-Net.")
+        return model