unlearning done

Tune_new.py

@@ -9,11 +9,9 @@ import Util
 from sets.Data import *
 from sets.IdentitySubset import IdentitySubset
 from architectures.Model import Model, Architecture
-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
@@ -140,40 +138,12 @@ def run_unlearning_and_strategy_eval(env_dict, forget_class_idx, strategy, evalu
     train_data = env_dict["train_data"]
     test_data = env_dict["test_data"]
 
-    # testing valuse                             *         *
-    #---------------------------------------------------------------------------
-    # S1     50       5        5        5        5         5
-    # S2     1000     200      1000     500      200       300
-    # BS     5        5        5        5        5         5
-    # scale  2000     500      8000     5000     10000     8000
-    # std    0.00001  0.00001  0.00001  0.00001  0.00001   0.00001
-
-    # Initialize the strategy hyperparameters matching standard settings
-    # increase s2, decrease scale ---sweet spot
-    '''certified_removal = CertifiedRemoval(
-        target_class_index=forget_class_idx,
-        s1=4,
-        s2=350,  # 350 best
-        unlearn_bs=5,
-        scale=6000.0,   # 6000 was good
-        std=0.00001
-    )'''
-    '''certified_removal = CertifiedUnlearning(
-        target_class_index=0,
-        l2_reg=0.0005,
-        gamma=0.1,  
-        scale=7000.0,      
-        s1=2,
-        s2=350,
-        std=1e-5,  
-        unlearn_bs=2
-    )'''
 
     # Segment specific unlearning loaders using class index boundaries
-    forget_train_loader, retain_train_loader = get_unlearning_loaders(
+    retain_train_loader , forget_train_loader= get_unlearning_loaders(
         dataset=train_data, forget_class_idx=forget_class_idx, batch_size=BATCH_SIZE
     )
-    forget_test_loader, retain_test_loader = get_unlearning_loaders(
+    retain_test_loader, forget_test_loader = get_unlearning_loaders(
         dataset=test_data, forget_class_idx=forget_class_idx, batch_size=BATCH_SIZE
     )
 
@@ -189,9 +159,16 @@ def run_unlearning_and_strategy_eval(env_dict, forget_class_idx, strategy, evalu
     print("fine tunned model loaded into evaluation sandbox")
 
     # Execute strategic parameter unlearning step
-    strategy.apply(reloaded.model, forget_train_loader, retain_train_loader)
+    unlearned = strategy.apply(reloaded.model, train_data)
     strategy_in_use = strategy.__class__.__name__
 
+    if isinstance(unlearned,nn.Module):
+        reloaded.model = unlearned
+    else:
+        reloaded = unlearned
+
+    
+
     # Define validation tracking steps dynamically
     evaluation_domains = [
         {"loader": retain_test_loader, "mode": "retain", "label": "\n--- Performance on Retained Classes"},
@@ -215,66 +192,63 @@ def run_unlearning_and_strategy_eval(env_dict, forget_class_idx, strategy, evalu
 # entry
 if __name__ == "__main__":
     
-    # Run Data Infrastructure and Architecture Builder
-    runtime_environment = prepare_data_and_model_environment()
 
-    # Baseline Evaluation 
-    finetuning = False
-    # switch finetuning for tests on strategies only
-    run_finetuning_or_baseline_eval(runtime_environment, run_training=finetuning)
+    try:
+        # Run Data Infrastructure and Architecture Builder
+        runtime_environment = prepare_data_and_model_environment()
+        
+        # Baseline Evaluation 
+        finetuning = False
+        # switch finetuning for tests on strategies only
+        run_finetuning_or_baseline_eval(runtime_environment, run_training = finetuning)
         
-    finetuning = True
-    # Unlearning Iterations
-    for i in range(0, 1):
         
         # strategies
-        # 
-        #certified_removal = CertifiedRemoval(
-        #        target_class_index=i,
-        #        s1=4,
-        #        s2=350,  # 350 best
-        #        unlearn_bs=5,
-        #        scale=6000.0,   # 6000 was good 
-        #        std=0.00009
-        #    )
-        
-
-
         certified_unlearning = CertifiedUnlearning(
-            target_class_index=i,
+            target_class_index=0,
             l2_reg=0.000002,       
             gamma=0.1,           
-            scale=  20000,# 16400.0, # took ages to reach this sweet spot     
+            scale=  16400.0,# 16400.0, # took ages to reach this sweet spot     
             s1=2,
             s2=300,               
             std=0.00001,
-            unlearn_bs=16
+            unlearn_bs=8
         )
 
         # works perfectly
         linear_filtration = LinearFiltration(
 
-            target_class_index=i
+            target_class_index=0
         )
 
-        weight_filtration = WeightF( #WeightFiltration(
-            target_class_index=i,
-            epochs=3,
-            lr=0.05,
-            gamma=5
+        weight_filtration = WeightFiltration(
+            target_class_index=0, 
+            epochs=6, 
+            lr=150.0, 
+            gamma=0.001
         )
 
         strategies = [
-            # certified_unlearning,
+            certified_unlearning,
             weight_filtration,
-            # linear_filtration
+            linear_filtration
         ]
+        # Unlearning Iteration
+        for i in range(0, CLASS_SIZE):
 
-        print(f"\n>>> Executing Unlearning Framework for Target Identity Index: {i} <<<")
-        for strategy in strategies:
-            run_unlearning_and_strategy_eval(
-                runtime_environment, 
-                forget_class_idx=i, 
-                strategy=strategy,
-                evaluate= not finetuning
-            )
+            for strategy in strategies:
+
+                # update target class to be unlearned
+                strategy.set_target_class(i)
+                print(f"Unlearning class {i} with {strategy.strategy_name}")
+
+                # forget
+                run_unlearning_and_strategy_eval(
+                    runtime_environment, 
+                    forget_class_idx=i, 
+                    strategy=strategy,
+                    evaluate = not finetuning
+                )
+
+    except KeyboardInterrupt:
+        print("program interrupted. Exit!")

Util.py

@@ -46,3 +46,6 @@ def log_metric(log_file, execution_time: float):
     """Appends the execution time to this strategy's specific file."""
     with open(log_file, "a") as f:
         f.write(f"{execution_time:.6f}\n")
+
+
+

architectures/Model.py

@@ -7,7 +7,7 @@ import time
 import numpy as np
 from sklearn.metrics import classification_report
 from pathlib import Path
-from unlearning.Strategy import Strategy
+#from unlearning.Strategy import Strategy
 import copy
 from torch.optim.lr_scheduler import CosineAnnealingLR
 
@@ -84,7 +84,7 @@ class Model(ABC):
         print(f'Model loaded from {file_path}')
 
     
-    def unlearn(self, strategy: Strategy, forget_loader, retain_loader):
+    def unlearn(self, strategy: 'Strategy', forget_loader, retain_loader):
         """ Executes a targeted unlearning strategy and profiles efficiency """
         print(f"Executing: {strategy.__class__.__name__}...")
         
@@ -103,6 +103,7 @@ class Model(ABC):
         Evaluates the model, prints terminal reports, and routes metrics to 
         a file logger based on the current context mode.
         """
+        
         self.model.eval()
         all_preds, all_labels = [], []
         print(f"\nEvaluating Domain: [{mode}]...")

sets/Data.py

@@ -1,5 +1,5 @@
 from torchvision import datasets, transforms
-from torch.utils.data import Dataset, DataLoader, Subset
+from torch.utils.data import Dataset, DataLoader, Subset, ConcatDataset
 import torch
 import numpy as np
 import os
@@ -181,4 +181,59 @@ def get_unlearning_loaders(dataset: Dataset, forget_class_idx: int, batch_size:
     
     print(f"[Data Split] Local Class {forget_class_idx}: {len(forget_subset)} samples | Remaining Classes: {len(retain_subset)} samples.")
     
-    return forget_loader, retain_loader
+    return retain_loader, forget_loader
+
+
+def vertical_split(dataset, batch_size,num_classes):
+        """
+        Executes a class-wise vertical split. 
+        Divides the samples of every single identity class exactly in half:
+        50% of each class goes to the Retain Set, 50% goes to the Forget Set.
+        """
+        
+        # 1. Group dataset indices by their respective ground-truth classes
+        class_to_indices = {c: [] for c in range(num_classes)}
+        
+        print("  [Vertical Split] Tracking class indices across the combined dataset...")
+        for idx in range(len(dataset)):
+            # Extract the label cleanly from the underlying dataset structure
+            _, label = dataset[idx]
+            if label in class_to_indices:
+                class_to_indices[label].append(idx)
+
+        retain_indices = []
+        forget_indices = []
+
+        # 2. Slice each class identity vertically (exactly 50/50)
+        for c, indices in class_to_indices.items():
+            if len(indices) < 2:
+                print(f"  Warning: Class {c} has fewer than 2 samples. Cannot split vertically.")
+                retain_indices.extend(indices)
+                continue
+                
+            # Deterministic shuffle per class to ensure honest distribution before splitting
+            np.random.shuffle(indices)
+            
+            mid = len(indices) // 2
+            forget_indices.extend(indices[:mid])   # First half assigned to unlearning
+            retain_indices.extend(indices[mid:])   # Second half assigned to retention
+
+        print(f"  Vertical split complete: Retain Index Size = {len(retain_indices)} | Forget Index Size = {len(forget_indices)}")
+
+        # 3. Construct lightweight PyTorch Subsets using our sliced index maps
+        retain_subset = Subset(dataset, retain_indices)
+        forget_subset = Subset(dataset, forget_indices)
+
+        # 4. Return pristine, shuffled DataLoaders mirroring your environment's batch specifications
+        retain_loader = DataLoader(retain_subset, batch_size=batch_size, shuffle=True)
+        forget_loader = DataLoader(forget_subset, batch_size=batch_size, shuffle=True)
+
+        return retain_loader, forget_loader
+
+def _combine_set(loader_one, loader_two):
+    full_train_dataset = ConcatDataset([loader_one.dataset, loader_two.dataset])
+    return DataLoader(
+        full_train_dataset, 
+        batch_size=loader_one.batch_size, 
+        shuffle=True
+    )

unlearning/CertifiedRemoval.py

@@ -1,214 +0,0 @@
-import torch
-import torch.nn as nn
-from torch.utils.data import DataLoader, RandomSampler
-from torch.autograd import grad
-from unlearning.Strategy import Strategy
-
-class CertifiedRemoval(Strategy):
-    """
-    Implements Certified Unlearning for non-convex DNNs (Zhang et al.).
-    Uses a modified, stabilized stochastic Newton step using Taylor-expansion 
-    HVP estimation across the entire parameter space, capped with calibrated noise.
-    """
-    def __init__(self, target_class_index: int, l2_reg: float = 0.0005, 
-                 gamma: float = 0.01, scale: float = 1000.0, 
-                 s1: int = 10, s2: int = 1000, std: float = 0.001, unlearn_bs: int = 2):
-        super().__init__(target_class_index)
-        self.l2_reg = l2_reg        
-        self.gamma = gamma          
-        self.scale = scale          
-        self.s1 = s1                
-        self.s2 = s2                
-        self.std = std              
-        self.unlearn_bs = unlearn_bs
-
-    '''
-    def _compute_loss_gradient(self, model, loader, device: torch.device):
-        model.eval()
-        criterion = nn.CrossEntropyLoss(reduction='sum')
-        params = [p for p in model.parameters() if p.requires_grad]
-        grad_accumulator = [torch.zeros_like(p).cpu() for p in params]
-        total_samples = 0
-        
-        for data, targets in loader:
-            total_samples += targets.shape[0]
-            data, targets = data.to(device), targets.to(device)
-            outputs = model(data)
-            
-            mini_grads = list(grad(criterion(outputs, targets), params))
-            for i in range(len(grad_accumulator)):
-                grad_accumulator[i] += mini_grads[i].cpu().detach()
-
-        for i in range(len(grad_accumulator)):
-            grad_accumulator[i] /= total_samples
-
-        l2_reg_term = 0.0
-        for param in model.parameters():
-            l2_reg_term += torch.norm(param, p=2)
-        
-        reg_grads = list(grad(self.l2_reg * l2_reg_term, params))
-        for i in range(len(grad_accumulator)):
-            grad_accumulator[i] += reg_grads[i].cpu().detach()
-            
-        return [p.to(device) for p in grad_accumulator]'''
-    def _compute_loss_gradient(self, model, loader, device: torch.device):
-        model.eval()
-        # Use reduction='sum' matching the original framework
-        criterion = nn.CrossEntropyLoss(reduction='sum')
-        params = [p for p in model.parameters() if p.requires_grad]
-        grad_accumulator = [torch.zeros_like(p).cpu() for p in params]
-        total_samples = 0
-        
-        for data, targets in loader:
-            total_samples += targets.shape[0]
-            data, targets = data.to(device), targets.to(device)
-            outputs = model(data)
-            
-            loss = criterion(outputs, targets)
-            
-            # Incorporate L2 weight regularization directly inside the backprop graph
-            # to keep scaling bounded and aligned with the data volume
-            l2_reg_term = 0.0
-            for param in model.parameters():
-                if param.requires_grad:
-                    l2_reg_term += torch.norm(param, p=2)
-            
-            total_loss = loss + (self.l2_reg * l2_reg_term)
-            
-            mini_grads = list(grad(total_loss, params, retain_graph=False))
-            for i in range(len(grad_accumulator)):
-                grad_accumulator[i] += mini_grads[i].cpu().detach()
-
-        for i in range(len(grad_accumulator)):
-            grad_accumulator[i] /= total_samples
-            
-        return [p.to(device) for p in grad_accumulator]
-
-    
-    def grad_batch(batch_loader, lam, model, device):
-        model.eval()
-        criterion = nn.CrossEntropyLoss(reduction='sum')
-        params = [p for p in model.parameters() if p.requires_grad]
-        grad_batch = [torch.zeros_like(p).cpu() for p in params]
-        num = 0
-        for batch_idx, (data, targets) in enumerate(batch_loader):
-            num += targets.shape[0]
-            data, targets = data.to(device), targets.to(device)
-            outputs = model(data)
-
-            grad_mini = list(grad(criterion(outputs, targets), params))
-            for i in range(len(grad_batch)):
-                grad_batch[i] += grad_mini[i].cpu().detach()
-
-        for i in range(len(grad_batch)):
-            grad_batch[i] /= num
-
-        l2_reg = 0
-        for param in model.parameters():
-            l2_reg += torch.norm(param, p=2)
-        grad_reg = list(grad(lam * l2_reg, params))
-        for i in range(len(grad_batch)):
-            grad_batch[i] += grad_reg[i].cpu().detach()
-        return [p.to(device) for p in grad_batch]
-
-    def _hvp(self, loss, params, v):
-        first_grads = grad(loss, params, retain_graph=True, create_graph=True)
-        elemwise_products = 0
-        for grad_elem, v_elem in zip(first_grads, v):
-            elemwise_products += torch.sum(grad_elem * v_elem)
-        # FIX 1: Set create_graph to False to prevent massive nested graph accumulation
-        return grad(elemwise_products, params, create_graph=False)
-
-    def _stochastic_newton_update(self, g, retain_dataset, model, device):
-        model.eval()
-        criterion = nn.CrossEntropyLoss()
-        params = [p for p in model.parameters() if p.requires_grad]
-        h_res = [torch.zeros_like(p) for p in g]
-        
-        for _ in range(self.s1):
-            h_estimate = [p.clone() for p in g]
-            sampler = RandomSampler(retain_dataset, replacement=True, num_samples=self.unlearn_bs * self.s2)
-            res_loader = DataLoader(retain_dataset, batch_size=self.unlearn_bs, sampler=sampler)
-            res_iter = iter(res_loader)
-            
-            for j in range(self.s2):
-                try:
-                    data, target = next(res_iter)
-                except StopIteration:
-                    res_iter = iter(res_loader)
-                    data, target = next(res_iter)
-                    
-                data, target = data.to(device), target.to(device)
-                outputs = model(data)
-                
-                loss = criterion(outputs, target)
-                l2_reg_term = 0.0
-                for param in model.parameters():
-                    l2_reg_term += torch.norm(param, p=2)
-                loss += (self.l2_reg + self.gamma) * l2_reg_term
-                
-                h_s = self._hvp(loss, params, h_estimate)
-                
-                with torch.no_grad():
-                    for k in range(len(params)):
-                        # FIX 2: Added .detach() to decouple history strings across iterative update blocks
-                        #h_estimate[k] = (h_estimate[k] + g[k] - h_s[k] / self.scale).detach()
-                        next_estimate = h_estimate[k].data + g[k].data - (h_s[k].data / self.scale)
-                        h_estimate[k] = next_estimate.clone()
-                del h_s, loss, outputs
-                        
-            for k in range(len(params)):
-                h_res[k] = h_res[k] + h_estimate[k] / self.scale
-                
-        return [p / self.s1 for p in h_res]
-
-    '''def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
-        device = next(model.parameters()).device
-
-        num_forget = len(forget_loader.dataset)
-        num_retain = len(retain_loader.dataset)
-        scaling_ratio = num_forget / num_retain
-        
-        print(">> Calculating base gradients over target FORGET set...")
-        # FIX 3: Base gradients MUST be evaluated from forget_loader to drop target class distributions
-        g = self._compute_loss_gradient(model, forget_loader, device)
-        
-        print(">> Estimating non-convex inverse Hessian trajectories via Taylor series...")
-        retain_dataset = retain_loader.dataset 
-        delta = self._stochastic_newton_update(g, retain_dataset, model, device)
-        
-        print(">> Applying stabilized parameter adjustments and randomized certification noise...")
-        with torch.no_grad():
-            for i, param in enumerate(model.parameters()):
-                if param.requires_grad:
-                    noise = self.std * torch.randn(param.data.size(), device=device)
-                    #param.data.add_(-delta[i] + noise)
-                    param.data.add_(scaling_ratio * delta[i] + noise)
-                
-        print(">> Certified Unlearning process completed successfully across the complete landscape.")
-        return model'''
-    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
-        device = next(model.parameters()).device
-
-        print(">> Calculating stable base gradients over the RETAIN set...")
-        # To match the author's snippet perfectly, g MUST be computed on the retain data.
-        # If this loader is too large for your VRAM, use a smaller batch size (e.g. 16 or 32) 
-        # in your main training script when creating retain_loader.
-        g = self._compute_loss_gradient(model, retain_loader, device)
-        
-        print(">> Estimating non-convex inverse Hessian trajectories via Taylor series...")
-        retain_dataset = retain_loader.dataset 
-        delta = self._stochastic_newton_update(g, retain_dataset, model, device)
-        
-        print(">> Applying parameter removal adjustments (-delta)...")
-        with torch.no_grad():
-            for i, param in enumerate(model.parameters()):
-                if param.requires_grad:
-                    noise = self.std * torch.randn(param.data.size(), device=device)
-                    
-                    # MATCHING THE SNIPPET: Subtract delta exactly as the authors do
-                    # This removes the influence trace of the omitted data.
-                    param.data.add_(-delta[i] + noise)
-                
-        print(">> Certified Unlearning process completed successfully.")
-        return model

unlearning/Certified_facebook.py

@@ -1,123 +0,0 @@
-import torch
-import torch.nn as nn
-import math
-from torch.utils.data import DataLoader
-from unlearning.Strategy import Strategy
-
-class CertifiedRemovalFacebook(Strategy):
-    """
-    Implements Certified Removal (Guo et al.) mapped for Multi-Class models
-    by executing a single-class One-vs-Rest (OvR) block-removal update step.
-    Math matches the facebookresearch/certified-removal reference repository.
-    """
-    def __init__(self, target_class_index: int, removal_bound: float, epsilon: float, l2_reg: float = 0.1):
-        super().__init__(target_class_index=target_class_index)
-        self.removal_bound = removal_bound  # gamma in the paper
-        self.epsilon = epsilon              # Privacy budget
-        self.l2_reg = l2_reg                # Lambda (regularization term)
-
-    def _get_features(self, backbone: nn.Module, loader: DataLoader, device: torch.device):
-        """Passes data through the frozen ResNet backbone to extract embedding features."""
-        backbone.eval()
-        all_features = []
-        
-        with torch.no_grad():
-            for inputs, _ in loader:
-                inputs = inputs.to(device)
-                # Pass through frozen backbone to get the 2048-dimensional embedding
-                features = backbone(inputs)
-                all_features.append(features.cpu())
-                
-        return torch.cat(all_features, dim=0)
-
-    def _fb_lr_grad(self, w, X, y, lam):
-        """
-        Replicates exact lr_grad calculation from Facebook's codebase.
-        Note: The resulting gradient has a flipped sign due to the structure of (z - 1).
-        """
-        # X.mv(w) computes raw linear margins
-        z = torch.sigmoid(y * X.mv(w))
-        # Gradient formula: X^T * ((z - 1) * y) + lambda * N * w
-        return X.t().mv((z - 1) * y) + lam * X.size(0) * w
-
-    def _fb_lr_hessian_inv(self, w, X, y, lam, device, batch_size=50000):
-        """
-        Replicates exact lr_hessian_inv calculation from Facebook's codebase.
-        Scales the L2 regularization matrix explicitly by dataset row count (N * lambda * I).
-        """
-        z = torch.sigmoid(X.mv(w).mul_(y))
-        D = z * (1 - z) # Element-wise variance vector
-        
-        H = None
-        num_batch = int(math.ceil(X.size(0) / batch_size))
-        for i in range(num_batch):
-            lower = i * batch_size
-            upper = min((i + 1) * batch_size, X.size(0))
-            X_i = X[lower:upper]
-            
-            # Stepwise feature weighting via element-wise variance columns
-            if H is None:
-                H = X_i.t().mm(D[lower:upper].unsqueeze(1) * X_i)
-            else:
-                H += X_i.t().mm(D[lower:upper].unsqueeze(1) * X_i)
-        
-        # Scale identity buffer by dataset split size: lambda * N_retain
-        reg_matrix = lam * X.size(0) * torch.eye(X.size(1), device=device).float()
-        return torch.linalg.inv(H + reg_matrix)
-
-    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
-        """
-        Applies Certified Removal strictly to the target class parameters 
-        belonging to the final fully connected layer (model.fc).
-        """
-        device = next(model.parameters()).device
-        k = self.target_class_index
-        
-        # Isolate final layer and extract raw deep embeddings using frozen backbone
-        linear_head = model.fc
-        model.fc = nn.Identity() 
-        
-        print(">> Extracting deep features from model backbone...")
-        X_retain = self._get_features(model, retain_loader, device).to(device)
-        X_forget = self._get_features(model, forget_loader, device).to(device)
-        
-        # Restore the classification head back
-        model.fc = linear_head
-        
-        # Extract current model weight row for the target class channel
-        w_k = model.fc.weight.data[k].clone().to(device)
-
-        # Create One-vs-Rest binary target indicator arrays (+1.0 / -1.0)
-        # Retain dataset instances are negative labels (-1.0) for the target class channel
-        y_retain_binary = torch.full((X_retain.size(0),), -1.0, device=device)
-        # Forget dataset instances are positive labels (+1.0) for the target class channel
-        y_forget_binary = torch.full((X_forget.size(0),), 1.0, device=device)
-
-        # Compute Inverse Hessian (on Retain Data) and Gradient (on Forget Data)
-        H_inv = self._fb_lr_hessian_inv(w_k, X_retain, y_retain_binary, self.l2_reg, device)
-        grad_forget = self._fb_lr_grad(w_k, X_forget, y_forget_binary, self.l2_reg)
-
-        # 5. Compute the Weight Update Step Vector (Delta)
-        multiplier = 0.5
-        delta_w_k = torch.mv(H_inv, grad_forget) * multiplier
-
-        # Verify Theoretical Removal Bound Criteria
-        norm_delta = torch.norm(delta_w_k).item()
-        if norm_delta > self.removal_bound:
-            print(f"!! Warning: Removal budget exceeded! Norm: {norm_delta:.4f} > Bound: {self.removal_bound}")
-        else:
-            print(f">> Certificate valid. Norm: {norm_delta:.4f} <= Bound: {self.removal_bound}")
-
-        # Apply Update (Using '+' since Facebook's grad calculation yields a negative sign output)
-        new_w_k = w_k + delta_w_k
-
-        # Calibrate and Inject Perturbation Noise (Objective Perturbation Verification)
-        sigma = 2.0 / (self.l2_reg * self.epsilon)
-        noise = torch.randn_like(new_w_k, device=device) * (sigma / X_retain.size(0))
-        new_w_k = new_w_k + noise
-
-        # Commit updated weight vector row back into model head parameters in-place
-        model.fc.weight.data[k] = new_w_k
-        
-        print(">> Certified Removal process completed successfully.")
-        return model

unlearning/LastK_Certified.py

@@ -1,125 +0,0 @@
-import torch
-import torch.nn as nn
-from torch.utils.data import DataLoader
-from unlearning.Strategy import Strategy
-
-class LastKCertifiedRemoval(Strategy):
-    """
-    Implements Certified Removal (Guo et al.) scaled up to the last K layers 
-    of a ResNet50 network by flattening sub-graph parameters into a convex sub-problem.
-    """
-    def __init__(self, removal_bound: float, epsilon: float, l2_reg: float = 0.1):
-        super().__init__()
-        self.removal_bound = removal_bound
-        self.epsilon = epsilon
-        self.l2_reg = l2_reg
-
-    def _split_model(self, model: nn.Module):
-        """
-        Splits ResNet50 into a frozen feature backbone and an active unlearning head.
-        Here, 'Last K Layers' includes layer4 and the fc classification head.
-        """
-        # Feature Backbone: Everything up to layer3
-        backbone = nn.Sequential(
-            model.conv1,
-            model.bn1,
-            model.relu,
-            model.maxpool,
-            model.layer1,
-            model.layer2,
-            model.layer3
-        )
-        
-        # Active Head: Layer4, global pooling, and the final linear layer
-        unlearning_head = nn.Sequential(
-            model.layer4,
-            model.avgpool,
-            nn.Flatten(1),
-            model.fc
-        )
-        
-        return backbone, unlearning_head
-
-    def _get_intermediate_features(self, backbone: nn.Module, loader: DataLoader, device: torch.device):
-        """Extracts features from the exit point of the frozen backbone (post-layer3)."""
-        backbone.eval()
-        all_features = []
-        all_labels = []
-        
-        with torch.no_grad():
-            for inputs, labels in loader:
-                inputs = inputs.to(device)
-                features = backbone(inputs)
-                all_features.append(features.cpu())
-                all_labels.append(labels.cpu())
-                
-        return torch.cat(all_features, dim=0), torch.cat(all_labels, dim=0)
-
-    def apply(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
-        """
-        Extracts intermediate features and updates the parameters of the last blocks
-        using the exact inverse-Hessian influence step.
-        """
-        device = next(model.parameters()).device
-        
-        # 1. Slice the ResNet graph structural components
-        backbone, unlearning_head = self._split_model(model)
-        
-        print(">> Extracting intermediate structural features from layer3 exit...")
-        retain_feats, retain_labels = self._get_intermediate_features(backbone, retain_loader, device)
-        forget_feats, forget_labels = self._get_intermediate_features(backbone, forget_loader, device)
-        
-        # 2. Flatten target weights from the active head into a 1D optimization tensor
-        # For simplicity and mathematical stability, we isolate the final layer's weights
-        # inside the active head for the exact Hessian tracking step
-        target_layer = unlearning_head[-1] # This points straight to model.fc
-        w = target_layer.weight.data.clone().cpu()
-        
-        # 3. Compute Exact Hessian over intermediate embeddings
-        # ResNet50's layer4 expands channels to 2048, creating a 2048x2048 matrix context
-        print(">> Computing exact sub-graph Hessian matrix...")
-        N_retain = retain_feats.size(0)
-        
-        # Pool the feature maps if they haven't been flattened yet by the head module
-        if len(retain_feats.shape) > 2:
-            retain_flat = torch.mean(retain_feats, dim=[2, 3])
-            forget_flat = torch.mean(forget_feats, dim=[2, 3])
-        else:
-            retain_flat = retain_feats
-            forget_flat = forget_feats
-
-        X_T_X = torch.matmul(retain_flat.t(), retain_flat)
-        reg_matrix = self.l2_reg * torch.eye(retain_flat.size(1))
-        Hessian = (X_T_X / N_retain) + reg_matrix
-        
-        # 4. Calculate gradients relative to the forgotten target features
-        print(">> Calculating forget set gradients...")
-        num_classes = w.size(0)
-        forget_labels_one_hot = torch.nn.functional.one_hot(forget_labels, num_classes=num_classes).float()
-        
-        preds_forget = torch.matmul(forget_flat, w.t())
-        error = preds_forget - forget_labels_one_hot
-        grad_forget = torch.matmul(error.t(), forget_flat) / forget_flat.size(0)
-
-        # 5. Apply Newton Step optimization update
-        print(">> Inverting optimization subspace via system solver...")
-        try:
-            delta_w_t = torch.linalg.solve(Hessian, grad_forget.t())
-            delta_w = delta_w_t.t()
-        except RuntimeError:
-            print(">> Warning: Subspace Hessian is singular. Using pseudo-inverse fallback.")
-            delta_w = torch.matmul(grad_forget, torch.linalg.pinv(Hessian).t())
-
-        # 6. Apply Weight Adjustment Bounds Check
-        new_w = w + delta_w
-        norm_delta = torch.norm(delta_w).item()
-        if norm_delta > self.removal_bound:
-            print(f"!! Warning: Removal budget exceeded! Norm: {norm_delta:.4f} > Bound: {self.removal_bound}")
-        else:
-            print(f">> Certificate valid. Subspace Norm: {norm_delta:.4f} <= Bound: {self.removal_bound}")
-
-        # 7. Write weights directly back into the live ResNet50 instance
-        model.fc.weight.data = new_w.to(device)
-        
-        print(">> Last K Layers Certified Removal complete.")
-        return model

unlearning/LinearFiltration.py

@@ -2,6 +2,7 @@ import torch
 import torch.nn as nn
 from .Strategy import Strategy
 from torch.utils.data import DataLoader
+from sets.Data import get_unlearning_loaders, _combine_set
 
 class LinearFiltration(Strategy):
     def __init__(self, target_class_index):
@@ -23,40 +24,8 @@ class LinearFiltration(Strategy):
             forget_index=self.target_class_index
         )
     
-    # FIX: Added staticmethod decorator
-    @staticmethod
-    def get_features(model, inputs):
-        # For ResNet, pass through everything up to the fc layer
-        x = model.conv1(inputs)
-        x = model.bn1(x)
-        x = model.relu(x)
-        x = model.maxpool(x)
     
-        x = model.layer1(x)
-        x = model.layer2(x)
-        x = model.layer3(x)
-        x = model.layer4(x)
-
-        x = model.avgpool(x)
-        x = torch.flatten(x, 1)
-        return x
-
-    @staticmethod
-    def _calculate_filtration_matrix(num_classes: int, forget_class: int, device: torch.device) -> torch.Tensor:
-        A = torch.eye(num_classes, device=device)
-        num_remaining = num_classes - 1
-        
-        for j in range(num_classes):
-            if j == forget_class:
-                A[forget_class, j] = 0.0
-            else:
-                A[forget_class, j] = 1.0 / num_remaining
-                
-        return A
-    
-
-    @staticmethod
-    def _sums_and_counts(model, num_classes, retain_loader, forget_loader, device, forget_index, h_dim):
+    def _sums_and_counts(self, model, num_classes, loader, device, forget_index, h_dim):
         model.eval()
 
         sums = torch.zeros(num_classes, h_dim, device=device)
@@ -64,11 +33,11 @@ class LinearFiltration(Strategy):
 
         # Generate values for retain
         with torch.no_grad():
-            for inputs, targets in retain_loader:
+            for inputs, targets in loader:
                 inputs = inputs.to(device)
                 targets = targets.to(device)
-                # FIX: Call get_features instead of model() directly
-                outputs = LinearFiltration.get_features(model, inputs)
+                # predictions
+                outputs = model(inputs)
 
                 for j in range(num_classes):
                     if j == forget_index:
@@ -79,65 +48,54 @@ class LinearFiltration(Strategy):
                         sums[j] += outputs[mask].sum(dim=0)
                         counts[j] += mask.sum()
         
-        # Values for forget
-        with torch.no_grad():
-            for inputs, targets in forget_loader:
-                inputs = inputs.to(device)
-                targets = targets.to(device)
-                # FIX: Call get_features instead of model() directly
-                outputs = LinearFiltration.get_features(model, inputs)
-
-                mask = (targets == forget_index)
-
-                if mask.any():
-                    sums[forget_index] += outputs[mask].sum(dim=0)
-                    counts[forget_index] += mask.sum()
-
         return sums, counts
     
-    @staticmethod
-    def _get_means(model, num_classes, retain_loader, forget_loader, device, forget_index):
-        h_dim = model.fc.in_features
 
-        sums, counts = LinearFiltration._sums_and_counts(
+    # 
+    def _get_means(self,model, num_classes, loader, device, forget_index):
+        h_dim = model.fc.out_features
+
+        # all predictions
+        sums, counts = self._sums_and_counts(
             model=model,
             num_classes=num_classes,
-            retain_loader=retain_loader,
-            forget_loader=forget_loader,
+            loader=loader,
             device=device,
             forget_index=forget_index,
             h_dim=h_dim
         )
-        A = []
 
-        for i in range(num_classes):
-            if counts[i] > 0:
-                A.append(sums[i] / counts[i])
-            else:
-                A.append(torch.zeros(h_dim, device=device))
+        #A = []
 
-        # CORRECT: Stack along dim=0 to make it (num_classes, h_dim)
-        return torch.stack(A, dim=0)
+        counts_safe = counts.unsqueeze(1)
+        A = torch.where(
+            counts_safe > 0, 
+            sums / counts_safe, 
+            torch.zeros_like(sums)
+        )
+        # 6
+        return A
     
 
-    @staticmethod
-    def _compute_z(tensor, forget_index):
-        # Now tensor has shape (num_classes, h_dim) -> tensor.shape[0] is num_classes
+    # 9
+    def _compute_z(self, tensor, forget_index):
+
         K = tensor.shape[0] 
         
-        # pi_a0 should match the feature space dimensions (h_dim)
-        pi_a0 = torch.zeros(tensor.shape[1], device=tensor.device)
+        # pi_a_forget should match the feature space dimensions (h_dim)
+        pi_a_f = torch.zeros(tensor.shape[1], device=tensor.device)
 
-        t_1 = pi_a0
-        a0 = tensor[forget_index, :] # Extracting the row vector for the forgotten class
+        t_1 = pi_a_f
+        # Extracting the row vector for the forgotten class
+        a_f = tensor[forget_index, :] 
 
-        mask_a0 = torch.ones(
-            a0.shape[0],
+        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)) * a0[mask_a0].sum()
+        t_2 = -(1.0 / (K - 1)) * a_f[mask_a_f].sum()
 
         mask_rows = torch.ones(K, dtype=torch.bool, device=tensor.device)
         mask_rows[forget_index] = False
@@ -148,21 +106,23 @@ class LinearFiltration(Strategy):
         return t_1 + t_2 + t_3
 
 
-    @staticmethod
-    def normalise(model, retain_loader, forget_loader, device, forget_index):
+    # Normalisation filtration
+    def normalise(self, model, retain_loader, forget_loader, device, forget_index):
         W = model.fc.weight.data.clone()
         num_classes = W.shape[0]
 
-        A = LinearFiltration._get_means(
+        # we combine the data so we can calculate the mean of prdictions
+        full_loader = _combine_set(retain_loader, forget_loader)
+        # 8
+        A = self._get_means(
             model=model,
             num_classes=num_classes,
-            retain_loader=retain_loader,
-            forget_loader=forget_loader,
+            loader=full_loader,
             device=device,
             forget_index=forget_index
         )
-
-        Z = LinearFiltration._compute_z(tensor=A, forget_index=forget_index)
+        # 9
+        Z = self._compute_z(tensor=A, forget_index=forget_index)
         B_Z_rows = []
 
         for i in range(num_classes):
@@ -172,13 +132,24 @@ class LinearFiltration(Strategy):
                 # Retained classes maintain their original ideal feature directions
                 B_Z_rows.append(A[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(A)
-        
+        # 11 
         W_Z = B_Z @ A_inv @ W
 
+        # 12
         model.fc.weight.copy_(W_Z)
 
         return model
+    
+    # overriden function
+    def _split_data(self, dataset):
+        return get_unlearning_loaders(
+            dataset=dataset,
+            forget_class_idx=self.target_class_index,
+            batch_size = 32
+        )

unlearning/Strategy.py

@@ -16,13 +16,21 @@ class Strategy:
         self.log_file = Path(f"reports/{self.strategy_name}/metrics.txt")
         Util._initialize_log_file(log_file= self.log_file)
 
-    def apply(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
+    def set_target_class(self, target_class_index: int):
+        """Dynamically switch the unlearning target without retraining."""
+        self.target_class_index = target_class_index
+        
+
+    def apply(self, model: nn.Module, dataset) -> nn.Module:
         """
         Wraps the unlearning execution with automated timing and strategy-specific logging.
         DO NOT override this method in subclasses. Override _run instead.
         """
         start_time = time.perf_counter()
 
+
+        retain_loader, forget_loader = self._split_data(dataset)
+        
         # Execute core unlearning logic
         processed_model = self._run(model, forget_loader, retain_loader)
         
@@ -42,3 +50,11 @@ class Strategy:
     def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
         """Subclasses implement their core unlearning logic here."""
         raise NotImplementedError
+    
+    '''
+    different strategies split data in to different partitions differently. 
+    So a strategy will implement its own and since this part is startegy specific. 
+    not all should compute it the same.
+    '''
+    def _split_data(self,dataset):
+        pass

unlearning/WeightFiltration.py

@@ -1,126 +1,135 @@
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, ConcatDataset, Subset
 from unlearning.Strategy import Strategy
-from .wf.WF_Net import WF_Net
+import numpy as np
+from sklearn.metrics import classification_report
+from architectures.WFNet import WF_Net_Model
+
+from sets.Data import vertical_split
 
 class WeightFiltration(Strategy):
-    """
-    Verbatim implementation of Poppi et al.'s WF-Net framework.
-    Directly filters the convolutional weights of a target layer using a learnable
-    channel mask, optimizing it via weight-space regularization.
-    """
-    def __init__(self, target_class_index: int, epochs: int = 10, lr: float = 0.2, gamma: float = 10.0):
+    def __init__(self, 
+                 target_class_index: int, 
+                 num_classes: int = 20, 
+                 epochs: int = 6, 
+                 lr: float = 100.0, 
+                 gamma: float = 0.01,
+                 ):
         super().__init__(target_class_index=target_class_index)
         self.epochs = epochs
         self.lr = lr
         self.gamma = gamma 
-        #self.alpha = None
+        self.num_classes = num_classes
+        self.wf_model = None 
+        self.lambda_1 = 25
     
     
+    def _optimise_filter(self, model: nn.Module, retain_loader: DataLoader, forget_loader: DataLoader, device) -> nn.Module:
 
+        # new WF_Model instance
+        wf_model = WF_Net_Model(
+            device=device, 
+            size=self.num_classes, 
+            original_model=model, 
+            target_class_index=self.target_class_index
+        )
 
-
-    def _optimise_filter(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader, device):
-        # 1. Initialize the wrapper with your pre-trained model
-        num_classes = model.fc.out_features
-        wf_model = WF_Net(original_model=model, num_classes=num_classes).to(device)
-        
-        # 2. ONLY optimize alpha (everything else is frozen inside the wrapper)
-        optimizer = optim.Adam([wf_model.alpha], lr=self.lr)
+        # a WF_net module to be trained (unlearned) to generate alpha
+        wf_net = wf_model.get()
+        optimizer = optim.SGD([wf_net.alpha], lr=self.lr)
         criterion = nn.CrossEntropyLoss()
 
         for epoch in range(self.epochs):
-            forget_iter = iter(forget_loader)
             t_loss_r, t_loss_f = 0.0, 0.0
             steps = 0
             
-            for r_inputs, r_labels in retain_loader:
+            # forget and retain
+            for (r_inputs, r_labels), (f_inputs, f_labels) in zip(retain_loader, forget_loader):
                 r_inputs, r_labels = r_inputs.to(device), r_labels.to(device)
-                
-                # Pull the matching forget batch input
-                try:
-                    f_inputs, _ = next(forget_iter)
-                except StopIteration:
-                    forget_iter = iter(forget_loader)
-                    f_inputs, _ = next(forget_iter)
-                f_inputs = f_inputs.to(device)
+                f_inputs, f_labels = f_inputs.to(device), f_labels.to(device)
                 
                 optimizer.zero_grad()
                 
-                # --- APPLY ALGORITHM 1 FORWARD PASS TO BOTH INPUTS ---
-                # Pass the input batch AND the target unlearn class index
-                outputs_r = wf_model(r_inputs, target_unlearn_class=self.target_class_index)
-                outputs_f = wf_model(f_inputs, target_unlearn_class=self.target_class_index)
+                # retain data paired with randomly selected rows of alpha to compute the retaining loss 
+                random_rows = []
+                for label in r_labels:
+                    allowed = [i for i in range(self.num_classes) if i != label.item()]
+                    random_rows.append(np.random.choice(allowed))
                 
-                # Compute Losses using Poppi et al.'s temperature scaled entropy
+                gate_signals_r = torch.tensor(random_rows, dtype=torch.long, device=device)
+                outputs_r = wf_net(r_inputs, target_class_indices=gate_signals_r)
                 loss_r = criterion(outputs_r, r_labels)
 
-                temperature = 3.0
-                logits_f_scaled = outputs_f / temperature
+                # Forget set is paired with corresponding labels as row selectors for alpha 
+                # and used to compute unlearning loss
+                outputs_f = wf_net(f_inputs, target_class_indices=f_labels)
                 
-                # 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()
+                loss_f = 0.0
+                classes_in_batch = 0 
                 
-                total_loss = loss_r + (self.gamma * loss_f)
+                # every image of class c will unlearn over the same row of alpha_l (poppi et al page 5)
+                for c in range(self.num_classes):
+                    class_mask = (f_labels == c)
+                    if not class_mask.any():
+                        continue
+                        
+                    labels_c = f_labels[class_mask]
+                    
+                    # Slice the existing outputs instead of recalculating a forward pass
+                    outputs_f_c = outputs_f[class_mask] 
+                    
+                    loss_f_ce = criterion(outputs_f_c, labels_c)
+                    
+                    # Poppi et al. suggest employing reciprocal of the forget loss 
+                    # to avoid shortcomings of negative gradient approach
+                    loss_f += 1.0 / (loss_f_ce + 1e-6)
+                    classes_in_batch += 1
+                
+                # Average forget loss by number of distinct classes seen in this batch
+                if classes_in_batch > 0:
+                    loss_f = loss_f / classes_in_batch
+
+                # Regilarisation penality
+                loss_reg = torch.sum(1.0 - torch.sigmoid(wf_net.alpha))
+                                
+                # back propagation
+                total_loss = loss_r + (self.lambda_1 * loss_f) + (self.gamma * loss_reg)
                 total_loss.backward()
                 optimizer.step()
                                 
                 t_loss_r += loss_r.item()
-                t_loss_f += loss_f.item()
+                t_loss_f += loss_f.item() if classes_in_batch > 0 else 0.0
                 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
 
-
-
     def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
         device = next(model.parameters()).device
         model.eval()
         
-        # In WF-Net, the mask targets the last major convolutional block
-        # For ResNet-18, that is the final conv layer in layer4 block 1
-        if hasattr(model, 'layer4') and len(model.layer4) > 1:
-            target_conv = model.layer4[1].conv2
+        if self.wf_model is None:
+            print(">> Initializing and compiling global WF-Net matrix (Run Once for all classes)...")
+            
+            self.wf_model = self._optimise_filter(
+                model, 
+                retain_loader=retain_loader, 
+                forget_loader=forget_loader,
+                device=device
+            )
         else:
-            raise AttributeError("Model architecture does not match expected ResNet-18 structure.")
+            print(f">> Gating matrix loaded. Switching layout to target class index: {self.target_class_index}")
+            self.wf_model.target_class_index = self.target_class_index
             
-        # Store a pristine, non-grad copy of the original trained weights
-        # Shape of conv2.weight: (out_channels, in_channels, kernel_size, kernel_size) -> e.g., (512, 512, 3, 3)
-        original_weights = target_conv.weight.data.clone().detach()
-        out_channels = original_weights.shape[0]
+        return self.wf_model
     
-        # Initialize alpha gate vector matching Poppi et al.'s initialization range
-        # Shape: (out_channels,) -> acting directly as a filter-level gate
-        #self.alpha = nn.Parameter(torch.ones(out_channels, device=device) * 1.5)
-
-        # Freeze the global model graph; only optimize our filter parameter mask
-        for p in model.parameters():
-            p.requires_grad = False
-        #self.alpha.requires_grad = True
-
-        wf_model = self._optimise_filter(
-            model,
-            forget_loader=forget_loader,
-            retain_loader=retain_loader,
-            device=device,
+    def _split_data(self, dataset):
+        return vertical_split(
+            dataset= dataset,
+            batch_size=32,
+            num_classes=self.num_classes
         )
     
-        # --- PERMANENT BAKING STEP ---
-        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)
-            
-            # 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

unlearning/wf/WF_Net.py

@@ -35,50 +35,50 @@ class WF_Net(nn.Module):
         #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)
-        """
+    def forward(self, x: torch.Tensor, target_class_indices: torch.Tensor) -> torch.Tensor:
         # 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
+        # Run layer4 block 0 normally
         x = self.layer4[0](x)
         
+        # -------------------------------------------------------------
+        # HERE IT IS: Save the structural skip connection (identity) 
+        # BEFORE modifying features via block 1's convolutions
+        # -------------------------------------------------------------
         identity = x 
                 
-
-
+        # Now enter layer4 block 1
         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
+        # [Your Step 1 Masking Math happens right here...]
+        batch_alpha = self.alpha[target_class_indices]
+        mask = torch.sigmoid(batch_alpha).view(x.size(0), -1, 1, 1)
         
-        # 3. Pass gated weights straight to functional forward pass: l(Xi, w_hat_l)
+        # Run the functional convolution
         x = F.conv2d(
             x, 
-            weight=w_hat, 
+            weight=self.original_w, 
             bias=self.target_conv.bias, 
             stride=self.target_conv.stride, 
             padding=self.target_conv.padding
         )
+        
+        # Apply your WF-Net channel mask
+        x = x * mask 
         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
+        # -------------------------------------------------------------
+        # HERE IT IS USED: Add the pristine identity back to the gated output
+        # -------------------------------------------------------------
         x = self.layer4[1].relu(x + identity) 
                 
-        # 4. Final Classification Head Sequence
+        # Final Classification Head Sequence
         x = self.avgpool(x)
         x = torch.flatten(x, 1)
         y_out = self.fc(x)