certified

.gitignore

@@ -1,3 +1,5 @@
+# Created by venv; see https://docs.python.org/3/library/venv.html
+*
 # Virtual Environment (the folders Git saw)
 bin/
 lib/

DataAnalyser.py

@@ -1,6 +1,6 @@
 
 #from Data import *
-from datasets.Casia import *
+from sets.Casia import *
 
 '''
 Because the size of samples per class had the biggest impact

Tune.py

@@ -41,6 +41,7 @@ EPOCHS = 10
 # Inception = 299
 RESOLUTION = 224
 
+FINETUNE = False # whether to fintune or just load finetuned model from dir
 # model architecture options are 
 #  - RESNET18
 #  - RESNET50
@@ -112,19 +113,24 @@ device = SetUp.get_device()
 for i in range(0,1):#CLASS_SIZE):
     FORGET_CLASS_IDX = i
     # Create model using Factory
-    model = Model.create(
-        arch = arch, 
-        device = device, 
-        size = CLASS_SIZE)
 
-    # we may need to load existing model or finetune
+    model = None
-    #model.train(
-    #        epochs = EPOCHS, 
-    #        loader = train_loader, 
-    #        rate = LR_RATE)
 
-        # save.
+    if FINETUNE:
-    #model.save(filename=arch.name.lower())
+        model = Model.create(
+            arch = arch, 
+            device = device, 
+            size = CLASS_SIZE)
+
+        # we may need to load existing model or finetune
+        model.train(
+                epochs = EPOCHS, 
+                loader = train_loader, 
+                rate = LR_RATE)
+
+            # save.
+        file_name = f"{arch.name.lower}_{dataset_name.name.lower()}"
+        model.save(filename=arch.name.lower())
 
 
     # done tuning
@@ -147,18 +153,21 @@ for i in range(0,1):#CLASS_SIZE):
 
     # Evaluate
     current_mode = "Finetuned"
-    #accuracy, report_dict = model.evaluate(
+    if FINETUNE:
-    #    loader = test_loader,
-    #    mode=current_mode
-    #)
 
-    Util._log_to_csv(
+        #current_mode = "Finetuned"
-        arch=model.__class__.__name__, 
+        accuracy, report_dict = model.evaluate(
-        mode = current_mode, 
+            loader = test_loader,
-        accuracy=accuracy, 
+            mode=current_mode
-        report_dict=report_dict, 
+        )
-        strategy="base"
+
-    )
+        Util._log_to_csv(
+            arch=model.__class__.__name__, 
+            mode = current_mode, 
+            accuracy=accuracy, 
+            report_dict=report_dict, 
+            strategy="base"
+        )
 
     # unlearning algorithms
     #linear_filtration = LinearFiltration(target_class_index=FORGET_CLASS_IDX)
@@ -167,7 +176,14 @@ for i in range(0,1):#CLASS_SIZE):
     #weight_filtration = WeightFiltration(num_classes = CLASS_SIZE,target_class_idx=FORGET_CLASS_IDX)
     #weight_filtration.apply(reloaded.model)
 
-    certified_removal = CertifiedRemoval(target_class_index=FORGET_CLASS_IDX,removal_bound=0.05, epsilon=0.5, l2_reg=15)
+    certified_removal = CertifiedRemoval(
+        target_class_index=FORGET_CLASS_IDX,
+        s1=2,
+        s2=500,
+        unlearn_bs=2,
+        scale=100.0,    # Drop scale to match lower s2 depth
+        std=0.00001)
+    #,removal_bound=0.05, epsilon=0.5, l2_reg=15)
     #certified_removal.apply(reloaded.model)
 
     # to be unlearned
@@ -200,6 +216,12 @@ for i in range(0,1):#CLASS_SIZE):
         #   loader = test_loader
         #)
 
+        if not FINETUNE:
+            reloaded.evaluate(
+            loader = test_loader,
+            mode=current_mode
+        )
+
         # Unlearning
         # train loaders passed here
         strategy.apply(reloaded.model, forget_train_loader, retain_train_loader)

Tune_new.py

@@ -10,6 +10,9 @@ 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
 
 # Global Hyperparameters
 CLASS_SIZE = 20
@@ -17,7 +20,7 @@ BATCH_SIZE = 16
 SAMPLE_SIZE = 30
 TRAINING_SAMPLE = 27
 RESOLUTION = 224
-ARCH = Architecture.RESNET50
+ARCH = Architecture.RESNET18
 
 
 # Data preparation and model setup
@@ -27,14 +30,17 @@ def prepare_data_and_model_environment():
     train-test class splits, and configures the architecture base.
     """
     device = SetUp.get_device()
-    dataset_name = Set_Name.CELEBA
+    dataset_name = Set_Name.CASIAFACES
+    if dataset_name == Set_Name.CASIAFACES:
+        SAMPLE_SIZE = 400
+        TRAINING_SAMPLE = 320
 
     dataset = get_set(set_name=dataset_name)
     print(f"> {dataset.__class__.__name__} dataset loaded")
 
     # Select target identities (deterministic top sample identities)
     selected_identities = select_top_ids(dataset=dataset, class_size=CLASS_SIZE)
-    print(f'> Selected {CLASS_SIZE} random identity classes from CelebA dataset.')
+    print(f'> Selected {CLASS_SIZE} random identity classes from {dataset_name.name} dataset.')
     print(f'> A class has {TRAINING_SAMPLE} train and {SAMPLE_SIZE - TRAINING_SAMPLE} test samples')
 
     # Isolate sample index partitions
@@ -81,6 +87,8 @@ def prepare_data_and_model_environment():
 
 # Fine tunning and evaluation
 def run_finetuning_or_baseline_eval(env_dict, run_training=False, lr_rate=0.0001, epochs=10):
+
+    
     """
     Handles model training (if flag is true) and logs the baseline fine-tuned 
     performance to file metrics.
@@ -91,13 +99,15 @@ def run_finetuning_or_baseline_eval(env_dict, run_training=False, lr_rate=0.0001
 
     test_loader = DataLoader(test_data, batch_size=BATCH_SIZE, shuffle=False)
 
-    # Optional training configuration switch
+    
-    if run_training:
+    train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)
-        train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)
+
-        print(f"Starting training on {env_dict['device']}...")
+    if not run_training:
-        model.train(epochs=epochs, loader=train_loader, rate=lr_rate)
+        return
-        model.save(filename=ARCH.name.lower())
+    #print(f"Starting training on {env_dict['device']}...")
-        print(f"Model saved to trained_models/{ARCH.name.lower()}.pth")
+    model.train(epochs=epochs, loader=train_loader, rate=lr_rate)
+    model.save(filename=ARCH.name.lower())
+    print(f"Model saved to trained_models/{ARCH.name.lower()}.pth")
 
     print(f"Total test images for these {CLASS_SIZE} classes: {len(test_data)}")
 
@@ -119,7 +129,7 @@ def run_finetuning_or_baseline_eval(env_dict, run_training=False, lr_rate=0.0001
 
 
 # Unlearning and strategy eval
-def run_unlearning_and_strategy_eval(env_dict, forget_class_idx):
+def run_unlearning_and_strategy_eval(env_dict, forget_class_idx, strategy, evaluate = False):
     """
     Reloads a clean model state, applies the isolated unlearning framework,
     and runs specific target evaluation domain checks.
@@ -128,13 +138,34 @@ def run_unlearning_and_strategy_eval(env_dict, forget_class_idx):
     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
-    certified_removal = CertifiedRemoval(
+    # increase s2, decrease scale ---sweet spot
+    '''certified_removal = CertifiedRemoval(
         target_class_index=forget_class_idx,
-        removal_bound=0.05, 
+        s1=4,
-        epsilon=0.5, 
+        s2=350,  # 350 best
-        l2_reg=15
+        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(
@@ -147,11 +178,17 @@ def run_unlearning_and_strategy_eval(env_dict, forget_class_idx):
     # Instantiate a clean copy of the model to keep weights isolated
     reloaded = Model.create(arch=ARCH, device=device, size=CLASS_SIZE)
     reloaded.load(arch=ARCH)
+
+    if evaluate:
+        reloaded.evaluate(
+            loader=retain_test_loader, mode="finetuned"
+        )
+
     print("fine tunned model loaded into evaluation sandbox")
 
     # Execute strategic parameter unlearning step
-    certified_removal.apply(reloaded.model, forget_train_loader, retain_train_loader)
+    strategy.apply(reloaded.model, forget_train_loader, retain_train_loader)
-    strategy_in_use = certified_removal.__class__.__name__
+    strategy_in_use = strategy.__class__.__name__
 
     # Define validation tracking steps dynamically
     evaluation_domains = [
@@ -180,10 +217,62 @@ if __name__ == "__main__":
     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=True)
+    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,
+            l2_reg=0.000002,       
+            gamma=0.1,           
+            scale=  20000,# 16400.0, # took ages to reach this sweet spot     
+            s1=2,
+            s2=300,               
+            std=0.00001,
+            unlearn_bs=16
+        )
+
+        # works perfectly
+        linear_filtration = LinearFiltration(
+
+            target_class_index=i
+        )
+
+        weight_filtration = WeightFiltration(
+            target_class_index=i,
+            epochs=3,
+            lr=0.5,
+            gamma=150
+        )
+
+        strategies = [
+            certified_unlearning,
+            # weight_filtration,
+            # linear_filtration
+        ]
+
         print(f"\n>>> Executing Unlearning Framework for Target Identity Index: {i} <<<")
-        run_unlearning_and_strategy_eval(runtime_environment, forget_class_idx=i)
+        for strategy in strategies:
+            run_unlearning_and_strategy_eval(
+                runtime_environment, 
+                forget_class_idx=i, 
+                strategy=strategy,
+                evaluate= not finetuning
+            )

architectures/Model.py

@@ -139,7 +139,7 @@ class Model(ABC):
 
 
 
-    # Using the factory patern here
+    # factory
     @staticmethod
     def create(arch, device, size):
         print(f'>> MODEL ARCHITECTURE  >> {arch.name}.')
@@ -151,6 +151,11 @@ class Model(ABC):
                 from architectures.ResNet18 import ResNet18
                 return ResNet18(device, size)
             
+            # ResNet34
+            case Architecture.RESNET34: 
+                from architectures.ResNet34 import ResNet34
+                return ResNet34(device, size)
+            
             # ResNet50
             case Architecture.RESNET50:
                 from architectures.ResNet50 import ResNet50
@@ -190,6 +195,7 @@ from enum import Enum, auto
 class Architecture(Enum):
     RESNET18 = auto()
     RESNET50 = auto()
+    RESNET34 = auto()
     INCEPTION = auto()
     DENSENET121 = auto()
     GOOGLENET = auto()

sets/CelebA.py

@@ -9,10 +9,10 @@ class CelebA(Data):
 
     def get_set(self):
         set = datasets.CelebA(
-            root = "./data", 
+            root = "../data", 
             split='all', 
             target_type='identity', 
-            download=True, 
+            download=False, 
             transform=None
         )
         # set the target first

sets/Extractor.py

@@ -75,7 +75,7 @@ def extract_selected_binary(rec_path, idx_path, output_dir, top_labels):
             current_count = save_counters[label]
             img_filename = f"{current_count}.jpg"
             img_path = os.path.join(target_folder, img_filename)
-            if(current_count > 200):
+            if(current_count > 405):
                 continue
             
             with open(img_path, 'wb') as img_f:
@@ -119,9 +119,9 @@ if __name__ == "__main__":
 '''
 if __name__ == "__main__":
     base_dir = os.path.dirname(os.path.abspath(__file__))
-    REC = os.path.join(base_dir, 'casia', 'train.rec')
+    REC = os.path.join(base_dir, '../data/casia-set', 'train.rec')
-    IDX = os.path.join(base_dir, 'casia', 'train.idx')
+    IDX = os.path.join(base_dir, '../data/casia-set', 'train.idx')
-    OUT = os.path.join(base_dir, 'casia-set')
+    OUT = os.path.join(base_dir, '../data/casia-set')
 
     # Step 1: Trust the binary, not the text file
     top_verified_labels = get_top_identities_binary(REC, IDX, top_n=50)

unlearning/CertifiedRemoval.py

@@ -1,127 +1,214 @@
 import torch
 import torch.nn as nn
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, RandomSampler
+from torch.autograd import grad
 from unlearning.Strategy import Strategy
 
 class CertifiedRemoval(Strategy):
     """
-    Implements Certified Removal (Guo et al.) adapted for deep architectures
+    Implements Certified Unlearning for non-convex DNNs (Zhang et al.).
-    like ResNet50 by isolating and updating the final classification layer.
+    Uses a modified, stabilized stochastic Newton step using Taylor-expansion 
+    HVP estimation across the entire parameter space, capped with calibrated noise.
     """
-    def __init__(self, removal_bound: float, epsilon: float, l2_reg: float = 0.1):
+    def __init__(self, target_class_index: int, l2_reg: float = 0.0005, 
-        super().__init__()
+                 gamma: float = 0.01, scale: float = 1000.0, 
-        self.removal_bound = removal_bound  # gamma in the paper
+                 s1: int = 10, s2: int = 1000, std: float = 0.001, unlearn_bs: int = 2):
-        self.epsilon = epsilon              # Privacy budget
+        super().__init__(target_class_index)
-        self.l2_reg = l2_reg                # Lambda regularization term
+        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_features(self, backbone: nn.Module, loader: DataLoader, device: torch.device):
+    '''
-        """Passes data through the frozen ResNet backbone to extract embedding features."""
+    def _compute_loss_gradient(self, model, loader, device: torch.device):
-        backbone.eval()
+        model.eval()
-        all_features = []
+        criterion = nn.CrossEntropyLoss(reduction='sum')
-        all_labels = []
+        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
         
-        with torch.no_grad():
+        for data, targets in loader:
-            for inputs, labels in loader:
+            total_samples += targets.shape[0]
-                inputs = inputs.to(device)
+            data, targets = data.to(device), targets.to(device)
-                # Pass through backbone to get the 2048-dimensional feature vector
+            outputs = model(data)
-                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)
+            mini_grads = list(grad(criterion(outputs, targets), params))
+            for i in range(len(grad_accumulator)):
+                grad_accumulator[i] += mini_grads[i].cpu().detach()
 
