facebook's implementation

unlearning/Certified_facebook.py

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