strategies tested

unlearning/CertifiedRemoval.py

@@ -1,77 +1,153 @@
 import torch
-import numpy as np
-from scipy.optimize import minimize
-from .Strategy import Strategy
 import torch.nn as nn
+from torch.utils.data import DataLoader
+from unlearning.Strategy import Strategy
 
 class CertifiedRemoval(Strategy):
-    """Implements Certified Removal for machine unlearning."""
-
-    def __init__(self, model, data, labels, removal_bound, epsilon):
+    """
+    Implements Certified Removal (Guo et al.) adapted for deep architectures
+    like ResNet50 by isolating and updating the final classification layer.
+    """
+    def __init__(self, removal_bound: float, epsilon: float, l2_reg: float = 0.1):
         super().__init__()
-        self.model = model
-        self.data = data
-        self.labels = labels
-        self.removal_bound = removal_bound
-        self.epsilon = epsilon
+        self.removal_bound = removal_bound  # gamma in the paper
+        self.epsilon = epsilon              # Privacy budget
+        self.l2_reg = l2_reg                # Lambda regularization term
 
-    def _run(self, model: nn.Module) -> nn.Module:
-        """Runs the certified removal algorithm."""
-        # 1. Linear Model Creation
-        # This is a simplification for demonstration purposes.  In a real implementation,
-        # you'd use more sophisticated methods to learn the parameters of the
-        # 'removal' model based on the example being removed.
-
-        def linear_model(x):
-            return torch.dot(x, torch.tensor([1, 1])) # Simplified Linear Model
-
-        # 2. Optimization for Parameter Adjustment
-        # Optimize the parameter values to minimize the loss while staying within bounds.
-        original_params = torch.tensor([0.0, 0.0])  # Initial parameters for linear model
-
-        def objective_function(params):
-            new_model = linear_model #use same function as defined above
-            return torch.sum(((new_model(self.data[0]) - self.labels)**2))
+    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 = []
+        all_labels = []
+        
+        with torch.no_grad():
+            for inputs, labels in loader:
+                inputs = inputs.to(device)
+                # Pass through backbone to get the 2048-dimensional feature vector
+                features = backbone(inputs)
+                all_features.append(features.cpu())
+                all_labels.append(labels.cpu())
                 
-        result = minimize(objective_function, original_params, method='L-BFGS-B', bounds=[(-self.removal_bound, self.removal_bound)], options={'maxiter': 100})
+        return torch.cat(all_features, dim=0), torch.cat(all_labels, dim=0)
 
-        if not result.success:
-            print("Warning: Optimization failed!")
-            print(result.message)
-            return model #Return original if optimization fails
+    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
+        """
+        Entry point expected by your Model.unlearn() architecture interface.
+        Applies Certified Removal strictly to the final linear layer (model.fc).
+        """
+        device = next(model.parameters()).device
+        
+        # Isolate the final NN (Fully connected) layer from the model
+        linear_head = model.fc
+        # Temporarily turn the fc layer into a identity pass-through
+        model.fc = nn.Identity() 
+        
+        print(">> Extracting deep features from model backbone...")
+        retain_features, retain_labels = self._get_features(model, retain_loader, device)
+        forget_features, forget_labels = self._get_features(model, forget_loader, device)
+        
+        # Restore the linear head back
+        model.fc = linear_head
+        
+        # Extract weights from the classification layer
+        # w shape: [num_classes, 2048]
+        w = model.fc.weight.data.clone().cpu()
+        
+        # Compute the Exact Hessian Matrix over the remaining (retained) features
+        # Formula: H = (X^T * X) / N + lambda * I
+        # this will be done on CPU. requires more ram so we cant afford to do it on VRAM
+        # 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))
+        hessian = self._compute_hessian(retain_features=retain_features, retain_features_size = N_retain)
+        
+        # Compute the gradient of the loss with respect to the forgotten data
+        # print(">> Calculating forget set gradients...")
+        # num_classes = w.size(0)
+        # Pass features through linear layer weights to get logits
+        # logits_forget = torch.matmul(forget_features, w.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()
+        
+        #preds_forget = torch.matmul(forget_features, w.t())
+        #error = preds_forget - forget_labels_one_hot
+        # error = preds_softmax - forget_labels_one_hot
+        # grad_forget shape: [num_classes, 2048]
+        grad_forget = self._compute_loss_gradient(
+            forget_labels=forget_labels,
+            forget_features=forget_features, 
+            model_weights=w)
+        #torch.matmul(error.t(), forget_features) / forget_features.size(0)
 