-    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
+        for i in range(len(grad_accumulator)):
-        """
+            grad_accumulator[i] /= total_samples
-        Entry point expected by your Model.unlearn() architecture interface.
+
-        Applies Certified Removal strictly to the final linear layer (model.fc).
+        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
 
-        # Isolate the final NN (Fully connected) layer from the model
+        num_forget = len(forget_loader.dataset)
-        linear_head = model.fc
+        num_retain = len(retain_loader.dataset)
-        # Temporarily turn the fc layer into a identity pass-through
+        scaling_ratio = num_forget / num_retain
-        model.fc = nn.Identity() 
         
-        print(">> Extracting deep features from model backbone...")
+        print(">> Calculating base gradients over target FORGET set...")
-        retain_features, retain_labels = self._get_features(model, retain_loader, device)
+        # FIX 3: Base gradients MUST be evaluated from forget_loader to drop target class distributions
-        forget_features, forget_labels = self._get_features(model, forget_loader, device)
+        g = self._compute_loss_gradient(model, forget_loader, device)
         
-        # Restore the linear head back
+        print(">> Estimating non-convex inverse Hessian trajectories via Taylor series...")
-        model.fc = linear_head
+        retain_dataset = retain_loader.dataset 
+        delta = self._stochastic_newton_update(g, retain_dataset, model, device)
         
-        # Extract weights from the classification layer
+        print(">> Applying stabilized parameter adjustments and randomized certification noise...")
-        # w shape: [num_classes, 2048]
+        with torch.no_grad():
-        w = model.fc.weight.data.clone().cpu()
+            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)
                 
-        # Compute the Exact Hessian Matrix over the remaining (retained) features
+        print(">> Certified Unlearning process completed successfully across the complete landscape.")
-        # Formula: H = (X^T * X) / N + lambda * I
+        return model'''
-        N_retain = retain_features.size(0)
+    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
-        hessian = self._compute_hessian(retain_features=retain_features, retain_features_size = N_retain)
+        device = next(model.parameters()).device
 
-        grad_forget = self._compute_loss_gradient(
+        print(">> Calculating stable base gradients over the RETAIN set...")
-            forget_labels=forget_labels,
+        # To match the author's snippet perfectly, g MUST be computed on the retain data.
-            forget_features=forget_features, 
+        # If this loader is too large for your VRAM, use a smaller batch size (e.g. 16 or 32) 
-            model_weights=w)
+        # in your main training script when creating retain_loader.
-        #torch.matmul(error.t(), forget_features) / forget_features.size(0)
+        g = self._compute_loss_gradient(model, retain_loader, device)
         
-        # Compute the Newton step update via solving: H * Delta_W^T = Grad_forget^T
+        print(">> Estimating non-convex inverse Hessian trajectories via Taylor series...")
-        delta_w = self._compute_newton_step(
+        retain_dataset = retain_loader.dataset 
-            tensor = hessian,
+        delta = self._stochastic_newton_update(g, retain_dataset, model, device)
-            gradient= grad_forget
-        )
-        # Apply the Certified Removal update rule: W_new = W + Delta_W
-        new_w = w + delta_w
-        # Calibrate noise based on your epsilon budget 
-        # (Guo et al. use a perturbation based on the regularization lambda and epsilon)
-        sigma = 2.0 / (self.l2_reg * self.epsilon)
-        noise = torch.randn_like(new_w) * (sigma / N_retain)
-        new_w = new_w + noise
         
-        # Theoretical Guarantee verification
+        print(">> Applying parameter removal adjustments (-delta)...")
-        norm_delta = torch.norm(delta_w).item()
+        with torch.no_grad():
-        if norm_delta > self.removal_bound:
+            for i, param in enumerate(model.parameters()):
-            print(f"!! Warning: Removal budget exceeded! Norm: {norm_delta:.4f} > Bound: {self.removal_bound}")
+                if param.requires_grad:
-        else:
+                    noise = self.std * torch.randn(param.data.size(), device=device)
-            print(f">> Certificate valid. Norm: {norm_delta:.4f} <= Bound: {self.removal_bound}")
                     
-        # Push updated parameters back into the model instance in-place
+                    # MATCHING THE SNIPPET: Subtract delta exactly as the authors do
-        model.fc.weight.data = new_w.to(device)
+                    # This removes the influence trace of the omitted data.
+                    param.data.add_(-delta[i] + noise)
                 
-        print(">> Certified Removal process completed successfully.")
+        print(">> Certified Unlearning process completed successfully.")
         return model
-    
-
-    # computing the hessian matrix
-    def _compute_hessian(self, retain_features, retain_features_size):
-        print(">> Computing exact Hessian matrix...")
-        # N_retain = retain_features.size(0)
-        X_T_X = torch.matmul(retain_features.t(), retain_features)
-        reg_matrix = self.l2_reg * torch.eye(retain_features.size(1))
-        return (X_T_X / retain_features_size) + reg_matrix
-    
-
-    def _compute_loss_gradient(self, forget_features, forget_labels, model_weights):
-        print(">> Calculating forget set gradients...")
-        num_classes = model_weights.size(0)
-        # Pass features through linear layer weights to get logits
-        logits_forget = torch.matmul(forget_features, model_weights.t()) 
-        # Apply softmax to get true class probabilities
-        preds_softmax = torch.softmax(logits_forget, dim=1)
-        
-        forget_labels_one_hot = torch.nn.functional.one_hot(forget_labels, num_classes=num_classes).float()
-        
-
-        error = preds_softmax - forget_labels_one_hot
-        # grad_forget shape: [num_classes, 2048]
-        return torch.matmul(error.t(), forget_features) / forget_features.size(0)
-    
-
-    def _compute_newton_step(self,tensor, gradient):
-        print(">> Solving Newton step via system optimization...")
-        try:
-            delta_w_t = torch.linalg.solve(tensor, gradient.t())
-            delta_w = delta_w_t.t()
-        except RuntimeError:
-            print(">> Warning: Hessian matrix is singular. Falling back to pseudo-inverse.")
-            delta_w = torch.matmul(gradient, torch.linalg.pinv(tensor).t())
-        return delta_w

unlearning/LinearFiltration.py

@@ -1,47 +1,184 @@
-
 import torch
 import torch.nn as nn
 from .Strategy import Strategy
 from torch.utils.data import DataLoader
 
 class LinearFiltration(Strategy):
-    def __init__(self,target_class_index):
+    def __init__(self, target_class_index):
-        super().__init__(target_class_index = target_class_index)
+        super().__init__(target_class_index=target_class_index)
 
     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
 
-        with torch.no_grad():
+        device = next(model.parameters()).device
-            W = model.fc.weight.data.clone()
-            num_classes = W.shape[0]
             
-            A = self._calculate_filtration_matrix(num_classes, self.target_class_index, W.device)
+        return self.normalise(
-            sanitized_W = torch.mm(A, W)
+            model=model, 
-            model.fc.weight.copy_(sanitized_W)
+            retain_loader=retain_loader, 
-            # Filter the bias (if the layer uses one)
+            forget_loader=forget_loader,
-            if model.fc.bias is not None:
+            device=device,
-                b = model.fc.bias.data.clone()
+            forget_index=self.target_class_index
-                # b is a 1D tensor of shape (num_classes), 
+        )
-                # so we use torch.mv (matrix-vector multiplication) or unsqueeze it
-                sanitized_b = torch.mv(A, b)
-                model.fc.bias.copy_(sanitized_b)
     
-        return model
+    # 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
         
-        # The row of the forgotten class should average all other classes
         for j in range(num_classes):
             if j == forget_class:
-                # we zero the forget class
                 A[forget_class, j] = 0.0
             else:
-                # and we distribute the output to the remaining
                 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):
+        model.eval()
+
+        sums = torch.zeros(num_classes, h_dim, device=device)
+        counts = torch.zeros(num_classes, device=device)
+
+        # Generate values for retain
+        with torch.no_grad():
+            for inputs, targets in retain_loader:
+                inputs = inputs.to(device)
+                targets = targets.to(device)
+                # FIX: Call get_features instead of model() directly
+                outputs = LinearFiltration.get_features(model, inputs)
+
+                for j in range(num_classes):
+                    if j == forget_index:
+                        continue
+                    mask = (targets == j)
+
+                    if mask.any():
+                        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(
+            model=model,
+            num_classes=num_classes,
+            retain_loader=retain_loader,
+            forget_loader=forget_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))
+        
+        # CORRECT: Stack along dim=0 to make it (num_classes, h_dim)
+        return torch.stack(A, dim=0)
+    
+
+    @staticmethod
+    def _compute_z(tensor, forget_index):
+        # Now tensor has shape (num_classes, h_dim) -> tensor.shape[0] is num_classes
+        K = tensor.shape[0] 
+        
+        # pi_a0 should match the feature space dimensions (h_dim)
+        pi_a0 = torch.zeros(tensor.shape[1], device=tensor.device)
+
+        t_1 = pi_a0
+        a0 = tensor[forget_index, :] # Extracting the row vector for the forgotten class
+
+        mask_a0 = torch.ones(
+            a0.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()
+
+        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
+
+
+    @staticmethod
+    def normalise(model, retain_loader, forget_loader, device, forget_index):
+        W = model.fc.weight.data.clone()
+        num_classes = W.shape[0]
+        
+        A = LinearFiltration._get_means(
+            model=model,
+            num_classes=num_classes,
+            retain_loader=retain_loader,
+            forget_loader=forget_loader,
+            device=device,
+            forget_index=forget_index
+        )
+
+        Z = LinearFiltration._compute_z(tensor=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(A[i])
+        
+        # Stack back along dim=0 to match (num_classes, h_dim)
+        B_Z = torch.stack(B_Z_rows, dim=0)
+
+        A_inv = torch.linalg.pinv(A)
+        
+        W_Z = B_Z @ A_inv @ W
+
+        model.fc.weight.copy_(W_Z)
+
+        return model

unlearning/WeightFiltration.py

@@ -3,97 +3,34 @@ 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 WeightFiltration(Strategy):
     """
