methods.py

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm
import numpy as np
import models
from typing import Tuple, Dict, List
import logging
from torch.optim.lr_scheduler import (
    CosineAnnealingLR,
    CosineAnnealingWarmRestarts,
    StepLR,
)
from torch.utils.data import Dataset, Subset
from torch.optim.lr_scheduler import _LRScheduler
import os
import torchvision

# from opts import OPT as opt
import pickle

from utils import accuracy
import time
from copy import deepcopy
import itertools
import copy


class UnlearningMethods:
    def __init__(self, model: nn.Module, learning_rate: float = 0.01):
        """
        Initialize unlearning methods

        Args:
            model: Neural network model
            forget_class: Class index to forget
            learning_rate: Learning rate for unlearning
        """
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = model.to(self.device)
        self.learning_rate = learning_rate
        print(f"Using device: {self.device}")

    def _monitor_progress(self, loader: DataLoader, epoch: int) -> None:
        """Monitor prediction distribution during unlearning"""
        self.model.eval()
        with torch.no_grad():
            all_probs = []
            for data, _ in loader:
                data = data.to(self.device)
                output = F.softmax(self.model(data), dim=1)
                all_probs.append(output)

            all_probs = torch.cat(all_probs, dim=0)
            avg_probs = all_probs.mean(dim=0)

            print(f"\nEpoch {epoch+1} prediction distribution:")
            for i, prob in enumerate(avg_probs):
                print(f"Class {i}: {prob.item()*100:.2f}%")
        self.model.train()

    def evaluate(
        self, test_loader: DataLoader, num_classes: int = 100
    ) -> Tuple[float, Dict[int, float], Dict[int, np.ndarray]]:
        """
        Evaluate model performance for each class

        Returns:
            Tuple of (overall_accuracy, class_accuracy, prediction_distribution)
        """
        self.model.eval()
        class_correct = [0] * num_classes
        class_total = [0] * num_classes
        class_predictions = [[] for _ in range(num_classes)]

        with torch.no_grad():
            for data, targets in tqdm(test_loader):
                data, targets = data.to(self.device), targets.to(self.device)
                outputs = self.model(data)
                _, predicted = outputs.max(1)

                for i in range(len(targets)):
                    label = targets[i].item()
                    pred = predicted[i].item()
                    class_total[label] += 1
                    if label == pred:
                        class_correct[label] += 1
                    class_predictions[label].append(pred)

        class_accuracy = {}
        prediction_distribution = {}
        correct = 0
        total = 0
        overall_accuracy = 0
        for i in range(len(class_total)):
            if class_total[i] > 0:
                correct += class_correct[i]
                total += class_total[i]
                accuracy = 100.0 * class_correct[i] / class_total[i]
                predictions = class_predictions[i]
                dist = np.bincount(predictions, minlength=num_classes) / len(
                    predictions
                )
                class_accuracy[i] = accuracy
                prediction_distribution[i] = dist
        overall_accuracy = 100.0 * correct / total
        return overall_accuracy, class_accuracy, prediction_distribution

    def retrain_learning(
        self,
        seed: int,
        forget_class: int,
        train_retain_loader: DataLoader,
        test_forget_loader: DataLoader,
        test_retain_loader: DataLoader,
        num_epochs: int = 5,
        save_path: str = None,
    ) -> None:
        # optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate)
        # Adam optimiser does not work for this task
        # use SGD

        criterion = nn.CrossEntropyLoss()
        optimizer = optim.SGD(
            self.model.parameters(),
            lr=self.learning_rate,
            momentum=0.9,
            weight_decay=5e-4,
            nesterov=True,
        )

        cosine = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)

        best_acc = 0
        # Training progress bar
        for epoch in tqdm(range(num_epochs)):
            self.model.train()
            train_correct = 0
            train_total = 0
            train_loss = 0
            train_pbar = tqdm(
                train_retain_loader, desc=f"Epoch {epoch+1}/{num_epochs} [Train]"
            )
            for data, target in train_pbar:
                data, target = data.to(self.device), target.to(self.device)
                optimizer.zero_grad()
                output = self.model(data)
                loss = criterion(output, target)
                loss.backward()
                optimizer.step()

                # Calculate training accuracy
                _, predicted = torch.max(output.data, 1)
                train_total += target.size(0)
                train_correct += (predicted == target).sum().item()
                train_loss += loss.item()

                # Update progress bar
                train_pbar.set_postfix(
                    {
                        "loss": f"{train_loss/(train_pbar.n+1):.4f}",
                        "acc": f"{100.*train_correct/train_total:.2f}%",
                    }
                )
            cosine.step()

            # Evaluate on test retain set only (gold model is trained only on retain data)
            self.model.eval()
            test_retain_correct = 0
            test_retain_total = 0
            with torch.no_grad():
                for data, target in test_retain_loader:
                    data, target = data.to(self.device), target.to(self.device)
                    output = self.model(data)
                    _, predicted = torch.max(output.data, 1)
                    test_retain_total += target.size(0)
                    test_retain_correct += (predicted == target).sum().item()

            test_retain_acc = 100.0 * test_retain_correct / test_retain_total

            # Print epoch summary
            train_acc = 100.0 * train_correct / train_total
            tqdm.write(f"\nEpoch {epoch+1}/{num_epochs}:")
            tqdm.write(
                f"Train Loss: {train_loss/len(train_retain_loader):.4f}, Train Acc: {train_acc:.2f}%, Test Retain Acc: {test_retain_acc:.2f}%"
            )
            logging.info(
                f"Epoch{epoch+1}'s loss: {train_loss/len(train_retain_loader):.4f}, Train Acc: {train_acc:.2f}%, Test Retain Acc: {test_retain_acc:.2f}%"
            )

            # Save best model
            if test_retain_acc > best_acc:
                best_acc = test_retain_acc
                if not os.path.exists(f"outputs/seed{seed}/models"):
                    os.makedirs(f"outputs/seed{seed}/models")
                torch.save(
                    self.model.state_dict(),
                    f"outputs/seed{seed}/models/cifar10_forget{forget_class}_gold_model_best.pth",
                )
                tqdm.write(
                    f"Saved best model with validation accuracy: {best_acc:.2f}%"
                )
            tqdm.write("-" * 60)

        # Save final model
        if save_path:
            # Use provided save path with _final suffix
            save_dir = os.path.dirname(save_path)
            final_save_path = save_path.replace(".pth", "_final.pth")
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            torch.save(self.model.state_dict(), final_save_path)
        else:
            # Fallback to original behavior
            torch.save(
                self.model.state_dict(),
                f"outputs/seed{seed}/models/cifar10_forget{forget_class}_gold_model_final.pth",
            )
        # return best_acc

    def vanilla_unlearning(
        self, forget_loader: DataLoader, retain_loader: DataLoader, num_epochs: int = 5
    ) -> None:
        """Simple entropy maximization unlearning"""
        self.model.train()
        # optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate)
        # Adam optimiser does not work for this task
        # use SGD
        optimizer = optim.SGD(
            self.model.parameters(), lr=self.learning_rate, momentum=0.9
        )

        for epoch in range(num_epochs):
            total_entropy = 0
            train_loss = 0
            batch_count = 0
            pbar = tqdm(
                forget_loader, desc=f"Vanilla Unlearning Epoch {epoch+1}/{num_epochs}"
            )
            for data, _ in pbar:
                data = data.to(self.device)

                optimizer.zero_grad()
                outputs = self.model(data)

                # Maximize entropy
                probs = F.softmax(outputs, dim=1)
                entropy = -(probs * torch.log(probs + 1e-10)).sum(dim=1).mean()
                loss = -entropy
                loss.backward()
                optimizer.step()

                total_entropy += entropy.item()
                pbar.set_postfix({"entropy": f"{entropy.item():.4f}"})

                train_loss += loss.item()
                batch_count += 1
            avg_loss = train_loss / batch_count
            logging.info(f"Epoch{epoch}'s loss: {avg_loss:.4f}")
            avg_entropy = total_entropy / len(forget_loader)
            print(f"Epoch {epoch+1}/{num_epochs}, Average Entropy: {avg_entropy:.4f}")
            # self._monitor_progress(forget_loader, epoch)

    def finetuning_unlearning(
        self, forget_loader: DataLoader, retain_loader: DataLoader, num_epochs: int = 50
    ) -> None:
        """Simple entropy maximization unlearning"""
        self.model.train()
        # optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate)
        # Adam optimiser does not work for this task
        # use SGD
        optimizer = optim.SGD(
            self.model.parameters(), lr=self.learning_rate, momentum=0.9
        )
        criterion = nn.CrossEntropyLoss()
        for epoch in range(num_epochs):
            train_loss = 0
            batch_count = 0
            pbar = tqdm(retain_loader, desc=f"Fine-tuning Epoch {epoch+1}/{num_epochs}")
            for data, target in pbar:
                data = data.to(self.device)
                target = target.to(self.device)

                optimizer.zero_grad()
                outputs = self.model(data)
                loss = criterion(outputs, target)
                loss.backward()
                optimizer.step()

                train_loss += loss.item()
                batch_count += 1
            avg_loss = train_loss / batch_count
            logging.info(f"Epoch{epoch}'s loss: {avg_loss:.4f}")
            # self._monitor_progress(forget_loader, epoch)

    def ETF_replace(
        self,
        forget_loader: DataLoader,
        retain_loader: DataLoader,
        forget_class: int,
    ) -> None:
        """
        ETF (Equiangular Tight Frame) replacement unlearning method.
        Assumes the model classifier has reached ETF structure and eliminates
        the forget class by applying an orthogonal transformation to the
        vector corresponding to the forgetting class.

        Args:
            forget_loader: DataLoader for forget set
            retain_loader: DataLoader for retain set
            forget_class: Class index to forget
        """
        self.model.train()

        # Get the classifier weights (assuming it's the last layer)
        classifier = None
        for name, module in self.model.named_modules():
            if isinstance(module, nn.Linear) and name.endswith(
                (".fc", ".classifier", ".head")
            ):
                classifier = module
                break

        if classifier is None:
            # Try to find the last linear layer
            for name, module in reversed(list(self.model.named_modules())):
                if isinstance(module, nn.Linear):
                    classifier = module
                    break

        if classifier is None:
            raise ValueError("Could not find classifier layer in the model")

        print(f"Found classifier: {classifier}")
        print(f"Classifier weight shape: {classifier.weight.shape}")

        # Get the current weight matrix
        W = classifier.weight.data.clone()  # Shape: [num_classes, feature_dim]
        num_classes, feature_dim = W.shape

        # Normalize the weight vectors to unit length
        W_norm = F.normalize(W, p=2, dim=1)

        # Get the forget class vector
        forget_vector = W_norm[forget_class]  # Shape: [feature_dim]

        # Create a projection matrix to remove the forget class component
        # from all other class vectors
        def create_projection_matrix(forget_vector, feature_dim):
            """
            Create a projection matrix P = I - vv^T/||v||^2 that projects out
            the component along the forget vector from any vector
            """
            v = F.normalize(forget_vector, p=2, dim=0)
            # P = I - vv^T
            P = torch.eye(feature_dim, device=self.device) - torch.outer(v, v)
            return P

        # Create projection matrix
        P = create_projection_matrix(forget_vector, feature_dim)

        # Apply projection to all other class vectors to remove forget class component
        new_W = W.clone()
        for i in range(num_classes):
            if i != forget_class:
                # Project out the forget class component from this class vector
                projected_vector = torch.matmul(P, W[i])
                new_W[i] = projected_vector

        # Set forget class vector to zero
        new_W[forget_class] = torch.zeros_like(W[forget_class])

        # Update the classifier weights
        with torch.no_grad():
            classifier.weight.data.copy_(new_W)

            # Also zero out the bias for the forget class if it exists
            if classifier.bias is not None:
                classifier.bias.data[forget_class] = 0.0

        print(f"Applied projection-based transformation to forget class {forget_class}")
        print(
            f"Original forget vector norm: {torch.norm(W_norm[forget_class], p=2):.4f}"
        )
        print(f"New forget vector norm: {torch.norm(new_W[forget_class], p=2):.4f}")

        # Additional analysis: Check if the model still has information about forget class
        print(
            f"\nForget class vector is now zero: {torch.allclose(new_W[forget_class], torch.zeros_like(new_W[forget_class]))}"
        )
        if classifier.bias is not None:
            print(
                f"Forget class bias is now zero: {classifier.bias.data[forget_class] == 0.0}"
            )

        # Verify that other class vectors are orthogonal to forget class vector
        new_W_norm = F.normalize(new_W, p=2, dim=1)
        print("\nVerifying orthogonality after projection:")
        for i in range(num_classes):
            if i != forget_class:
                dot_product = torch.dot(new_W_norm[forget_class], new_W_norm[i])
                print(f"Dot product with class {i}: {dot_product:.6f}")

        # Check the rank of the transformed weight matrix
        rank = torch.linalg.matrix_rank(new_W)
        print(
            f"\nRank of transformed weight matrix: {rank} (should be {num_classes-1})"
        )

        self.model.eval()

    def _find_classifier(self) -> nn.Linear:
        # Try common names first
        classifier = None
        for name, module in self.model.named_modules():
            if isinstance(module, nn.Linear) and name.endswith(
                (".fc", ".classifier", ".head")
            ):
                classifier = module
                break
        if classifier is None:
            # Fallback: last Linear in the network
            for name, module in reversed(list(self.model.named_modules())):
                if isinstance(module, nn.Linear):
                    classifier = module
                    break
        if classifier is None:
            raise ValueError("Could not find classifier (nn.Linear) in the model.")
        return classifier

    @torch.no_grad()
    def _build_etf_projection(self, w_forget: torch.Tensor) -> torch.Tensor:
        """
        Make P = I - v v^T where v = normalized forget vector (dim=D).
        """
        v = F.normalize(w_forget, p=2, dim=0)
        D = v.numel()
        P = torch.eye(D, device=self.device) - torch.outer(v, v)
        return P

    def _capture_features(
        self, model: nn.Module, x: torch.Tensor, classifier: nn.Linear, detach: bool
    ) -> torch.Tensor:
        """
        Capture the input to the classifier (penultimate features) via a forward hook.
        If detach=True, returns a detached tensor (for teacher/targets).
        """
        feats = {}

        def hook(m, inputs, output):
            # inputs is a tuple; the first item is the feature tensor [B, D]
            feats["f"] = inputs[0] if not detach else inputs[0].detach()

        handle = classifier.register_forward_hook(hook)
        _ = model(x)  # forward pass triggers the hook
        handle.remove()

        if "f" not in feats:
            raise RuntimeError("Failed to capture features. Hook did not run.")
        return feats["f"]

    def _compute_etf_loss(
        self, student_features: torch.Tensor, targets: torch.Tensor, loss_function: str
    ) -> torch.Tensor:
        """
        Compute loss between student features and ETF-projected targets.

        Args:
            student_features: features from student model [B, D]
            targets: ETF-projected target features [B, D]
            loss_function: type of loss to compute

        Returns:
            loss tensor
        """
        if loss_function == "cosine":
            # Original cosine similarity loss (can be unstable)
            f_s_n = F.normalize(student_features, p=2, dim=1)
            targets_n = F.normalize(targets, p=2, dim=1)
            cos_sim = (f_s_n * targets_n).sum(dim=1)
            return 1.0 - cos_sim.mean()

        elif loss_function == "mse_normalized":
            # MSE between normalized vectors (more stable)
            f_s_n = F.normalize(student_features, p=2, dim=1)
            targets_n = F.normalize(targets, p=2, dim=1)
            return F.mse_loss(f_s_n, targets_n)

        elif loss_function == "l2_regression":
            # Direct L2 regression without normalization
            return F.mse_loss(student_features, targets)

        elif loss_function == "huber":
            # Huber (smooth L1) loss - robust against outliers
            return F.smooth_l1_loss(student_features, targets)

        elif loss_function == "angular":
            # Angular distance loss (arccos of cosine similarity)
            f_s_n = F.normalize(student_features, p=2, dim=1)
            targets_n = F.normalize(targets, p=2, dim=1)
            cos_sim = (f_s_n * targets_n).sum(dim=1)
            # Clamp to avoid numerical issues with arccos
            cos_sim = torch.clamp(cos_sim, -1.0 + 1e-7, 1.0 - 1e-7)
            return torch.acos(cos_sim).mean()

        else:
            raise ValueError(
                f"Unknown loss function: {loss_function}. "
                f"Supported options: cosine, mse_normalized, l2_regression, huber, angular"
            )

    # ---------- main method ----------
    def etf_finetune(
        self,
        retain_loader: DataLoader = None,
        forget_class: int = 0,
        num_epochs: int = 20,
        freeze_classifier: bool = True,
        freeze_bn_stats: bool = True,
        log_every: int = 50,
        loss_function: str = "l2_regression",
        forget_loader: DataLoader = None,
    ) -> None:
        """
        Fine-tune the feature extractor so its features match ETF-projected features.

        Steps:
        1) Build projection P = I - v v^T using v = normalized classifier weight of the forget class.
        2) Create a frozen teacher model T that produces target features f_T(x); targets = P f_T(x).
        3) Optimize original model's backbone to match its features to targets using specified loss.

        Args:
            retain_loader: deprecated; not used in this method.
            forget_loader: forget data loader to use for fine-tuning (required)
            forget_class: index of the class to forget (used to form the projection)
            num_epochs: fine-tuning epochs
            freeze_classifier: if True, do not update classifier weights/bias
            freeze_bn_stats: if True, run teacher in eval() and also set BN in student to eval() during training
            log_every: steps per logging
            loss_function: loss function to use for feature alignment. Options:
                - "cosine": 1 - cosine_similarity (original, can be unstable)
                - "mse_normalized": MSE between normalized vectors (more stable)
                - "l2_regression": direct L2 loss without normalization
                - "huber": smooth L1 loss (robust against outliers)
                - "angular": arccos of cosine similarity (smoother than cosine)
        """
        self.model.train()

        # 1) locate classifier and build projection
        classifier = self._find_classifier()
        with torch.no_grad():
            W = classifier.weight.data.clone()  # [C, D]
            v_forget = F.normalize(W[forget_class], p=2, dim=0)  # [D]
        P = self._build_etf_projection(v_forget)  # [D, D]

        # 2) frozen teacher model (same weights now)
        teacher = copy.deepcopy(self.model).to(self.device)
        for p in teacher.parameters():
            p.requires_grad_(False)
        teacher.eval()

        # 3) prepare optimizer over *feature extractor only*
        # freeze classifier if requested
        if freeze_classifier:
            for p in classifier.parameters():
                p.requires_grad_(False)

        # optionally freeze BN running-stat updates in the student
        if freeze_bn_stats:
            # Put student in train mode but freeze BN stats by forcing BN layers to eval
            self.model.train()
            for m in self.model.modules():
                if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
                    m.eval()  # no running stat updates

        # pick params that require grad
        params = [p for p in self.model.parameters() if p.requires_grad]
        if len(params) == 0:
            raise ValueError("No trainable parameters selected for fine-tuning.")

        optimizer = optim.SGD(
            params,
            lr=self.learning_rate,
            momentum=0.9,
            weight_decay=5e-4,
            nesterov=True,
        )
        scheduler = CosineAnnealingLR(optimizer, T_max=num_epochs)

        step = 0
        # strictly use forget_loader; retain data must not be used here
        actual_loader = forget_loader
        if actual_loader is None:
            raise ValueError(
                "etf_finetune requires forget_loader (retain data is not used)."
            )
        for epoch in range(num_epochs):
            pbar = tqdm(
                actual_loader, desc=f"ETF-Finetune Epoch {epoch+1}/{num_epochs}"
            )
            running_loss, running_sim = 0.0, 0.0
            for i, (x, _) in enumerate(pbar):
                x = x.to(self.device)

                # teacher targets: f_T(x) -> P f_T(x)
                with torch.no_grad():
                    teacher_cls = None
                    # find teacher's classifier (same structure)
                    for name, module in teacher.named_modules():
                        if isinstance(module, nn.Linear) and name.endswith(
                            (".fc", ".classifier", ".head")
                        ):
                            teacher_cls = module
                            break
                    if teacher_cls is None:
                        for name, module in reversed(list(teacher.named_modules())):
                            if isinstance(module, nn.Linear):
                                teacher_cls = module
                                break
                    f_t = self._capture_features(
                        teacher, x, teacher_cls, detach=True
                    )  # [B, D]
                    targets = f_t @ P.T  # [B, D]

                # student features (need grads)
                f_s = self._capture_features(
                    self.model, x, classifier, detach=False
                )  # [B, D]

                # compute loss using specified loss function
                loss = self._compute_etf_loss(f_s, targets, loss_function)

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                running_loss += loss.item()

                # compute cosine similarity for logging (regardless of loss function)
                with torch.no_grad():
                    f_s_n = F.normalize(f_s, p=2, dim=1)
                    targets_n = F.normalize(targets, p=2, dim=1)
                    cos_sim = (f_s_n * targets_n).sum(dim=1).mean().item()
                    running_sim += cos_sim

                step += 1

                if (i + 1) % log_every == 0:
                    avg_loss = running_loss / log_every
                    avg_sim = running_sim / log_every
                    pbar.set_postfix(
                        {"loss": f"{avg_loss:.4f}", "cos_sim": f"{avg_sim:.4f}"}
                    )
                    logging.info(
                        f"[ETF-Finetune] epoch {epoch+1} step {i+1}: loss={avg_loss:.4f} ({loss_function}), cos_sim={avg_sim:.4f}"
                    )
                    running_loss, running_sim = 0.0, 0.0

            scheduler.step()

        # restore classifier grad flags (optional—only if you want to resume normal training later)
        if freeze_classifier:
            for p in classifier.parameters():
                p.requires_grad_(True)

    def random_label_unlearning(
        self,
        forget_loader: DataLoader,
        retain_loader: DataLoader,
        num_epochs: int = 1,
        noise_scale: float = 0.1,
    ) -> None:
        """Unlearning by training on randomly labeled forget set data"""
        self.model.train()
        optimizer = optim.SGD(
            self.model.parameters(), lr=self.learning_rate, momentum=0.9
        )
        criterion = nn.CrossEntropyLoss()

        for epoch in range(num_epochs):
            train_loss = 0
            batch_count = 0

            # Train on forget set with random labels
            pbar = tqdm(
                forget_loader,
                desc=f"Random Label Unlearning Epoch {epoch+1}/{num_epochs}",
            )
            for data, _ in pbar:
                data = data.to(self.device)
                batch_size = data.size(0)

                # Generate random labels for forget set (0-8 for PathMNIST)
                random_labels = torch.randint(0, 9, (batch_size,), device=self.device)
                # random number 2 or 8
                # random_labels = torch.randint(2, 8, (batch_size,), device=self.device)

                optimizer.zero_grad()
                # noisy_data = data + torch.randn_like(data) * noise_scale
                outputs = self.model(data)
                loss = criterion(outputs, random_labels)
                loss.backward()
                optimizer.step()

                train_loss += loss.item()
                batch_count += 1

                # Update progress bar
                pbar.set_postfix({"loss": f"{loss.item():.4f}"})

            """
            # Train on retain set with correct labels
            for data, labels in retain_loader:
                data, labels = data.to(self.device), labels.to(self.device)
                
                optimizer.zero_grad()
                outputs = self.model(data)
                loss = criterion(outputs, labels)
                loss.backward()
                optimizer.step()
                
                train_loss += loss.item()
                batch_count += 1
            """

            # Monitor progress
            avg_loss = train_loss / batch_count
            self._monitor_progress(retain_loader, epoch)

            tqdm.write(f"\nEpoch {epoch+1}/{num_epochs}:")
            tqdm.write(f"Average Loss: {avg_loss:.4f}")
            tqdm.write("-" * 60)

    def upper_bound_unlearning(
        self,
        train_forget_loader: DataLoader,
        train_retain_loader: DataLoader,
        test_forget_loader: DataLoader,
        test_retain_loader: DataLoader,
        num_epochs: int = 5,
    ) -> None:
        """
        Upper bound unlearning method that trains a 9-class model on 9 classes.
        This is the corrected version of the gold method that fixes the output dimension mismatch.
        """
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.SGD(
            self.model.parameters(),
            lr=self.learning_rate,
            momentum=0.9,
            weight_decay=5e-4,
            nesterov=True,
        )

        cosine = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)

        best_acc = 0
        # Training progress bar
        for epoch in tqdm(range(num_epochs)):
            self.model.train()
            train_correct = 0
            train_total = 0
            train_loss = 0
            train_pbar = tqdm(
                train_retain_loader,
                desc=f"Epoch {epoch+1}/{num_epochs} [Upper Bound Train]",
            )
            for data, target in train_pbar:
                data, target = data.to(self.device), target.to(self.device)
                optimizer.zero_grad()
                output = self.model(data)
                loss = criterion(output, target)
                loss.backward()
                optimizer.step()

                # Calculate training accuracy
                _, predicted = torch.max(output.data, 1)
                train_total += target.size(0)
                train_correct += (predicted == target).sum().item()
                train_loss += loss.item()

                # Update progress bar
                train_pbar.set_postfix(
                    {
                        "loss": f"{train_loss/(train_pbar.n+1):.4f}",
                        "acc": f"{100.*train_correct/train_total:.2f}%",
                    }
                )
            cosine.step()

            # Evaluate on test set
            self.model.eval()
            test_retain_correct = 0
            test_retain_total = 0
            with torch.no_grad():
                for data, target in test_retain_loader:
                    data, target = data.to(self.device), target.to(self.device)
                    output = self.model(data)
                    _, predicted = torch.max(output.data, 1)
                    test_retain_total += target.size(0)
                    test_retain_correct += (predicted == target).sum().item()

            test_retain_acc = 100.0 * test_retain_correct / test_retain_total

            # Print epoch summary
            train_acc = 100.0 * train_correct / train_total
            tqdm.write(f"\nEpoch {epoch+1}/{num_epochs}:")
            tqdm.write(
                f"Train Loss: {train_loss/len(train_retain_loader):.4f}, Train Acc: {train_acc:.2f}%, Test Retain Acc: {test_retain_acc:.2f}%"
            )
            logging.info(
                f"Epoch{epoch+1}'s loss: {train_loss/len(train_retain_loader):.4f}, Train Acc: {train_acc:.2f}%, Test Retain Acc: {test_retain_acc:.2f}%"
            )
            tqdm.write("-" * 60)

        tqdm.write("Upper bound unlearning completed successfully!")
        logging.info("Upper bound unlearning completed successfully!")

    def boundary_shrink_unlearning(
        self,
        forget_loader: DataLoader,
        num_classes: int,
        epsilon: float = 5,  # step size for crossing boundary (image inputs in [0,1])
        epochs: int = 5,  # few epochs typically suffice
        lr: float = 1e-3,  # fine-tune LR (head-only by default)
        head_only: bool = True,  # train only classifier head to avoid collateral drift
        clamp_min: float = 0.0,  # clamp inputs after FGSM
        clamp_max: float = 1.0,
        momentum: float = 0.9,
        weight_decay: float = 0.0,
        adv_batch_size: int = None,  # optional smaller batch for FGSM gen if memory tight
    ) -> None:
        """
        Boundary Shrink unlearning:
          1) For each forget sample x, take an FGSM step across the boundary.
          2) Relabel x to the model's predicted (incorrect) class of the adversarial point.
          3) Fine-tune on (x, new_label) to shrink the forget region.

        Args:
            forget_loader: DataLoader over forget set D_f (with true labels, but we'll relabel)
            num_classes: total number of classes
            epsilon: FGSM magnitude to cross the boundary
            epochs: #epochs to fine-tune on relabeled forget set
            lr: learning rate for fine-tune
            head_only: if True, only optimize the final linear classifier
            clamp_min/max: clamp bounds for perturbed inputs
            adv_batch_size: if set, regenerates a loader with this batch size for FGSM step
        """
        self.model.eval()

        # ---- 1) Generate nearest-boundary adversarial points and new labels ----
        criterion = nn.CrossEntropyLoss(reduction="mean")

        # Optionally rebuild a loader for FGSM to control memory
        gen_loader = forget_loader
        if adv_batch_size is not None:
            gen_loader = DataLoader(
                forget_loader.dataset, batch_size=adv_batch_size, shuffle=False
            )

        new_images = []
        new_targets = []

        for data, y_true in tqdm(gen_loader, desc="Generating cross-boundary labels"):
            data = data.to(self.device)
            y_true = y_true.to(self.device)

            # require grad on inputs for FGSM
            data_adv = data.clone().detach().requires_grad_(True)

            logits = self.model(data_adv)
            loss = criterion(logits, y_true)
            self.model.zero_grad(set_to_none=True)
            loss.backward()

            # FGSM step across the boundary
            x_adv = data_adv + epsilon * data_adv.grad.sign()
            x_adv = torch.clamp(x_adv.detach(), clamp_min, clamp_max)

            # predict label of adversarial sample (nearest-but-incorrect)
            with torch.no_grad():
                pred_adv = self.model(x_adv).argmax(dim=1)

            # We keep the ORIGINAL images but use the adversarial label (per paper heuristic)
            new_images.append(data.detach().cpu())
            new_targets.append(pred_adv.detach().cpu())

        X = torch.cat(new_images, dim=0)
        Y = torch.cat(new_targets, dim=0)
        D_shrink = torch.utils.data.TensorDataset(X, Y)
        shrink_loader = DataLoader(
            D_shrink, batch_size=forget_loader.batch_size, shuffle=True
        )

        # ---- 2) Fine-tune on relabeled forget set to shrink decision region ----
        # Find classifier head if training head-only
        params = list(self.model.parameters())
        train_params = params
        if head_only:
            # try to guess the final linear layer
            classifier = None
            last_linear = None
            for name, m in self.model.named_modules():
                if isinstance(m, nn.Linear):
                    last_linear = m
            classifier = last_linear
            if classifier is None:
                raise ValueError(
                    "BoundaryShrink: couldn't find a final nn.Linear layer for head_only=True."
                )
            # freeze all except head
            for p in self.model.parameters():
                p.requires_grad = False
            for p in classifier.parameters():
                p.requires_grad = True
            train_params = [p for p in classifier.parameters() if p.requires_grad]

        self.model.train()
        optimizer = optim.SGD(
            train_params, lr=lr, momentum=momentum, weight_decay=weight_decay
        )

        for ep in range(epochs):
            ep_loss = 0.0
            n_batches = 0
            pbar = tqdm(
                shrink_loader, desc=f"Boundary Shrink fine-tune {ep+1}/{epochs}"
            )
            for x, y in pbar:
                x = x.to(self.device)
                y = y.to(self.device)

                optimizer.zero_grad(set_to_none=True)
                out = self.model(x)
                loss = criterion(out, y)
                loss.backward()
                optimizer.step()

                ep_loss += loss.item()
                n_batches += 1
                pbar.set_postfix({"loss": f"{ep_loss/max(n_batches,1):.4f}"})

        # unfreeze if we froze
        if head_only:
            for p in self.model.parameters():
                p.requires_grad = True

        self.model.eval()

    def delete_unlearning(
        self,
        forget_loader: DataLoader,
        num_classes: int,
        epochs: int = 5,
        lr: float = 1e-3,
        fix_head: bool = True,
        temperature: float = 1.0,  # distillation temperature
        momentum: float = 0.9,
        weight_decay: float = 0.0,
    ) -> None:
        """
        DELETE: Decoupled Distillation To Erase (mask-distillation on teacher logits).
        - Freeze a copy of the original model as teacher.
        - For each batch from the FORGET set, mask forget-class logits on the teacher,
          renormalize (softmax), and distill to the current model.
        - Optimizes forgetting + retention simultaneously, using only D_f.

        Args:
            forget_loader: DataLoader over D_f (contains (x, y) with y the class to forget for each sample)
            num_classes: total number of classes
            epochs: fine-tuning epochs on D_f
            lr: learning rate
            fix_head: freeze the final linear layer if True
            temperature: distillation temperature
        """
        # 1) Make a frozen teacher (original model snapshot at start)
        import copy

        teacher = copy.deepcopy(self.model).to(self.device)
        for p in teacher.parameters():
            p.requires_grad = False
        teacher.eval()

        # 2) Optionally fix the head (freeze final linear layer)
        train_params = list(self.model.parameters())
        classifier = None
        if fix_head:
            last_linear = None
            for name, m in self.model.named_modules():
                if isinstance(m, nn.Linear):
                    last_linear = m
            classifier = last_linear
            if classifier is None:
                raise ValueError(
                    "DELETE: couldn't find a final nn.Linear for fix_head=True."
                )
            # Freeze only the head (final linear layer)
            for p in classifier.parameters():
                p.requires_grad = False
            # Keep all other parameters trainable
            train_params = [p for p in self.model.parameters() if p.requires_grad]

        optimizer = optim.SGD(
            train_params, lr=lr, momentum=momentum, weight_decay=weight_decay
        )
        # KLDivLoss expects log-probs as input; we'll use log_softmax(student/T) vs soft teacher probs.
        kldiv = nn.KLDivLoss(reduction="batchmean")

        self.model.train()
        T = temperature

        for ep in range(epochs):
            running = 0.0
            n = 0
            pbar = tqdm(forget_loader, desc=f"DELETE (mask-distill) {ep+1}/{epochs}")
            for x, y_forget in pbar:
                x = x.to(self.device)
                y_forget = y_forget.to(self.device)  # per-sample forget class index

                # --- Teacher forward + mask ---
                with torch.no_grad():
                    t_logits = teacher(x) / T  # (B, K)
                    # build mask that zeroes all non-forget entries? Paper sets forget logit -> -inf before softmax
                    # Here we set forget logit to -inf so its prob becomes 0, then softmax & renorm on others.
                    mask = torch.zeros_like(t_logits)
                    mask.scatter_(1, y_forget.view(-1, 1), float("-inf"))
                    masked_logits = t_logits + mask
                    soft_targets = F.softmax(
                        masked_logits, dim=1
                    )  # (B, K), prob(forget)=0, others renormalized

                # --- Student forward ---
                s_logits = self.model(x) / T
                loss = kldiv(F.log_softmax(s_logits, dim=1), soft_targets) * (
                    T * T
                )  # standard T^2 scaling

                optimizer.zero_grad(set_to_none=True)
                loss.backward()
                optimizer.step()

                n += 1
                running += loss.item()
                pbar.set_postfix({"loss": f"{running/max(n,1):.4f}"})

        # unfreeze head if it was frozen
        if fix_head:
            for p in classifier.parameters():
                p.requires_grad = True

        self.model.eval()

    @torch.no_grad()
    def _extract_features_and_probs(self, loader: DataLoader):
        """Return features Z (N,d), probs P (N,C), labels y (N,) from penultimate layer."""
        self.model.eval()
        # Find last Linear (classifier) and a feature hook
        last_linear = None
        for m in self.model.modules():
            if isinstance(m, nn.Linear):
                last_linear = m
        if last_linear is None:
            raise ValueError("No final nn.Linear head found.")

        feats_list, probs_list, y_list = [], [], []
        handle = None

        # small helper to grab penultimate activations
        penult = {}

        def hook(module, inp, out):
            penult["z"] = inp[0].detach()

        handle = last_linear.register_forward_hook(hook)

        for x, y in loader:
            x = x.to(self.device)
            logits = self.model(x)  # hook fills penult['z']
            z = penult["z"].detach()  # [B,d]
            p = F.softmax(logits, dim=1).detach()
            feats_list.append(z.cpu())
            probs_list.append(p.cpu())
            y_list.append(y.clone().cpu())

        handle.remove()
        Z = torch.cat(feats_list, 0)  # (N,d)
        P = torch.cat(probs_list, 0)  # (N,C)
        y = torch.cat(y_list, 0)  # (N,)
        return Z, P, y

    def _conj_grad(self, Avp, b, iters=100, tol=1e-6):
        """Conjugate gradient to solve A x = b given a function Avp(v)."""
        x = torch.zeros_like(b)
        r = b.clone()
        p = r.clone()
        rr = torch.dot(r, r)
        for _ in range(iters):
            Ap = Avp(p)
            alpha = rr / (torch.dot(p, Ap) + 1e-12)
            x += alpha * p
            r -= alpha * Ap
            rr_new = torch.dot(r, r)
            if torch.sqrt(rr_new) < tol:
                break
            beta = rr_new / (rr + 1e-12)
            p = r + beta * p
            rr = rr_new
        return x

    def approx_data_deletion_head(
        self,
        forget_loader: DataLoader,
        hessian_loader: DataLoader,
        l2: float = 1e-4,
        cg_iters: int = 200,
        cg_tol: float = 1e-5,
    ):
        """
        Izzo-style approximate deletion adapted to deep nets:
        treat the FINAL LINEAR HEAD as a multi-class logistic model.
        Compute Δw ≈ H^{-1} * g_forget and update w <- w - Δw (head-only).
        Args:
            forget_loader: data to delete (D_f)
            hessian_loader: small retain subset to estimate H (D_r_sub)  (few thousand points suffice)
            l2: L2 regularization on head (ridge term)
        """
        self.model.eval()

        # locate final head (W,b)
        head = None
        for m in self.model.modules():
            if isinstance(m, nn.Linear):
                head = m
        if head is None:
            raise ValueError("No final nn.Linear head found.")
        C = head.out_features
        d = head.in_features

        # 1) Gather features & probs for Hessian estimate (on retain subset)
        Z_H, P_H, _ = self._extract_features_and_probs(
            hessian_loader
        )  # (N_h,d), (N_h,C)
        Z_H = Z_H.to(self.device)
        P_H = P_H.to(self.device)

        # 2) Gather features & probs for gradient over FORGET set
        Z_F, P_F, y_F = self._extract_features_and_probs(
            forget_loader
        )  # (N_f,d), (N_f,C), (N_f,)
        Z_F = Z_F.to(self.device)
        P_F = P_F.to(self.device)
        y_F = y_F.to(self.device)

        # one-hot labels for forget set
        YF = torch.zeros((y_F.numel(), C), device=self.device)
        YF.scatter_(1, y_F.view(-1, 1), 1.0)

        # Flatten head params: θ = [W(:); b]
        def pack(W, b):
            return torch.cat([W.reshape(-1), b])

        def unpack(theta):
            W = theta[: d * C].view(C, d)
            b = theta[d * C :]
            return W, b

        W0 = head.weight.detach().clone()
        b0 = (
            head.bias.detach().clone()
            if head.bias is not None
            else torch.zeros(C, device=self.device)
        )
        theta0 = pack(W0, b0)

        # 3) Gradient wrt θ over forget set: g_f = ∑ (p - y) kron z  (multi-class logistic)
        # Compute class-wise residuals R = (P_F - YF)  [N_f, C]
        R_F = P_F - YF  # N_f x C
        # Gradients:
        # dL/dW_c = ∑_i R_F[i,c] * z_i
        G_W = R_F.T @ Z_F  # C x d
        G_b = R_F.sum(dim=0)  # C
        g_forget = pack(G_W, G_b)  # (C*d + C,)

        # 4) Define Hessian-vector product on θ using retain subset
        # For multi-class softmax logistic with L2, H(θ) ≈ ∑_i (J_i^T S_i J_i) + l2 I
        # We implement Avp via a forward pass + vector-Jacobian products.
        def Avp(v_flat: torch.Tensor) -> torch.Tensor:
            Wv, bv = unpack(v_flat)

            # logits on retain subset: L = Z_H @ W^T + b
            logits = Z_H @ W0.T + b0
            P = F.softmax(logits, dim=1)  # N_h x C

            # For each sample i, S_i = diag(p_i) - p_i p_i^T
            # Efficient Avp without forming full S: use vector trick for multinomial LR.
            # We need: H_v_W and H_v_b matching shapes.
            # Define U = Z_H @ Wv.T + bv   -> N_h x C  (directional logits)
            U = Z_H @ Wv.T + bv  # N_h x C

            # M = (S applied to U): M = U - (U * P).sum(1, keepdim=True)  (because Sx = x - <x,1> p  with softmax geometry)
            # Correct formula for multinomial: S_i U_i = (U_i - <U_i, 1> * P_i) * P_i  ?  Let's use standard: J^T S J ≈ Z^T R Z blockwise.
            # Practical, stable approximation: class-wise independent blocks using diagonal of S:
            # R_diag = P*(1-P)  (C-wise), ignoring off-diagonals. Works well in practice and mirrors repo's low-cost aim.
            Rdiag = P * (1 - P)  # N_h x C

            # Compute H_v under diagonal approximation:
            # For each class c:
            #   H_v_W[c] = ∑_i Rdiag[i,c] * z_i z_i^T * Wv[c]   + l2 * Wv[c]
            #             + ∑_i Rdiag[i,c] * z_i * bv[c]
            # But the cross W-b blocks give: H_v_b[c] = ∑_i Rdiag[i,c] * (z_i^T Wv[c] + bv[c]) + l2 * bv[c] (on bias)
            # Implement efficiently:
            # Precompute A_c = (Rdiag[:,c][:,None] * Z_H)  (N_h,d)
            H_W = torch.zeros_like(Wv)
            H_b = torch.zeros_like(bv)
            for c in range(P.shape[1]):
                rc = Rdiag[:, c].unsqueeze(1)  # N_h x 1
                A = rc * Z_H  # N_h x d
                # term for Wv[c]:
                H_W[c] = A.t() @ (Z_H @ Wv[c].unsqueeze(1)).squeeze(1)
                # cross with bv[c]:
                H_W[c] += A.sum(0) * bv[c]
                # bias block:
                H_b[c] = ((Z_H @ Wv[c].unsqueeze(1)).squeeze(1) * Rdiag[:, c]).sum() + (
                    rc.squeeze(1) * bv[c]
                ).sum()

            # add L2
            H_W += l2 * Wv
            H_b += l2 * bv
            return pack(H_W, H_b)

        # 5) Solve H Δθ = g_forget   (then θ_new = θ0 - Δθ)
        b_vec = g_forget.detach()
        # move to device & float
        b_vec = b_vec.to(self.device).float()

        def Avp_flat(v_flat):
            return Avp(v_flat.float())

        delta = self._conj_grad(Avp_flat, b_vec, iters=cg_iters, tol=cg_tol)

        theta_new = theta0 - delta
        W_new, b_new = unpack(theta_new)

        # 6) Update head weights
        with torch.no_grad():
            head.weight.copy_(W_new)
            if head.bias is not None:
                head.bias.copy_(b_new)

        self.model.eval()
        print(
            "[Approx-Deletion] Head updated via H^{-1} g_forget (diagonal-softmax Hessian approx)."
        )

    def feature_injection_score(
        self,
        loader: DataLoader,
        trigger: torch.Tensor,
        trigger_scale: float = 1.0,
    ) -> float:
        """
        Izzo et al. 'feature-injection' test (simplified):
        Add a fixed 'trigger' direction to inputs and measure how much predictions align with
        a target shift. Returns average change in predicted logit margin.
        Args:
            trigger: tensor broadcastable to inputs (e.g., small patch or noise vector)
        """
        self.model.eval()
        deltas = []
        for x, _ in loader:
            x = x.to(self.device)
            x2 = torch.clamp(x + trigger_scale * trigger.to(self.device), 0, 1)
            with torch.no_grad():
                m1 = self.model(x)
                m2 = self.model(x2)
            # logit-norm change as a proxy
            deltas.append((m2 - m1).norm(dim=1).mean().item())
        return float(np.mean(deltas))

    # ====== Instance-wise Unlearning (Cha et al., AAAI'24) ======
    def generate_adv_pool_for_unlearning(
        self,
        forget_loader: DataLoader,
        num_classes: int,
        n_per_instance: int = 20,  # paper: 20 for CIFAR-10; 200 for larger sets
        eps: float = 8 / 255,  # l_inf budget; tune per dataset
        alpha: float = 0.1,  # step size (paper uses 1e-1)
        steps: int = 100,  # paper used 100 PGD steps
        clamp_min: float = 0.0,
        clamp_max: float = 1.0,
        targeted: bool = True,
        rng: np.random.RandomState = None,
    ):
        """
        Build an adversarial example pool from the forget set with random targeted labels.
        Returns a TensorDataset (adv_images, adv_targets).
        """
        self.model.eval()
        criterion = nn.CrossEntropyLoss(reduction="sum")
        if rng is None:
            rng = np.random.RandomState(0)

        adv_imgs, adv_tgts = [], []
        for x, y_true in tqdm(forget_loader, desc="Building adversarial pool"):
            x = x.to(self.device)
            y_true = y_true.to(self.device)

            # sample random targets != y_true for each replication
            B = x.size(0)
            for _ in range(n_per_instance):
                rand_targets = []
                for yi in y_true.tolist():
                    t = rng.randint(0, num_classes - 1)
                    if t >= yi:
                        t += 1
                    rand_targets.append(t)
                y_t = torch.tensor(rand_targets, device=self.device, dtype=torch.long)

                x_adv = x.clone().detach()
                x_adv.requires_grad_(True)

                for _ in range(steps):
                    logits = self.model(x_adv)
                    # targeted attack: minimize CE(logits, y_t)
                    loss = criterion(logits, y_t)
                    self.model.zero_grad(set_to_none=True)
                    loss.backward()
                    with torch.no_grad():
                        x_adv = (
                            x_adv - alpha * x_adv.grad.sign()
                        )  # minus -> towards target
                        # project to l_inf ball around original x
                        x_adv = torch.max(torch.min(x_adv, x + eps), x - eps)
                        x_adv = torch.clamp(x_adv, clamp_min, clamp_max)
                        x_adv.requires_grad_(True)

                adv_imgs.append(x_adv.detach().cpu())
                adv_tgts.append(y_t.detach().cpu())

        adv_images = torch.cat(adv_imgs, dim=0)
        adv_targets = torch.cat(adv_tgts, dim=0)
        return torch.utils.data.TensorDataset(adv_images, adv_targets)

    def compute_mas_importance(self, forget_loader: DataLoader):
        """
        Memory Aware Synapses (MAS) importance from *forget* instances:
          omega_i = sum_x || d ||f(x)||_2 / d theta_i ||
        Returns: dict(param -> importance tensor normalized within each layer to [0,1]).
        """
        self.model.eval()
        # store initial requires_grad and prepare
        for p in self.model.parameters():
            p.grad = None

        # accumulate grads of output L2-norm wrt params
        omegas = {
            p: torch.zeros_like(p, device=self.device)
            for p in self.model.parameters()
            if p.requires_grad
        }

        for x, _ in tqdm(forget_loader, desc="MAS: accumulating importance"):
            x = x.to(self.device)
            self.model.zero_grad(set_to_none=True)
            out = self.model(x)  # [B, C]
            l2 = (out**2).sum(dim=1).mean()  # mean ||f(x)||_2^2
            grads = torch.autograd.grad(
                l2, omegas.keys(), retain_graph=False, allow_unused=True
            )
            for p, g in zip(list(omegas.keys()), grads):
                if g is not None:
                    omegas[p] += g.detach().abs()

        # per-parameter normalize to [0,1] within each layer
        for p in omegas:
            w = omegas[p]
            m, M = w.min(), w.max()
            if (M - m) > 0:
                omegas[p] = (w - m) / (M - m)
            else:
                omegas[p] = torch.zeros_like(w)

        return omegas

    def instancewise_unlearning(
        self,
        forget_loader: DataLoader,
        num_classes: int,
        mode: str = "misclassify",  # "misclassify" or "relabel"
        relabel_fn=None,  # function: (x, y_true) -> y_new (for mode="relabel")
        use_adv: bool = True,
        use_importance: bool = True,
        adv_pool: Dataset = None,  # pass precomputed pool (else we build one)
        adv_kwargs: dict = None,  # kwargs to generate_adv_pool_for_unlearning
        lambda_adv: float = 1.0,
        lambda_imp: float = 1e-4,  # strength for (1-omega)*||theta-theta0||^2
        max_epochs: int = 100,
        lr: float = 1e-3,
        momentum: float = 0.9,
        weight_decay: float = 1e-5,
        early_stop: bool = True,
    ):
        """
        Cha et al. (AAAI'24) instance-wise unlearning for pre-trained classifiers.
        - mode="misclassify": drive forget accuracy to 0% (negative-CE on true labels).
        - mode="relabel": make model predict new labels for forget items (positive-CE).
        Regularizers:
          - Adversarial examples from forget items (targeted PGD pool).
          - MAS importance penalty on parameter drift: sum (1-omega_i) * ||θ_i - θ_i^0||^2.
        """
        assert mode in ("misclassify", "relabel")
        self.model.train()

        # snapshot initial weights θ0
        theta0 = {
            p: p.detach().clone() for p in self.model.parameters() if p.requires_grad
        }

        # build adv pool if needed
        if use_adv and adv_pool is None:
            adv_kwargs = adv_kwargs or {}
            adv_pool = self.generate_adv_pool_for_unlearning(
                forget_loader=forget_loader,
                num_classes=num_classes,
                **adv_kwargs,
            )
        adv_loader = None
        if use_adv and adv_pool is not None:
            adv_loader = DataLoader(
                adv_pool,
                batch_size=forget_loader.batch_size,
                shuffle=True,
                drop_last=False,
            )

        # MAS importances (from forget items)
        omegas = None
        if use_importance:
            omegas = self.compute_mas_importance(forget_loader)

        # optimizer
        optimizer = optim.SGD(
            self.model.parameters(), lr=lr, momentum=momentum, weight_decay=weight_decay
        )
        ce = nn.CrossEntropyLoss()

        def forget_objective(logits, y_true):
            if mode == "misclassify":
                # maximize CE on true labels => minimize negative CE
                return -ce(logits, y_true)
            else:
                # relabel: get target labels via relabel_fn
                with torch.no_grad():
                    y_new = (
                        relabel_fn(logits.detach(), y_true)
                        if relabel_fn is not None
                        else (torch.randint_like(y_true, low=0, high=num_classes))
                    )
                    # ensure different from y_true
                    mask = y_new == y_true
                    if mask.any():
                        y_new[mask] = (y_new[mask] + 1) % num_classes
                return ce(logits, y_new)

        # simple metric on forget set
        def forget_accuracy():
            self.model.eval()
            correct, total = 0, 0
            with torch.no_grad():
                for x, y in forget_loader:
                    x, y = x.to(self.device), y.to(self.device)
                    pred = self.model(x).argmax(1)
                    correct += (pred == y).sum().item()
                    total += y.numel()
            self.model.train()
            return 100.0 * correct / max(total, 1)

        # training loop
        for ep in range(max_epochs):
            ep_loss, n_batches = 0.0, 0
            pbar = tqdm(forget_loader, desc=f"L2U ({mode}) epoch {ep+1}/{max_epochs}")
            adv_iter = iter(adv_loader) if adv_loader is not None else None

            for x_f, y_f in pbar:
                x_f, y_f = x_f.to(self.device), y_f.to(self.device)
                optimizer.zero_grad(set_to_none=True)

                # unlearning objective on the raw forget items
                logits_f = self.model(x_f)
                loss = forget_objective(logits_f, y_f)

                # adversarial-regularization term (classification on adv pool with their target labels)
                if adv_iter is not None:
                    try:
                        x_a, y_a = next(adv_iter)
                    except StopIteration:
                        adv_iter = iter(adv_loader)
                        x_a, y_a = next(adv_iter)
                    x_a, y_a = x_a.to(self.device), y_a.to(self.device)
                    logits_a = self.model(x_a)
                    loss_adv = ce(logits_a, y_a)
                    loss = loss + lambda_adv * loss_adv

                # MAS-style importance penalty
                if omegas is not None and lambda_imp > 0:
                    reg = 0.0
                    for p in self.model.parameters():
                        if p.requires_grad and p in theta0 and p in omegas:
                            w = (
                                1.0 - omegas[p]
                            )  # penalize changes to weights UNIMPORTANT for the forget set
                            reg = reg + (w * (p - theta0[p]) ** 2).sum()
                    loss = loss + lambda_imp * reg

                loss.backward()
                optimizer.step()

                ep_loss += loss.item()
                n_batches += 1
                pbar.set_postfix({"loss": f"{ep_loss/max(n_batches,1):.4f}"})

            # early stopping per paper
            if early_stop:
                acc_f = forget_accuracy()
                if (mode == "misclassify" and acc_f <= 0.0 + 1e-6) or (
                    mode == "relabel" and acc_f >= 100.0 - 1e-6
                ):
                    break

        self.model.eval()

    def boundary_expand_unlearning(
        self,
        forget_loader: DataLoader,
        epochs: int = 3,  # usually just a few
        lr: float = 1e-3,  # small LR is enough
        momentum: float = 0.9,
        weight_decay: float = 0.0,
        freeze_backbone: bool = True,
    ) -> None:
        """
        Boundary Expanding unlearning (shadow-class trick):
          1) Append one output neuron to the final linear classifier (shadow class).
          2) Fine-tune on forget samples with label = shadow index.
          3) Prune the shadow neuron; model returns to original output size.

        Args:
            forget_loader: DataLoader over the forget set (labels are ignored)
            epochs: epochs to fine-tune on shadow class
            lr, momentum, weight_decay: optimizer hyperparams
            freeze_backbone: if True, only train the head during expansion
        """

        # ---- helpers to find & replace the final nn.Linear head ----
        def _get_last_linear_and_name(model):
            last_name, last_linear = None, None
            for name, m in model.named_modules():
                if isinstance(m, nn.Linear):
                    last_name, last_linear = name, m
            if last_linear is None:
                raise ValueError(
                    "BoundaryExpand: couldn't find a final nn.Linear layer."
                )
            return last_name, last_linear

        def _get_parent_and_attr(model, module_name: str):
            # module_name like "fc" or "classifier.1" etc.
            parts = module_name.split(".")
            parent = model
            for p in parts[:-1]:
                if p.isdigit():
                    parent = parent[int(p)]
                else:
                    parent = getattr(parent, p)
            return parent, parts[-1]

        # ---- 0) locate the head and expand it by 1 output (shadow) ----
        self.model.eval()
        head_name, head = _get_last_linear_and_name(self.model)
        in_f, out_f = head.in_features, head.out_features
        shadow_index = out_f  # new class index at the end

        new_head = nn.Linear(in_f, out_f + 1, bias=head.bias is not None).to(
            self.device
        )
        with torch.no_grad():
            new_head.weight[:out_f].copy_(head.weight.data)
            if head.bias is not None:
                new_head.bias[:out_f].copy_(head.bias.data)
            # init the shadow neuron (xavier/zeros is fine)
            nn.init.xavier_uniform_(new_head.weight[out_f : out_f + 1])
            if new_head.bias is not None:
                new_head.bias[out_f].zero_()

        parent, attr = _get_parent_and_attr(self.model, head_name)
        setattr(parent, attr, new_head)

        # ---- 1) (optional) freeze backbone so only the head adapts ----
        train_params = list(self.model.parameters())
        if freeze_backbone:
            for p in self.model.parameters():
                p.requires_grad = False
            for p in new_head.parameters():
                p.requires_grad = True
            train_params = [p for p in new_head.parameters() if p.requires_grad]

        # ---- 2) fine-tune on forget set with shadow label ----
        self.model.train()
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.SGD(
            train_params, lr=lr, momentum=momentum, weight_decay=weight_decay
        )

        for ep in range(epochs):
            ep_loss, n_batches = 0.0, 0
            pbar = tqdm(
                forget_loader, desc=f"Boundary Expand fine-tune {ep+1}/{epochs}"
            )
            for x, _ in pbar:
                x = x.to(self.device)
                y_shadow = torch.full(
                    (x.size(0),), shadow_index, device=self.device, dtype=torch.long
                )

                optimizer.zero_grad(set_to_none=True)
                out = self.model(x)
                loss = criterion(out, y_shadow)
                loss.backward()
                optimizer.step()

                ep_loss += loss.item()
                n_batches += 1
                pbar.set_postfix({"loss": f"{ep_loss/max(n_batches,1):.4f}"})

        # ---- 3) prune the shadow neuron and restore original head size ----
        pruned_head = nn.Linear(in_f, out_f, bias=new_head.bias is not None).to(
            self.device
        )
        with torch.no_grad():
            pruned_head.weight.copy_(new_head.weight[:out_f])
            if pruned_head.bias is not None:
                pruned_head.bias.copy_(new_head.bias[:out_f])

        setattr(parent, attr, pruned_head)

        # unfreeze if we froze
        if freeze_backbone:
            for p in self.model.parameters():
                p.requires_grad = True

        self.model.eval()