-        new_params = result.x
-        # 3. New Model Creation
+        # Compute the Newton step update via solving: H * Delta_W^T = Grad_forget^T
+        delta_w = self._compute_newton_step(
+            tensor = hessian,
+            gradient= grad_forget
+        )
+        # print(">> Solving Newton step via system optimization...")
+        # try:
+        #     delta_w_t = torch.linalg.solve(Hessian, grad_forget.t())
+        #     delta_w = delta_w_t.t()
+        # except RuntimeError:
+        #     print(">> Warning: Hessian matrix is singular. Falling back to pseudo-inverse.")
+        #     delta_w = torch.matmul(grad_forget, torch.linalg.pinv(Hessian).t())
 
-        new_model = lambda x: torch.dot(x, new_params)
-        return new_model
+        # 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
+        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. Norm: {norm_delta:.4f} <= Bound: {self.removal_bound}")
 
-if __name__ == '__main__':
-    # Example Usage - Synthetic Data for Demonstration
-    np.random.seed(42)  # For reproducibility
-    n_samples = 100
-    X = np.random.randn(n_samples, 2)
-    y = (X[:, 0] + X[:, 1] > 0).astype(int)
+        # Push updated parameters back into the model instance in-place
+        model.fc.weight.data = new_w.to(device)
+        
+        print(">> Certified Removal process completed successfully.")
+        return model
+    
 
-    # Create a simple linear model for demonstration
-    model = nn.Linear(2, 1)  # Simple linear classifier - PyTorch Version
-    optimizer = torch.optim.SGD(model.parameters(), lr=0.01) # Optimizer for training the linear 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
+    
 
-    # Train a Linear Model
-    for _ in range(100): #training loop
-        optimizer.zero_grad()
-        predictions = model(X)
-        loss = torch.sum((predictions - y)**2)
-        loss.backward()
-        optimizer.step()
+    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()
+        
 
-    # Define parameters for Certified Removal
-    removal_bound = 1.0
-    epsilon = 0.1
+        error = preds_softmax - forget_labels_one_hot
+        # grad_forget shape: [num_classes, 2048]
+        return torch.matmul(error.t(), forget_features) / forget_features.size(0)
+    
 
-    # Create the CertifiedRemoval object with the trained model, data and labels
-    certified_removal_obj = CertifiedRemoval(model, X, y, removal_bound, epsilon)
-
-    # Run Certified Removal
-    new_model = certified_removal_obj.apply(model)
+    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/LastK_Certified.py

@@ -0,0 +1,125 @@
+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,13 +2,14 @@
 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_idx: int):
-        super().__init__()  # Automatically configures 'NormalizingLinearFiltration_metrics.txt'
+        super().__init__()
         self.target_class_idx = target_class_idx
 
-    def _run(self, model: nn.Module) -> nn.Module:
+    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
         model.eval()
         for param in model.parameters():
             param.requires_grad = False
@@ -20,29 +21,28 @@ class LinearFiltration(Strategy):
             A = self._calculate_filtration_matrix(num_classes, self.target_class_idx, W.device)
             sanitized_W = torch.mm(A, W)
             model.fc.weight.copy_(sanitized_W)
+            # Filter the bias (if the layer uses one)
+            if model.fc.bias is not None:
+                b = model.fc.bias.data.clone()
+                # 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
 