-    Implements Poppi et al.'s Weight Filtering framework for linear layers.
+    Verbatim implementation of Poppi et al.'s WF-Net framework.
-    Uses a standard functional hook to guarantee native PyTorch autograd tracking.
+    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,num_classes: int, epochs: int = 10, lr: float = 0.2, gamma: float = 10.0):
+    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)
+        super().__init__(target_class_index=target_class_index)
-        self.num_classes = num_classes
         self.epochs = epochs
         self.lr = lr
         self.gamma = gamma 
-        self.alpha = None
+        #self.alpha = None
-        self.hook_handle = None
 
 
-    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
-        device = next(model.parameters()).device
-        model.eval() 
-        
-        # Locate layer4 for dynamic optimization
-        target_layer = model.layer4 if hasattr(model, 'layer4') else None
-        fc_layer = model.fc if hasattr(model, 'fc') and isinstance(model.fc, nn.Linear) else None
-        
-        if target_layer is None or fc_layer is None:
-            raise AttributeError("Model does not have the required layers.")
-        
-        # Match alpha dimensions to the channels outputted by layer4
-        num_features = fc_layer.weight.shape[1] 
-        self.alpha = nn.Parameter(torch.ones(self.num_classes, num_features, device=device) * 1.5)
-
-        # Freeze everything except our channel mask
-        for p in model.parameters():
-            p.requires_grad = False
-        self.alpha.requires_grad = True
-
-        # Hook into layer4 dynamically to run the untraining optimization
-        self.hook_handle = target_layer.register_forward_hook(self._get_hook())
-        # optimise the filter to maintain accuracy on retain set
-        # and decrease accuracy on forget set
-        self._optimise_filter(model, forget_loader, retain_loader, device)
-
-        # Remove the runtime hook
-        self.hook_handle.remove()
-        
-        # Transfer the channel suppression permanently into model.fc
-        with torch.no_grad():
-            mask = torch.sigmoid(self.alpha[self.target_class_index]) # Shape: (num_features,)
-            
-            # Suppress the channels ONLY for the target class row in fc
-            fc_layer.weight[self.target_class_index].copy_(
-                fc_layer.weight[self.target_class_index] * mask
-            )
-            print(f">> Baked deep channel filter into Class {self.target_class_index} weights.")
-            
-        return model
-
-    def _get_hook(self):
-        """
-        Filters the internal feature map channels of layer4.
-        The mask scales the channels across the batch.
-        """
-        def functional_hook(module, layer_input, layer_output):
-            # layer_output shape: (batch, channels, height, width) -> e.g., (16, 2048, 7, 7)
-            # self.alpha shape: (num_classes, channels) -> e.g., (20, 2048)
-            
-            # Extract 1D mask for the target class: (channels,)
-            mask = torch.sigmoid(self.alpha[self.target_class_index])
-            
-            # Reshape mask to (1, channels, 1, 1) so it broadcasts over batch, height, and width
-            mask = mask.view(1, -1, 1, 1)
-            
-            # Scale the internal feature maps before they move to the next layer
-            return layer_output * mask
-            
-        return functional_hook
 
 
-    def _optimise_filter(self, model, forget_loader, retain_loader, device):
+
-        optimizer = optim.Adam([self.alpha], lr=self.lr)
+    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)
         criterion = nn.CrossEntropyLoss()
 
