optimised

unlearning/CertifiedUnlearning.py

@@ -0,0 +1,344 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, RandomSampler
+from torch.autograd import grad
+from unlearning.Strategy import Strategy
+
+from sets.Data import *
+
+# Single-Batch Certified Unlearning for DNNs
+
+class CertifiedUnlearning(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 = 50000.0, 
+                 s1: int = 2, s2: int = 350, 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 get_params(self, model: nn.Module, named):
+        """
+        Safely collects named parameter tuples while skipping 
+        InceptionV3 auxiliary layers and tracking gradients.
+        """
+        inner_model = getattr(model, "model", model)
+        
+        # Check if the current architecture is an Inception variant
+        is_inception = inner_model.__class__.__name__.lower() == "inception3"
+
+        params_list = []
+        for name, p in inner_model.named_parameters():
+            if p.requires_grad:
+                # Discard the disconnected auxiliary training branch weights
+                if is_inception and "AuxLogits" in name:
+                    continue
+                # CRITICAL: Append as a tuple so it can be unpacked as (name, param)
+                params_list.append((name, p))
+                
+        return params_list if named else [e[1] for e in params_list]
+    
+    '''
+    def _compute_loss_gradient(self, model, loader, device: torch.device):
+
+        model.eval()
+        criterion = nn.CrossEntropyLoss(reduction='sum')
+        params = self.get_params(model, False) # [p for name, p in model.named_parameters() if p.requires_grad and "AuxLogits" not in name]
+       
+
+        grad_accumulator = [torch.zeros_like(p, device = device) for p in params]
+        total_samples = 0'''
+        
+        # Accumulate true data cross-entropy gradients
+    '''
+        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)
+            
+            mini_grads = list(grad(loss, params, retain_graph=False))
+            for i in range(len(grad_accumulator)):
+                grad_accumulator[i] += mini_grads[i].cpu().detach()
+
+        # Empirical data mean conversion
+        for i in range(len(grad_accumulator)):
+            grad_accumulator[i] /= total_samples
+
+        # L2 weight regularization
+        l2_reg_term = 0.0
+        for param in params:
+            if param.requires_grad:
+                l2_reg_term += torch.sum(param ** 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]
+        '''
+    '''
+        with torch.set_grad_enabled(True):
+            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)
+                
+                mini_grads = grad(loss, params, retain_graph=False)
+                for i in range(len(grad_accumulator)):
+                    grad_accumulator[i] += mini_grads[i]
+
+        # Empirical data mean conversion
+        for i in range(len(grad_accumulator)):
+            grad_accumulator[i] /= total_samples
+
+        # OPTIMIZATION 2: Analytical L2 Regularization Gradient instead of autograd
+        # d/dx (l2_reg * x^2) = 2 * l2_reg * x
+        for i, param in enumerate(params):
+            grad_accumulator[i] += 2 * self.l2_reg * param.detach()
+            
+        return grad_accumulator
+
+    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)
+        elemwise_products = sum(torch.sum(g_elem * v_elem) for g_elem, v_elem in zip(first_grads, v))
+        return grad(elemwise_products, params, create_graph=False)'''
+    '''
+    def _stochastic_newton_update(self, g, dataset, model, device):
+        model.eval()
+        criterion = nn.CrossEntropyLoss()
+        params = self.get_params(model, False) # [p for p in model.parameters() if p.requires_grad]
+        h_res = [torch.zeros_like(p) for p in g]
+
+        # progress
+        total_steps = self.s1 * self.s2
+        step_interval = max(1, total_steps // 100) 
+        
+        global_step = 0
+        current_pct = 0
+
+        sampler = RandomSampler(dataset, replacement=True, num_samples=self.unlearn_bs * self.s2)
+        res_loader = DataLoader(dataset, batch_size=self.unlearn_bs, sampler=sampler)
+        res_iter = iter(res_loader)
+
+        for _ in range(self.s1):
+            h_estimate = [p.clone() for p in g]
+            sampler = RandomSampler(dataset, replacement=True, num_samples=self.unlearn_bs * self.s2)
+            res_loader = DataLoader(dataset, batch_size=self.unlearn_bs, sampler=sampler)
+            res_iter = iter(res_loader)
+            
+            for _ in range(self.s2):
+
+                global_step += 1
+
+                if global_step % step_interval == 0 and current_pct < 100:
+                    current_pct += 1
+                    print(f"\rProgress: {current_pct}% done", end="", flush=True)
+
+                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 = sum(p.pow(2).sum() for p in params)
+                'for param in params:
+                    #if param.requires_grad:
+                    l2_reg_term += torch.sum(param ** 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)):
+                        h_estimate[k].copy_(h_estimate[k] + g[k] - (h_s[k] / self.scale))
+                        #h_res[k] += h_estimate[k] / self.scale
+                        #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
+            with torch.no_grad():
+                for k in range(len(params)):
+                    h_res[k] += h_estimate[k] / self.scale
+                
+        return [p / self.s1 for p in h_res]
+    '''
+
+    def _compute_loss_gradient(self, model, loader, device: torch.device):
+        model.eval()
+        criterion = nn.CrossEntropyLoss(reduction='sum')
+        params = self.get_params(model, False)
+
+        # OPTIMIZATION 1: Keep accumulator on GPU device directly
+        grad_accumulator = [torch.zeros_like(p, device=device) for p in params]
+        total_samples = 0
+        
+        with torch.set_grad_enabled(True):
+            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)
+                
+                mini_grads = grad(loss, params, retain_graph=False)
+                for i in range(len(grad_accumulator)):
+                    grad_accumulator[i] += mini_grads[i]
+
+        # Empirical data mean conversion
+        for i in range(len(grad_accumulator)):
+            grad_accumulator[i] /= total_samples
+
+        # OPTIMIZATION 2: Analytical L2 Regularization Gradient instead of autograd
+        # d/dx (l2_reg * x^2) = 2 * l2_reg * x
+        for i, param in enumerate(params):
+            grad_accumulator[i] += 2 * self.l2_reg * param.detach()
+            
+        return grad_accumulator
+
+    def _hvp(self, loss, params, v):
+        first_grads = grad(loss, params, retain_graph=True, create_graph=True)
+        elemwise_products = sum(torch.sum(g_elem * v_elem) for g_elem, v_elem in zip(first_grads, v))
+        return grad(elemwise_products, params, create_graph=False)
+
+    def _stochastic_newton_update(self, g, dataset, model, device):
+        model.eval()
+        criterion = nn.CrossEntropyLoss()
+        params = self.get_params(model, False)
+        h_res = [torch.zeros_like(p, device=device) for p in g]
+
+        total_steps = self.s1 * self.s2
+        step_interval = max(1, total_steps // 100) 
+        
+        global_step = 0
+        current_pct = 0
+
+        # Create DataLoader outside or use optimal sampling
+        sampler = RandomSampler(dataset, replacement=True, num_samples=self.unlearn_bs * self.s2 * self.s1)
+        res_loader = DataLoader(dataset, batch_size=self.unlearn_bs, sampler=sampler)
+        res_iter = iter(res_loader)
+
+        for _ in range(self.s1):
+            h_estimate = [p.clone() for p in g]
+            
+            for _ in range(self.s2):
+                global_step += 1
+
+                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)
+                
+                # OPTIMIZATION 3: Clean up graph creation for loss & L2
+                outputs = model(data)
+                loss = criterion(outputs, target)
+                
+                l2_reg_term = sum(p.pow(2).sum() for p in params)
+                loss += (self.l2_reg + self.gamma) * l2_reg_term
+                
+                h_s = self._hvp(loss, params, h_estimate)
+                
+                # OPTIMIZATION 4: Avoid deprecated .data, use detach() and in-place ops
+                with torch.no_grad():
+                    for k in range(len(params)):
+                        h_estimate[k].copy_(h_estimate[k] + g[k] - (h_s[k] / self.scale))
+                        
+
+                # feed back on status
+                if global_step % step_interval == 0 and current_pct < 100:
+                    current_pct += 1
+                    print(f"\rProgress: {current_pct}% done", end="", flush=True)
+                        
+            with torch.no_grad():
+                for k in range(len(params)):
+                    h_res[k] += h_estimate[k] / self.scale
+                
+        return [p / self.s1 for p in h_res]
+
+
+    def _certify(self, model, device, delta, full_certification):
+        certification = "full " if full_certification else "partial"
+        print(f"Performing {certification} certification")
+        
+        delta_idx = 0 
+        
+        # named_parameters to monitor layer positions
+        for name, param in self.get_params(model, True):
+            if param.requires_grad:
+                noise = self.std * torch.randn(param.data.size(), device=device)
+
+                if full_certification:
+                    param.data.add_(delta[delta_idx] + noise)
+                else:
+                    # option for applying certification only to last layers
+                    # deprecated
+                    if "layer4" in name or "fc" in name:
+                        param.data.add_(delta[delta_idx] + noise)
+                    else:
+                        # Keep early low-level vision filters entirely pristine
+                        pass
+                
+                # Move to the next calculated Hessian vector block only after a valid update step
+                delta_idx += 1
+                
+        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 Forget set...")
+        g = self._compute_loss_gradient(model, forget_loader, device)
+        
+        print(">> Estimating non-convex inverse Hessian trajectories via Taylor series...")
+        dataset = retain_loader.dataset 
+        delta = self._stochastic_newton_update(g, dataset, model, device)
+        
+        print(">> Applying parameter removal adjustments (-delta)...")
+        model = self._certify(
+            model= model, 
+            device = device, 
+            delta = delta, 
+            full_certification = True
+        )
+                
+        print(">> Certified Unlearning process completed successfully.")
+        return model
+    
+    
+    # overriden function
+    def _split_data(self, dataset):
+        # Certified unlearning does require both forget and retain sets
+        # split horizontaly. one class to forget and the rest to retain
+        return get_unlearning_loaders(
+            dataset=dataset,
+            forget_class_idx=self.target_class_index,
+            batch_size = 32
+        )

