wf_net

Tune_new.py

@@ -13,6 +13,8 @@ from unlearning.CertifiedRemoval import CertifiedRemoval
 from unlearning.CertifiedUnlearning import CertifiedUnlearning
 from unlearning.LinearFiltration import LinearFiltration
 from unlearning.WeightFiltration import WeightFiltration
+from unlearning.WF import WeightF
+
 
 # Global Hyperparameters
 CLASS_SIZE = 20
@@ -255,16 +257,16 @@ if __name__ == "__main__":
             target_class_index=i
         )
 
-        weight_filtration = WeightFiltration(
+        weight_filtration = WeightF( #WeightFiltration(
             target_class_index=i,
             epochs=3,
-            lr=0.5,
-            gamma=150
+            lr=0.05,
+            gamma=5
         )
 
         strategies = [
-            certified_unlearning,
-            # weight_filtration,
+            # certified_unlearning,
+            weight_filtration,
             # linear_filtration
         ]
 

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

unlearning/wf/WF_Net.py

@@ -0,0 +1,86 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class WF_Net(nn.Module):
+    """
+    Implements Poppi et al.'s WF Model structure.
+    Wraps a pre-trained ResNet-18 and dynamically applies 
+    weight-space gating matrix multiplication during the forward step.
+    """
+    def __init__(self, original_model: nn.Module, num_classes: int):
+        super().__init__()
+        # Extract the sequence of blocks/layers L from the original model
+        self.conv1 = original_model.conv1
+        self.bn1 = original_model.bn1
+        self.relu = original_model.relu
+        self.maxpool = original_model.maxpool
+        
+        self.layer1 = original_model.layer1
+        self.layer2 = original_model.layer2
+        self.layer3 = original_model.layer3
+        self.layer4 = original_model.layer4
+        
+        self.avgpool = original_model.avgpool
+        self.fc = original_model.fc
+        
+        # Target layer for filtering: layer4 block 1 conv2
+        # We extract its static tensor data out of the autograd parameter pool
+        self.target_conv = self.layer4[1].conv2
+        self.original_w = nn.Parameter(self.target_conv.weight.data.clone().detach(), requires_grad=False)
+        
+        # Require: Alpha gating matrix. Shape: (num_classes, out_channels)
+        # Initialized to 1.5 as per Poppi et al.'s verbatim specification
+        out_channels = self.original_w.shape[0]
+        #self.alpha = nn.Parameter(torch.ones(num_classes, out_channels) * 1.5)
+        self.alpha = nn.Parameter(torch.ones(num_classes, out_channels))
+
+    def forward(self, x: torch.Tensor, target_unlearn_class: int) -> torch.Tensor:
+        """
+        Implements Algorithm 1: General forward step of a WF model
+        Inputs: 
+          x: Input tensor (Xin)
+          target_unlearn_class: The class index we are actively filtering out (Yunl)
+        """
+        # 1. Run through early sequence of layers undisturbed
+        x = self.maxpool(self.relu(self.bn1(self.conv1(x))))
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        
+        # Run layer4 block 0 and block 1 conv1 normally
+        x = self.layer4[0](x)
+
+        identity = x
+
+
+
+        x = self.layer4[1].conv1(x)
+        x = self.layer4[1].bn1(x)
+        x = self.layer4[1].relu(x)
+        
+        # 2. CORE WF-NET MATH: w_hat_l <- alpha_l[Yunl] ⊙ w_l
+        # Extract 1D vector for target class and reshape to (out_channels, 1, 1, 1) for 4D convolution broadcasting
+        mask = torch.sigmoid(self.alpha[target_unlearn_class]).view(-1, 1, 1, 1)
+        w_hat = self.original_w * mask
+        
+        # 3. Pass gated weights straight to functional forward pass: l(Xi, w_hat_l)
+        x = F.conv2d(
+            x, 
+            weight=w_hat, 
+            bias=self.target_conv.bias, 
+            stride=self.target_conv.stride, 
+            padding=self.target_conv.padding
+        )
+        x = self.layer4[1].bn2(x)
+        
+        # Handle residual shortcut skip connection manually since we opened up block 1
+        # In ResNet-18 layer4, block 1 has no downsample shortcut layer; it's a direct identity add
+        x = self.layer4[1].relu(x + identity) 
+        
+        # 4. Final Classification Head Sequence
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        y_out = self.fc(x)
+        
+        return y_out