-        print(f"[{self.__class__.__name__}] Unlearning Class {self.target_class_index} with gamma={self.gamma}...")
-
-        # To optimise this loop we will watch improvements after each optimisation
-        temp_forget_loss = None
-        # this can be adjusted to optimise the best escape point
-        # it is the value we set to evaluate performance improvement after each itteration.
-        # if improvement is less than this, then we break itteration.
-        threshold = 0.05
-
         for epoch in range(self.epochs):
             forget_iter = iter(forget_loader)
             t_loss_r, t_loss_f = 0.0, 0.0
@@ -102,6 +39,7 @@ class WeightFiltration(Strategy):
             for r_inputs, r_labels in retain_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:
@@ -111,10 +49,19 @@ class WeightFiltration(Strategy):
                 
                 optimizer.zero_grad()
                 
-                # Compute Losses
+                # --- APPLY ALGORITHM 1 FORWARD PASS TO BOTH INPUTS ---
-                # The hook handles the weight filtering smoothly behind the scenes
+                # Pass the input batch AND the target unlearn class index
-                loss_r = criterion(model(r_inputs), r_labels)
+                outputs_r = wf_model(r_inputs, target_unlearn_class=self.target_class_index)
-                loss_f = -torch.sum((torch.ones_like(model(f_inputs)) / self.num_classes) * torch.log_softmax(model(f_inputs), dim=-1))
+                outputs_f = wf_model(f_inputs, target_unlearn_class=self.target_class_index)
+                
+                # Compute Losses using Poppi et al.'s temperature scaled entropy
+                loss_r = criterion(outputs_r, r_labels)
+                
+                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)
+                )
                 
                 total_loss = loss_r + (self.gamma * loss_f)
                 total_loss.backward()
