1{
2  "op": "replace",
3  "filename": "opacus/bench/unlearning/custom_unlearning.py",
4  "start_line": 1,
5  "end_line": 25,
6  "content": "\"\"\"Editable unlearning method for MLS-Bench.\"\"\"\n\nimport torch\nimport torch.nn.functional as F\n\n# ============================================================\n# EDITABLE\n# ============================================================\nclass UnlearningMethod:\n    \"\"\"Complementary-Label Distillation with Confidence-Gated Forgetting.\n\n    Novel contribution: Instead of naive gradient ascent (negative_gradient),\n    entropy maximization (bad_teacher), or logit flattening (scrub), we use\n    complementary label targets — uniform over all classes EXCEPT the true\n    forget class — to guide the model toward the distribution a model never\n    trained on the forget class would exhibit. The forgetting pressure is\n    adaptively scaled by the model's current confidence on the true (forget)\n    class, ensuring strong initial erasure that naturally tapers as forgetting\n    succeeds, preserving retain accuracy.\n    \"\"\"\n\n    def __init__(self):\n        self.retain_weight = 1.0\n        self.forget_complement_weight = 1.0\n        self.forget_neg_weight = 0.5\n        self.forget_entropy_weight = 0.2\n\n    def unlearn_step(self, model, retain_batch, forget_batch, optimizer, step, epoch):\n        retain_x, retain_y = retain_batch\n        forget_x, forget_y = forget_batch\n\n        # === Retain objective: standard cross-entropy ===\n        retain_logits = model(retain_x)\n        retain_loss = F.cross_entropy(retain_logits, retain_y)\n\n        # === Forget objective: multi-pronged ===\n        forget_logits = model(forget_x)\n        num_classes = forget_logits.shape[1]\n        forget_probs = F.softmax(forget_logits, dim=1)\n\n        # 1. Complementary-label distillation: target is uniform over non-true classes\n        #    This actively teaches the model to classify forget data as other classes\n        complement_targets = torch.ones_like(forget_logits) / (num_classes - 1)\n        complement_targets.scatter_(1, forget_y.unsqueeze(1), 0.0)\n        log_probs = F.log_softmax(forget_logits, dim=1)\n        forget_complement = -(complement_targets * log_probs).sum(dim=1).mean()\n\n        # 2. Negative gradient: ascend on standard CE for the forget class\n        forget_ce = F.cross_entropy(forget_logits, forget_y)\n\n        # 3. Entropy maximization: ensure predictions spread out, don't collapse\n        #    to a single wrong class (which would be detectable by MIA)\n        forget_entropy = -(forget_probs * torch.log(forget_probs.clamp_min(1e-8))).sum(dim=1).mean()\n\n        # Confidence-gated adaptive weight: scale forgetting pressure by\n        # current confidence on the true class. When model still confidently\n        # predicts the forget class, apply strong pressure; when already\n        # forgotten, reduce to avoid harming retain performance.\n        with torch.no_grad():\n            confidence = forget_probs.gather(1, forget_y.unsqueeze(1)).squeeze(1).mean()\n            adaptive = confidence.clamp(min=0.05)\n\n        loss = (self.retain_weight * retain_loss\n                + adaptive * self.forget_complement_weight * forget_complement\n                - adaptive * self.forget_neg_weight * forget_ce\n                - self.forget_entropy_weight * forget_entropy)\n\n        optimizer.zero_grad()\n        loss.backward()\n        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)\n        optimizer.step()\n\n        return {\n            \"loss\": loss.item(),\n            \"retain_loss\": retain_loss.item(),\n            \"forget_ce\": forget_ce.item(),\n            \"confidence\": confidence.item()\n        }\n# ============================================================\n# END EDITABLE\n# ============================================================\n"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/bench/unlearning/custom_unlearning.py",
4  "start_line": 1,
5  "end_line": 24,
6  "content": "\"\"\"Editable unlearning method for MLS-Bench.\"\"\"\n\nimport torch\nimport torch.nn.functional as F\n\n# ============================================================\n# EDITABLE\n# ============================================================\nclass UnlearningMethod:\n    \"\"\"Class-Targeted Distributional Unlearning (CTDU).\n    \n    Novel contributions:\n    1. Class-Targeted Soft Cross-Entropy: Instead of pushing forget-set predictions\n       toward a uniform distribution over ALL classes, we target a distribution that\n       is uniform over non-forget classes with zero mass on the forget class. This\n       better emulates the behavior of a model that was never trained on the forget\n       class, which would spread predictions across known (non-forget) classes.\n    \n    2. Stochastic Complementary Relabeling: Each step, forget examples are randomly\n       assigned to non-forget class labels. Since labels change every step, the model\n       cannot form consistent patterns for forget features, effectively overwriting\n       old associations with stochastic noise. The average gradient over steps matches\n       the class-targeted distribution, providing coherent forgetting direction.\n    \n    3. Adaptive Forget Pressure: Forgetting weight decays over epochs - aggressive\n       early (when model still remembers) and conservative later (to preserve retain\n       accuracy and avoid MIA-detectable overcompensation).\n    \"\"\"\n    \n    def __init__(self):\n        self.num_classes = None\n        self.forget_class = None\n        self.target = None\n        \n    def unlearn_step(self, model, retain_batch, forget_batch, optimizer, step, epoch):\n        retain_x, retain_y = retain_batch\n        forget_x, forget_y = forget_batch\n        \n        # Forward pass both batches\n        retain_logits = model(retain_x)\n        forget_logits = model(forget_x)\n        \n        # Initialize on first call\n        if self.num_classes is None:\n            self.num_classes = retain_logits.shape[1]\n            self.forget_class = forget_y[0].item()\n            nc, fc = self.num_classes, self.forget_class\n            # Target: uniform over non-forget classes, zero on forget class\n            self.target = torch.full((nc,), 1.0 / (nc - 1), device=retain_x.device)\n            self.target[fc] = 0.0\n        \n        nc = self.num_classes\n        fc = self.forget_class\n        \n        # Move target to correct device if needed\n        if self.target.device != retain_x.device:\n            self.target = self.target.to(retain_x.device)\n        \n        # === 1. Retain Loss: Standard Cross-Entropy ===\n        retain_loss = F.cross_entropy(retain_logits, retain_y)\n        \n        # === 2. Class-Targeted Soft Cross-Entropy (Novel) ===\n        # Push forget predictions to uniform over non-forget classes\n        forget_log_probs = F.log_softmax(forget_logits, dim=1)\n        # Soft CE: -sum(target * log_probs), equivalent to KL + constant\n        soft_ce = -(self.target.unsqueeze(0) * forget_log_probs).sum(dim=1).mean()\n        \n        # === 3. Stochastic Complementary Relabeling (Novel) ===\n        # Randomly assign forget examples to non-forget classes each step\n        bs = forget_y.size(0)\n        rand_idx = torch.randint(0, nc - 1, (bs,), device=forget_y.device)\n        # Shift indices >= forget_class to skip the forget class\n        rand_labels = rand_idx + (rand_idx >= fc).long()\n        relabel_loss = F.cross_entropy(forget_logits, rand_labels)\n        \n        # === 4. Gradient Ascent on Forget CE ===\n        forget_ce = F.cross_entropy(forget_logits, forget_y)\n        \n        # === Adaptive Weighting ===\n        # More aggressive forgetting early, more conservative later\n        progress = epoch / 20.0  # 20 epochs total\n        forget_pressure = 1.0 - 0.3 * progress  # 1.0 -> 0.7\n        \n        # Combined loss\n        loss = (retain_loss \n                + forget_pressure * 1.0 * soft_ce        # Class-targeted forgetting\n                + forget_pressure * 0.3 * relabel_loss    # Stochastic relabeling\n                - forget_pressure * 0.3 * forget_ce)      # Gradient ascent\n        \n        optimizer.zero_grad()\n        loss.backward()\n        # Gradient clipping for stability (negative CE can cause large gradients)\n        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)\n        optimizer.step()\n        \n        return {\n            \"loss\": loss.item(),\n            \"retain_loss\": retain_loss.item(),\n            \"soft_ce\": soft_ce.item(),\n            \"relabel_loss\": relabel_loss.item(),\n            \"forget_ce\": forget_ce.item()\n        }\n# ============================================================\n# END EDITABLE\n# ============================================================\n"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/bench/unlearning/custom_unlearning.py",
4  "start_line": 1,
5  "end_line": 24,
6  "content": "\"\"\"Editable unlearning method for MLS-Bench.\"\"\"\n\nimport torch\nimport torch.nn.functional as F\n\n# ============================================================\n# EDITABLE\n# ============================================================\nclass UnlearningMethod:\n    \"\"\"Confidence-Adaptive Entropy Maximization (CAEM) for Machine Unlearning.\n\n    Novel contribution: Per-sample adaptive weighting of the forgetting objective\n    based on current memorization level (softmax confidence on the true class).\n    Samples that are still well-memorized receive stronger unlearning gradients,\n    leading to more efficient and thorough forgetting.\n\n    The forgetting objective combines:\n    1. KL divergence to uniform with per-sample adaptive weights (confidence-scaled)\n    2. Gradient ascent on forget cross-entropy (early epochs for rapid unlearning)\n\n    The retain objective uses standard cross-entropy.\n    \"\"\"\n\n    def __init__(self):\n        self.forget_class = None\n        self.kl_weight = 1.0\n        self.ga_weight = 1.0\n        self.total_epochs = 20\n\n    def unlearn_step(self, model, retain_batch, forget_batch, optimizer, step, epoch):\n        retain_x, retain_y = retain_batch\n        forget_x, forget_y = forget_batch\n\n        # Auto-detect forget class on first call\n        if self.forget_class is None:\n            self.forget_class = forget_y[0].item()\n\n        # Forward pass\n        retain_logits = model(retain_x)\n        forget_logits = model(forget_x)\n        num_classes = retain_logits.shape[1]\n\n        # === Retain objective: standard cross-entropy ===\n        retain_loss = F.cross_entropy(retain_logits, retain_y)\n\n        # === Forget objective 1: Confidence-adaptive KL to uniform ===\n        forget_probs = F.softmax(forget_logits, dim=1)\n        # Memorization score: probability assigned to the true (forget) class\n        memo_score = forget_probs[:, self.forget_class].detach()\n\n        # Per-sample KL divergence to uniform distribution\n        uniform_target = torch.ones_like(forget_logits) / num_classes\n        forget_log_probs = F.log_softmax(forget_logits, dim=1)\n        per_sample_kl = F.kl_div(\n            forget_log_probs, uniform_target, reduction='none'\n        ).sum(dim=1)\n\n        # Adaptive weights: more memorized samples get stronger unlearning\n        # Weight ranges from 1.0 (already forgotten) to 4.0 (fully memorized)\n        adaptive_weights = 1.0 + 3.0 * memo_score\n        kl_loss = (adaptive_weights * per_sample_kl).mean()\n\n        # === Forget objective 2: Gradient ascent on CE (early phase only) ===\n        forget_ce = F.cross_entropy(forget_logits, forget_y)\n        progress = epoch / self.total_epochs\n        # Decay GA from full weight to 0 over first 40% of training\n        ga_scale = max(0.0, 1.0 - 2.5 * progress)\n\n        # Combined loss\n        loss = (retain_loss\n                + self.kl_weight * kl_loss\n                - self.ga_weight * ga_scale * forget_ce)\n\n        optimizer.zero_grad()\n        loss.backward()\n        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)\n        optimizer.step()\n\n        return {\n            \"loss\": loss.item(),\n            \"retain_loss\": retain_loss.item(),\n            \"kl_loss\": kl_loss.item(),\n            \"forget_ce\": forget_ce.item(),\n            \"avg_memo_score\": memo_score.mean().item(),\n        }\n# ============================================================\n# END EDITABLE\n# ============================================================\n"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/bench/unlearning/custom_unlearning.py",
4  "start_line": 86,
5  "end_line": 89,
6  "content": "# ============================================================\n# END EDITABLE\n# ============================================================\n"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/bench/unlearning/custom_unlearning.py",
4  "start_line": 1,
5  "end_line": 88,
6  "content": "\"\"\"Editable unlearning method for MLS-Bench.\"\"\"\n\nimport torch\nimport torch.nn.functional as F\n\n# ============================================================\n# EDITABLE\n# ============================================================\nclass UnlearningMethod:\n    \"\"\"Confidence-Adaptive Entropy Maximization with Logit Regularization (CAEM-LR).\n\n    Novel contribution: Per-sample adaptive weighting of the forgetting loss\n    based on the model's current confidence (memorization proxy). Samples that\n    are still well-memorized receive proportionally stronger unlearning pressure,\n    making forgetting more efficient and targeted. This is combined with:\n    1. KL divergence to uniform distribution (distributional forgetting)\n    2. Logit magnitude regularization (reduces MIA-detectable confidence patterns)\n    3. No gradient ascent (avoids the instability of ascending the CE loss)\n    \"\"\"\n\n    def __init__(self):\n        self.forget_class = None\n        # KL-to-uniform weight (controls forgetting strength)\n        self.kl_weight = 0.3\n        # Logit-squared regularization weight (additional MIA defense)\n        self.logit_reg_weight = 0.05\n\n    def unlearn_step(self, model, retain_batch, forget_batch, optimizer, step, epoch):\n        retain_x, retain_y = retain_batch\n        forget_x, forget_y = forget_batch\n\n        # Auto-detect forget class on first call\n        if self.forget_class is None:\n            self.forget_class = forget_y[0].item()\n\n        # Forward pass (single pass per data split)\n        retain_logits = model(retain_x)\n        forget_logits = model(forget_x)\n        num_classes = retain_logits.shape[1]\n\n        # === Retain objective: standard cross-entropy ===\n        retain_loss = F.cross_entropy(retain_logits, retain_y)\n\n        # === Forget objective 1: Confidence-adaptive KL to uniform ===\n        # Compute per-sample memorization score (detached for weighting only)\n        with torch.no_grad():\n            forget_probs = F.softmax(forget_logits, dim=1)\n            memo_score = forget_probs[:, self.forget_class]\n            # Adaptive weight: 1.0 (already forgotten) to 3.0 (fully memorized)\n            adaptive_weights = 1.0 + 2.0 * memo_score\n\n        # KL(uniform || model) per sample\n        uniform_target = torch.ones_like(forget_logits) / num_classes\n        forget_log_probs = F.log_softmax(forget_logits, dim=1)\n        per_sample_kl = F.kl_div(\n            forget_log_probs, uniform_target, reduction='none'\n        ).sum(dim=1)\n\n        # Weighted mean: more memorized samples get stronger gradients\n        kl_loss = (adaptive_weights * per_sample_kl).mean()\n\n        # === Forget objective 2: Logit magnitude regularization ===\n        # Pushes all logits toward zero → low confidence → MIA resistance\n        logit_reg = (forget_logits ** 2).mean()\n\n        # === Combined loss (no gradient ascent for stability) ===\n        loss = (retain_loss\n                + self.kl_weight * kl_loss\n                + self.logit_reg_weight * logit_reg)\n\n        optimizer.zero_grad()\n        loss.backward()\n        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)\n        optimizer.step()\n\n        return {\n            \"loss\": loss.item(),\n            \"retain_loss\": retain_loss.item(),\n            \"kl_loss\": kl_loss.item(),\n            \"logit_reg\": logit_reg.item(),\n            \"avg_memo\": memo_score.mean().item(),\n        }\n# ============================================================\n# END EDITABLE\n# ============================================================\n"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/bench/unlearning/custom_unlearning.py",
4  "start_line": 1,
5  "end_line": 85,
6  "content": "\"\"\"Editable unlearning method for MLS-Bench.\"\"\"\n\nimport torch\nimport torch.nn.functional as F\n\n# ============================================================\n# EDITABLE\n# ============================================================\nclass UnlearningMethod:\n    \"\"\"Confidence-Adaptive Entropy Maximization with BN-Aware Unlearning (CAEM-BN).\n\n    Novel contributions:\n    1. Per-sample adaptive weighting of the forgetting loss based on the model's\n       current confidence on the true (forget) class — a memorization proxy.\n       More-memorized samples receive proportionally stronger unlearning gradients,\n       yielding more efficient and targeted forgetting.\n    2. BatchNorm freezing during unlearning to prevent forget-data statistics from\n       corrupting running mean/var, which otherwise degrades retain accuracy at\n       test time when BN uses running stats.\n    3. Dual forgetting objective:\n       (a) KL(uniform || model) pushes softmax toward uniform (distributional)\n       (b) Logit-squared regularization pushes logits toward zero (MIA defense\n           by reducing detectable confidence patterns)\n    \"\"\"\n\n    def __init__(self):\n        self.forget_class = None\n        # Moderate KL weight — balances forget quality vs retain preservation\n        self.kl_weight = 0.2\n        # Small logit-squared reg — complements KL by reducing logit magnitude\n        self.logit_reg_weight = 0.03\n\n    def unlearn_step(self, model, retain_batch, forget_batch, optimizer, step, epoch):\n        retain_x, retain_y = retain_batch\n        forget_x, forget_y = forget_batch\n\n        # Auto-detect forget class on first call\n        if self.forget_class is None:\n            self.forget_class = forget_y[0].item()\n\n        # ── BN freezing: use eval mode so running stats are NOT updated ──\n        # Prevents forget-data batch statistics from corrupting the running\n        # mean/var that will be used during evaluation. BN affine parameters\n        # (weight, bias) are still trainable via gradient descent.\n        for module in model.modules():\n            if isinstance(module, (torch.nn.BatchNorm1d, torch.nn.BatchNorm2d)):\n                module.eval()\n\n        # ── Forward pass ──\n        retain_logits = model(retain_x)\n        forget_logits = model(forget_x)\n        num_classes = retain_logits.shape[1]\n\n        # ── Retain objective: standard cross-entropy ──\n        retain_loss = F.cross_entropy(retain_logits, retain_y)\n\n        # ── Forget objective 1: confidence-adaptive KL to uniform ──\n        # Compute per-sample memorization score (detached — used only for weighting)\n        with torch.no_grad():\n            forget_probs = F.softmax(forget_logits, dim=1)\n            memo_score = forget_probs[:, self.forget_class]\n            # Adaptive weight: 1.0 (fully forgotten) → 3.0 (fully memorized)\n            adaptive_weights = 1.0 + 2.0 * memo_score\n\n        # KL(uniform || model) per sample\n        uniform_target = torch.ones_like(forget_logits) / num_classes\n        forget_log_probs = F.log_softmax(forget_logits, dim=1)\n        per_sample_kl = F.kl_div(\n            forget_log_probs, uniform_target, reduction='none'\n        ).sum(dim=1)\n        kl_loss = (adaptive_weights * per_sample_kl).mean()\n\n        # ── Forget objective 2: logit magnitude regularization ──\n        logit_reg = (forget_logits ** 2).mean()\n\n        # ── Combined loss ──\n        loss = (retain_loss\n                + self.kl_weight * kl_loss\n                + self.logit_reg_weight * logit_reg)\n\n        optimizer.zero_grad()\n        loss.backward()\n        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)\n        optimizer.step()\n\n        return {\n            \"loss\": loss.item(),\n            \"retain_loss\": retain_loss.item(),\n            \"kl_loss\": kl_loss.item(),\n            \"logit_reg\": logit_reg.item(),\n            \"avg_memo\": memo_score.mean().item(),\n        }\n# ============================================================\n# END EDITABLE\n# ============================================================\n"
7}