-    '''@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_classes = 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_classes
-        return A'''
     @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 the inputs of all other classes
+        # 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

unlearning/Strategy.py

@@ -2,6 +2,9 @@
 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."""
@@ -9,21 +12,10 @@ class Strategy:
     def __init__(self):
         # Dynamically set file name based on the class name (e.g., 'NormalizingLinearFiltration.txt')
         self.strategy_name = self.__class__.__name__
-        self.log_file = f"reports/{self.strategy_name}/metrics.txt"
-        self._initialize_log_file()
+        self.log_file = Path(f"reports/{self.strategy_name}/metrics.txt")
+        Util._initialize_log_file(log_file= self.log_file)
 
-    def _initialize_log_file(self):
-        """Creates a unique log file for this strategy with a header if it doesn't exist."""
-        if not os.path.exists(self.log_file):
-            with open(self.log_file, "w") as f:
-                f.write("execution_time_sec\n")
-
-    def log_metric(self, execution_time: float):
-        """Appends the execution time to this strategy's specific file."""
-        with open(self.log_file, "a") as f:
-            f.write(f"{execution_time:.6f}\n")
-
-    def apply(self, model: nn.Module) -> nn.Module:
+    def apply(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
         """
         Wraps the unlearning execution with automated timing and strategy-specific logging.
         DO NOT override this method in subclasses. Override _run instead.
@@ -31,17 +23,21 @@ class Strategy:
         start_time = time.perf_counter()
         
         # Execute core unlearning logic
-        processed_model = self._run(model)
+        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
-        self.log_metric(execution_time)
+        Util.log_metric(
+            log_file=self.log_file, 
+            execution_time=execution_time
+        )
+
         print(f"[{self.strategy_name}] Completed in {execution_time:.6f} seconds. Saved to {self.log_file}")
         
         return processed_model
 
-    def _run(self, model: nn.Module) -> nn.Module:
+    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
         """Subclasses implement their core unlearning logic here."""
         raise NotImplementedError

unlearning/WeightFiltration.py

@@ -0,0 +1,140 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from unlearning.Strategy import Strategy
+
+class WeightFiltration(Strategy):
+    """
+    Implements Poppi et al.'s Weight Filtering framework for linear layers.
+    Uses a standard functional hook to guarantee native PyTorch autograd tracking.
+    """
+    def __init__(self, num_classes: int, target_class_idx: int, epochs: int = 10, lr: float = 0.2, gamma: float = 10.0):
+        super().__init__()
+        self.num_classes = num_classes
+        self.target_class_idx = target_class_idx
+        self.epochs = epochs
+        self.lr = lr
+        self.gamma = gamma 
+        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_idx]) # Shape: (num_features,)
+            
+            # Suppress the channels ONLY for the target class row in fc
+            fc_layer.weight[self.target_class_idx].copy_(
+                fc_layer.weight[self.target_class_idx] * mask
+            )
+            print(f">> Baked deep channel filter into Class {self.target_class_idx} 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_idx])
+            
+            # 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)
+        criterion = nn.CrossEntropyLoss()
+
+        print(f"[{self.__class__.__name__}] Unlearning Class {self.target_class_idx} 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
+            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()
+                
+                # Compute Losses
+                # The hook handles the weight filtering smoothly behind the scenes
+                loss_r = criterion(model(r_inputs), r_labels)
+                loss_f = -torch.sum((torch.ones_like(model(f_inputs)) / self.num_classes) * torch.log_softmax(model(f_inputs), dim=-1))
+                
+                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
+            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:
+
+                improvement = temp_forget_loss - forget_loss
+                # if optimisation reaches a point of diminishing returns (improvements is less than threshold)
+                # we break the loop
+                if improvement < threshold:
+                    break
+            # else we update the lasst recorded loss.
+            temp_forget_loss = forget_loss