@@ -122,17 +69,56 @@ class WeightFiltration(Strategy):
                 
                 t_loss_r += loss_r.item()
                 t_loss_f += loss_f.item()
-
                 steps += 1
-            forget_loss = t_loss_f / steps
-            print(f"  Epoch {epoch+1}/{self.epochs} | Retain Loss: {t_loss_r/steps:.4f} | Forget Loss: {forget_loss:.4f}")
                 
-            if temp_forget_loss is not None:
+            print(f"  Epoch {epoch+1}/{self.epochs} | Retain Loss: {t_loss_r/steps:.4f} | Forget Loss: {t_loss_f/steps:.4f}")
 
-                improvement = temp_forget_loss - forget_loss
+        return wf_model
-                # if optimisation reaches a point of diminishing returns (improvements is less than threshold)
+            
-                # we break the loop
+
-                if improvement < threshold:
+
-                    break
+    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
-            # else we update the lasst recorded loss.
+        device = next(model.parameters()).device
-            temp_forget_loss = forget_loss
+        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
+        else:
+            raise AttributeError("Model architecture does not match expected ResNet-18 structure.")
+            
+        # 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]
+
+        # 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,
+        )
+
+        # --- PERMANENT BAKING STEP ---
+        # Disconnect the dynamic parameter and freeze the optimal gated state permanently into the architecture
+    
+        with torch.no_grad():
+            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
+            for p in model.parameters():
+                p.requires_grad = True
+                
+        print(f">> Permanently altered {out_channels} convolutional filters in layer4 via WF-Net.")
+        return model