unlearning/LinearFiltration.py

@@ -0,0 +1,168 @@
+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):
+        super().__init__(target_class_index=target_class_index)
+        self.A = None
+
+    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
+        model.eval()
+        # Freeze internal params
+        for param in model.parameters():
+            param.requires_grad = False
+
+        device = next(model.parameters()).device
+            
+        return self.normalise(
+            model=model, 
+            retain_loader=retain_loader, 
+            forget_loader=forget_loader,
+            device=device,
+            forget_index=self.target_class_index
+        )
+    
+
+    def _get_classifier(self, model: nn.Module) -> nn.Linear:
+
+        inner_model = getattr(model, "model", model)
+        
+        # looking for standard naming conventions in named modules
+        for name, module in inner_model.named_modules():
+            # Check if it's our target linear layer
+            if (name == "fc" or name == "classifier") and isinstance(module, nn.Linear):
+                return module
+                
+            # Handle models (like EfficientNet) where the classifier is a Sequential block
+            if name == "classifier" and isinstance(module, nn.Sequential):
+                for sub_module in reversed(list(module.children())):
+                    if isinstance(sub_module, nn.Linear):
+                        return sub_module
+
+        # scan backwards for the last Linear layer
+        for module in reversed(list(inner_model.modules())):
+            if isinstance(module, nn.Linear):
+                return module
+                
+        raise RuntimeError(f"Could not locate a linear classification head for {model.__class__.__name__}")
+    
+  
+
+    def _compute_A(self, model, num_classes, loader, device):
+        model.eval()
+        
+        
+        # Initialize tracking tensors
+        sums = torch.zeros(num_classes, num_classes, device=device)
+        counts = torch.zeros(num_classes, device=device)
+
+        with torch.no_grad():
+            for inputs, targets in loader:
+                inputs, targets = inputs.to(device), targets.to(device)
+                
+                # the logit predictions
+                outputs = model(inputs) 
+                
+                # One-hot encode targets to act as a routing mask
+                one_hot = torch.nn.functional.one_hot(targets, num_classes=num_classes).float()
+                
+                # add 
+                sums += torch.t(one_hot) @ outputs
+                
+                # Sum columns of one-hot to get counts per class in this batch
+                counts += one_hot.sum(dim=0)
+
+        # means
+        counts_safe = counts.unsqueeze(1)
+        print(f"COUNTS IS >>>>>>>>> {counts_safe}")
+        self.A = torch.where(
+            counts_safe > 0, 
+            sums / counts_safe, 
+            torch.zeros_like(sums)
+        )
+
+    # 9
+    def _compute_z(self, tensor, forget_index):
+
+        K = tensor.shape[0] 
+        
+        
+        pi_a_f = torch.zeros(tensor.shape[1], device=tensor.device)
+
+        t_1 = pi_a_f
+        # row vector for the forgotten class
+        a_f = tensor[forget_index, :] 
+
+        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)) * a_f[mask_a_f].sum()
+
+        mask_rows = torch.ones(K, dtype=torch.bool, device=tensor.device)
+        mask_rows[forget_index] = False
+
+        r_A = tensor[mask_rows, :]
+        t_3 = (1.0 / ((K - 1)) ** 2) * r_A.sum()
+
+        return t_1 + t_2 + t_3
+
+
+    # Normalisation filtration
+    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
+        # because mean of all prdictions is the same for all
+        if self.A is None:
+            self._compute_A(
+                model = model,
+                num_classes = num_classes,
+                loader = full_loader,
+                device = device
+            )
+
+        # 9
+        Z = self._compute_z(tensor=self.A, forget_index=forget_index)
+        B_Z_rows = []
+
+        for i in range(num_classes):
+            if i == forget_index:
+                B_Z_rows.append(Z)
+            else:
+                # Retained classes maintain their original ideal feature directions
+                B_Z_rows.append(self.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(self.A)
+        # 11 
+        W_Z = B_Z @ A_inv @ W
+
+        # 12
+        clf = self._get_classifier(model)
+        clf.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

@@ -0,0 +1,62 @@
+
+import torch.nn as nn
+import time
+import os
+from pathlib import Path
+from torch.utils.data import DataLoader
+import Util
+
+class Strategy:
+    """Abstract base class for unlearning algorithms with automated, strategy-specific logging."""
+    
+    def __init__(self, target_class_index):
+        # Dynamically set file name based on the class name (e.g., 'NormalizingLinearFiltration.txt')
+        self.strategy_name = self.__class__.__name__
+        self.target_class_index = target_class_index
+
+    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:
+        log_file = Path(f"reports/{self.strategy_name}/{model.__class__.__name__}/time_metrics.txt")
+        Util._initialize_log_file(log_file=log_file)
+        """
+        Wraps the unlearning execution with automated timing and strategy-specific logging.
+        DO NOT override this method in subclasses. Override _run instead.
+        """
+    
+
+
+        retain_loader, forget_loader = self._split_data(dataset)
+        
+        # record start time to evaluate efficiency
+        start_time = time.perf_counter()
+        # Execute core unlearning logic
+        processed_model = self._run(model, forget_loader, retain_loader)
+        
+        end_time = time.perf_counter()
+        execution_time = end_time - start_time
+        
+        # Log to the strategy's specific file
+        Util.log_metric(
+            log_file=log_file, 
+            execution_time=execution_time
+        )
+
+        print(f"[{self.strategy_name}] Completed in {execution_time:.6f} seconds. Saved to {log_file}")
+        
+        return processed_model
+
+    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/WF.py

@@ -0,0 +1,107 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from unlearning.Strategy import Strategy
+from .wf.WF_Net import WF_Net
+
+class WeightF(Strategy):
+    """
+    Verbatim implementation of Poppi et al.'s WF-Net framework modified
+    for static, single-class unlearning extraction.
+    """
+    def __init__(self, target_class_index: int, epochs: int = 10, lr: float = 0.2, gamma: float = 10.0):
+        super().__init__(target_class_index=target_class_index)
+        self.epochs = epochs
+        self.lr = lr
+        self.gamma = gamma 
+
+    def _optimise_filter(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader, device):
+        num_classes = model.fc.out_features
+        wf_model = WF_Net(original_model=model, num_classes=num_classes).to(device)
+                
+        # Optimize only the specific alpha masks
+        optimizer = optim.Adam([wf_model.alpha], lr=self.lr)
+        criterion = nn.CrossEntropyLoss() # Default reduction is 'mean'
+        
+        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:
+                r_inputs, r_labels = r_inputs.to(device), r_labels.to(device)
+                                
+                try:
+                    f_inputs, _ = next(forget_iter)
+                except StopIteration:
+                    forget_iter = iter(forget_loader)
+                    f_inputs, _ = next(forget_iter)
+                f_inputs = f_inputs.to(device)
+                                
+                optimizer.zero_grad()
+                                
+                # Forward Pass
+                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 Loss (Mean over batch)
+                loss_r = criterion(outputs_r, r_labels)
+                                
+                # Forget Loss (Corrected to Mean over batch)
+                temperature = 1.0
+                logits_f_scaled = outputs_f / temperature
+                
+                # 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
+                
+    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
+        device = next(model.parameters()).device
+        model.eval()                 
+
+        if hasattr(model, 'layer4') and len(model.layer4) > 1:
+            target_conv = model.layer4[1].conv2
+        else:
+            raise AttributeError("Model architecture does not match expected ResNet-18 structure.")
+                    
+        original_weights = target_conv.weight.data.clone().detach()
+        out_channels = original_weights.shape[0]
+        
+        # Freeze global network layers
+        for p in model.parameters():
+            p.requires_grad = False
+            
+        wf_model = self._optimise_filter(
+            model,
+            forget_loader=forget_loader,
+            retain_loader=retain_loader,
+            device=device,
+        )
+        
+        # --- 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/WeightFiltration.py

@@ -0,0 +1,139 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, ConcatDataset, Subset
+from unlearning.Strategy import Strategy
+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):
+    def __init__(self, 
+                 target_class_index: int, 
+                 arch,
+                 num_classes: int = 20, 
+                 epochs: int = 6, 
+                 lr: float = 100.0, 
+                 gamma: float = 0.01,
+                 lambda_1 = 25
+                 
+                 ):
+        super().__init__(target_class_index=target_class_index)
+        self.epochs = epochs
+        self.lr = lr
+        self.gamma = gamma 
+        self.num_classes = num_classes
+        self.wf_model = None 
+        self.lambda_1 = lambda_1
+        self.arch = arch
+    
+
+    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, 
+            arch=self.arch,
+            size=self.num_classes, 
+            original_model=model, 
+            target_class_index=self.target_class_index
+        )
+
+        wf_net = wf_model.get()
+        optimizer = optim.SGD([wf_net.alpha], lr=self.lr)
+        # Use reduction='none' so we can manipulate individual item losses
+        criterion_none = nn.CrossEntropyLoss(reduction='none')
+        criterion_mean = nn.CrossEntropyLoss()
+
+        for epoch in range(self.epochs):
+            t_loss_r, t_loss_f = 0.0, 0.0
+            steps = 0
+
+            # 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)
+                f_inputs, f_labels = f_inputs.to(device), f_labels.to(device)
+                
+                optimizer.zero_grad()
+                
+                # retain data paired with randomly selected rows of alpha to compute the retaining loss 
+                random_offset = torch.randint(0, self.num_classes - 1, size=r_labels.shape, device=device)
+                gate_signals_r = torch.where(random_offset >= r_labels, random_offset + 1, random_offset)
+                
+                outputs_r = wf_net(r_inputs, target_class_indices=gate_signals_r)
+                loss_r = criterion_mean(outputs_r, r_labels)
+
+                # 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)
+                
+                # Calculate loss for every single item in the batch at once
+                per_item_forget_loss = criterion_none(outputs_f, f_labels)
+                
+                # Use a scatter/sum approach to get class-wise losses without a Python loop
+                # Create a mask of unique classes present in this batch
+                unique_classes, inverse_indices = torch.unique(f_labels, return_inverse=True)
+                classes_in_batch = unique_classes.size(0)
+                
+                if classes_in_batch > 0:
+                    # average CE loss per class
+                    class_loss_sums = torch.zeros(classes_in_batch, device=device)
+                    class_loss_sums.scatter_add_(0, inverse_indices, per_item_forget_loss)
+                    
+                    class_counts = torch.zeros(classes_in_batch, device=device)
+                    class_counts.scatter_add_(0, inverse_indices, torch.ones_like(per_item_forget_loss))
+                    
+                    mean_class_ce_loss = class_loss_sums / class_counts
+                    
+                    # Poppi et al. suggest employing reciprocal of the forget loss 
+                    # to avoid shortcomings of negative gradient approach
+                    loss_f = torch.mean(1.0 / (mean_class_ce_loss + 1e-6))
+                else:
+                    loss_f = torch.tensor(0.0, device=device)
+
+                # Regularisation penalty
+                loss_reg = torch.sum(1.0 - torch.sigmoid(wf_net.alpha))
+                                
+                # Backpropagation
+                total_loss = loss_r + (self.lambda_1 * loss_f) + (self.gamma * loss_reg)
+                total_loss.backward()
+                optimizer.step()
+                                
+                # Keep tracking stats
+                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
+    
+
+    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
+        device = next(model.parameters()).device
+        model.eval()
+        
+        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:
+            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
+            
+        return self.wf_model
+    
+    def _split_data(self, dataset):
+        return vertical_split(
+            dataset= dataset,
+            batch_size=32,
+            num_classes=self.num_classes
+        )
+