1{
2  "op": "replace",
3  "filename": "opacus/bench/membership/custom_membership_defense.py",
4  "start_line": 1,
5  "end_line": 23,
6  "content": "\"\"\"Editable membership-inference defense for MLS-Bench.\"\"\"\n\nimport torch\nimport torch.nn.functional as F\n\n# ============================================================\n# EDITABLE\n# ============================================================\nclass MembershipDefense:\n    \"\"\"Confidence-Bounded Entropy-Regularized Training (CBERT).\n    \n    Novel multi-component defense against membership inference attacks:\n    \n    1. Soft Logit Bounding: Uses scaled tanh to create a differentiable \n       upper bound on logit magnitudes, directly limiting prediction \n       confidence without hard truncation.\n    \n    2. Confidence-Adaptive Smoothing: Per-sample label smoothing that \n       increases with model confidence, creating a negative feedback loop \n       that prevents memorization of individual samples.\n    \n    3. Entropy Maximization: Encourages spreading probability mass across \n       classes while maintaining correct classification.\n    \n    4. Logit Norm Penalty: L2 regularization on raw logits to ensure \n       bounded representations complement the soft clipping.\n    \"\"\"\n\n    def __init__(self):\n        # Soft logit bounding scale\n        self.logit_bound = 3.0\n        # Confidence-adaptive smoothing parameters\n        self.base_smoothing = 0.15\n        self.max_smoothing = 0.45\n        self.confidence_threshold = 0.5\n        self.smoothing_sharpness = 8.0\n        # Entropy bonus weight\n        self.entropy_weight = 0.15\n        # L2 penalty on raw logits\n        self.logit_l2 = 0.008\n\n    def compute_loss(self, logits, labels, model, epoch):\n        B, C = logits.shape\n        \n        # 1. Soft logit bounding via scaled tanh\n        # Creates a differentiable ceiling on logit magnitudes,\n        # directly limiting the maximum achievable confidence\n        bounded = self.logit_bound * torch.tanh(logits / self.logit_bound)\n        \n        # 2. Confidence-adaptive per-sample label smoothing\n        # More confident predictions receive stronger smoothing,\n        # creating a self-regulating negative feedback loop\n        with torch.no_grad():\n            probs_det = torch.softmax(bounded, dim=1)\n            confidence = probs_det.max(dim=1).values\n            # Sigmoid gating: smooth transition from base to max smoothing\n            alpha = self.base_smoothing + (self.max_smoothing - self.base_smoothing) * \\\n                    torch.sigmoid(self.smoothing_sharpness * (confidence - self.confidence_threshold))\n        \n        # Create per-sample soft targets: blend one-hot with uniform\n        one_hot = F.one_hot(labels, C).float()\n        uniform = torch.ones_like(one_hot) / C\n        soft_targets = (1 - alpha.unsqueeze(1)) * one_hot + alpha.unsqueeze(1) * uniform\n        \n        # Cross-entropy with per-sample soft targets\n        log_probs = F.log_softmax(bounded, dim=1)\n        ce = -(soft_targets * log_probs).sum(dim=1).mean()\n        \n        # 3. Entropy maximization on bounded logits\n        probs = torch.softmax(bounded, dim=1)\n        entropy = -(probs * torch.log(probs.clamp_min(1e-8))).sum(dim=1).mean()\n        \n        # 4. L2 penalty on raw logits to keep representations bounded\n        logit_reg = self.logit_l2 * (logits ** 2).mean()\n        \n        return ce - self.entropy_weight * entropy + logit_reg\n# ============================================================\n# END EDITABLE\n# ============================================================\n"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_membership_defense.py",
4  "start_line": 1,
5  "end_line": 31,
6  "content": "\"\"\"Editable membership-inference defense for MLS-Bench.\"\"\"\n\nimport torch\nimport torch.nn.functional as F\n\n# ============================================================\n# EDITABLE\n# ============================================================\nclass MembershipDefense:\n    \"\"\"Variance-Constrained Relaxed Entropy Loss (VCRE).\n\n    Novel contribution: Combines relaxed loss training with per-sample\n    loss variance regularization, entropy shaping, and logit magnitude\n    control to minimize membership inference leakage.\n\n    Key insight: MIA attacks exploit the distributional gap between\n    per-sample losses of members vs non-members. By (1) relaxing\n    gradients for overly well-fit samples, (2) penalizing the variance\n    of per-sample losses to encourage uniform treatment, (3) maximizing\n    prediction entropy, and (4) penalizing logit magnitudes, we\n    systematically close this gap without sacrificing generalization.\n    \"\"\"\n\n    def __init__(self):\n        # Relaxation parameters (proven in RelaxLoss)\n        self.target_ce = 1.0\n        self.alpha = 0.5\n        # Label smoothing strength\n        self.label_smoothing = 0.15\n        # Entropy bonus weight\n        self.entropy_weight = 0.15\n        # Novel: loss variance penalty weight\n        self.var_weight = 0.3\n        # Logit magnitude penalty weight\n        self.logit_reg_weight = 0.01\n\n    def compute_loss(self, logits, labels, epoch):\n        # Per-sample cross-entropy with label smoothing\n        ce_per = F.cross_entropy(logits, labels, reduction='none',\n                                  label_smoothing=self.label_smoothing)\n\n        # Relaxed loss: for well-fit samples (low CE), reverse the gradient\n        # direction to prevent memorization; for under-fit samples, train normally\n        relaxed = torch.where(\n            ce_per > self.target_ce,\n            ce_per,\n            self.alpha * (self.target_ce - ce_per),\n        )\n        main_loss = relaxed.mean()\n\n        # Novel: Per-sample loss variance regularization\n        # By penalizing the spread of per-sample losses, we encourage the model\n        # to treat all training samples uniformly. This directly reduces the\n        # statistical signal that MIA attacks use (the gap between the loss\n        # distributions of members vs non-members).\n        loss_var = ce_per.var()\n\n        # Maximum entropy regularization: push output distributions toward\n        # higher entropy, reducing overconfidence on training samples\n        probs = F.softmax(logits, dim=1)\n        entropy = -(probs * torch.log(probs.clamp_min(1e-8))).sum(dim=1).mean()\n\n        # Logit magnitude penalty: prevents the model from producing extreme\n        # logit values that lead to high-confidence predictions exploitable by MIA\n        logit_reg = (logits ** 2).mean()\n\n        return (main_loss\n                + self.var_weight * loss_var\n                - self.entropy_weight * entropy\n                + self.logit_reg_weight * logit_reg)\n# ============================================================\n# END EDITABLE\n# ============================================================\n"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_membership_defense.py",
4  "start_line": 1,
5  "end_line": 31,
6  "content": "\"\"\"Editable membership-inference defense for MLS-Bench.\"\"\"\n\nimport torch\nimport torch.nn.functional as F\n\n# ============================================================\n# EDITABLE\n# ============================================================\nclass MembershipDefense:\n    \"\"\"Variance-Regularized Confidence-Bounded Loss (VRCBL).\n\n    Novel contribution: Per-sample loss variance regularization combined with\n    adaptive flooding, label smoothing, and entropy maximization. The variance\n    penalty directly targets the MIA attack signal by encouraging uniform\n    loss across all training samples, closing the gap between memorized and\n    non-memorized examples that attackers exploit.\n    \"\"\"\n\n    def __init__(self):\n        self.label_smoothing = 0.15\n        self.flood_ratio = 0.15        # fraction of max entropy as flood level\n        self.variance_weight = 1.0      # weight for per-sample loss variance penalty\n        self.entropy_weight = 0.15      # weight for entropy bonus\n        self.logit_reg_weight = 0.005   # weight for logit L2 penalty\n\n    def compute_loss(self, logits, labels, epoch):\n        num_classes = logits.size(1)\n\n        # --- Per-sample cross-entropy with label smoothing ---\n        ce_per = F.cross_entropy(\n            logits, labels,\n            label_smoothing=self.label_smoothing,\n            reduction='none'\n        )\n        ce_mean = ce_per.mean()\n\n        # --- Adaptive flooding: floor on mean loss to prevent memorization ---\n        # Flood level scales with log(num_classes) so it adapts to dataset\n        flood_level = self.flood_ratio * torch.log(\n            torch.tensor(float(num_classes), device=logits.device)\n        )\n        ce_flooded = (ce_mean - flood_level).abs() + flood_level\n\n        # --- Per-sample loss variance penalty (NOVEL) ---\n        # Encourages uniform loss distribution across samples, directly\n        # reducing the signal that MIA attacks exploit (loss gap between\n        # memorized train samples and unseen test samples)\n        loss_var = ce_per.var()\n\n        # --- Entropy bonus: encourages less peaked predictions ---\n        probs = torch.softmax(logits, dim=1)\n        entropy = -(probs * torch.log(probs.clamp_min(1e-8))).sum(dim=1).mean()\n\n        # --- Logit L2 regularization: bounds output confidence ---\n        logit_reg = (logits ** 2).mean()\n\n        # --- Epoch-adaptive regularization (stronger in later epochs) ---\n        # Later epochs are more prone to memorization\n        adaptive = min(1.0, 0.5 + 0.5 * epoch / 50.0)\n\n        total = (\n            ce_flooded\n            + self.variance_weight * adaptive * loss_var\n            - self.entropy_weight * adaptive * entropy\n            + self.logit_reg_weight * logit_reg\n        )\n        return total\n# ============================================================\n# END EDITABLE\n# ============================================================\n"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_membership_defense.py",
4  "start_line": 68,
5  "end_line": 71,
6  "content": "# ============================================================\n# END EDITABLE\n# ============================================================\n"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_membership_defense.py",
4  "start_line": 1,
5  "end_line": 31,
6  "content": "\"\"\"Editable membership-inference defense for MLS-Bench.\"\"\"\n\nimport math\nimport torch\nimport torch.nn.functional as F\n\n# ============================================================\n# EDITABLE\n# ============================================================\nclass MembershipDefense:\n    \"\"\"Smooth Confidence Flooding with Multi-Regularization (SCF-MR).\n    \n    Novel defense combining:\n    1. Softplus-based smooth flooding: Unlike RelaxLoss's hard piecewise function,\n       uses softplus(threshold - per_sample_ce) to create a differentiable penalty\n       that smoothly increases as samples become too well-learned. This avoids \n       gradient discontinuities while effectively preventing memorization.\n    2. Label smoothing for baseline confidence calibration.\n    3. Entropy maximization bonus to spread probability mass.\n    4. Logit norm regularization to directly limit confidence magnitude.\n    \n    The flood level automatically adapts to the number of classes.\n    \"\"\"\n\n    def __init__(self):\n        self.label_smoothing = 0.1\n        self.flood_weight = 0.5\n        self.flood_beta = 3.0       # sharpness of flooding transition\n        self.entropy_weight = 0.1\n        self.logit_norm_weight = 0.005\n\n    def compute_loss(self, logits, labels, epoch):\n        num_classes = logits.size(1)\n        \n        # Class-adaptive flood level: scales with log(num_classes)\n        # For 10 classes: ~0.69, for 100 classes: ~1.38\n        flood_level = 0.3 * math.log(num_classes)\n        \n        # Per-sample cross-entropy with label smoothing\n        ce_per = F.cross_entropy(logits, labels, reduction='none',\n                                  label_smoothing=self.label_smoothing)\n        ce = ce_per.mean()\n        \n        # Smooth flooding: softplus creates a differentiable penalty for \n        # samples with CE below the flood level, preventing memorization\n        flood_penalty = F.softplus(flood_level - ce_per, beta=self.flood_beta).mean()\n        \n        # Predictive entropy bonus: encourages less peaked distributions\n        probs = torch.softmax(logits, dim=1)\n        entropy = -(probs * torch.log(probs.clamp_min(1e-8))).sum(dim=1).mean()\n        \n        # Logit norm regularization: directly limits confidence magnitude\n        logit_norm = logits.pow(2).mean()\n        \n        loss = (ce \n                + self.flood_weight * flood_penalty \n                - self.entropy_weight * entropy \n                + self.logit_norm_weight * logit_norm)\n        \n        return loss\n# ============================================================\n# END EDITABLE\n# ============================================================\n"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_membership_defense.py",
4  "start_line": 1,
5  "end_line": 64,
6  "content": "\"\"\"Editable membership-inference defense for MLS-Bench.\"\"\"\n\nimport math\nimport torch\nimport torch.nn.functional as F\n\n# ============================================================\n# EDITABLE\n# ============================================================\nclass MembershipDefense:\n    \"\"\"Smooth Relaxation with Confidence Calibration (SRC).\n\n    A novel training-time membership inference defense with three mechanisms:\n\n    1. Sigmoid-Smoothed Loss Relaxation: Unlike RelaxLoss's hard piecewise\n       switching between normal and inverted gradients, we use a sigmoid to\n       smoothly interpolate. Well-learned samples (low per-sample CE) get\n       partially inverted gradients pushing them toward lower confidence,\n       while poorly-learned samples receive standard CE gradients.\n       This eliminates the gradient discontinuity at the threshold, leading\n       to smoother optimization. The blending is self-curricularizing:\n       early in training all samples have high CE so blend~0 (pure CE),\n       and regularization automatically activates as samples become memorized.\n\n    2. Label Smoothing: Standard soft targets prevent overconfident predictions.\n\n    3. Direct Max-Confidence Penalty: Soft upper bound on maximum prediction\n       probability via softplus, directly reducing the signal available to\n       confidence-based MIA attacks.\n\n    All thresholds adapt to the number of classes via log-scaling.\n    \"\"\"\n\n    def __init__(self):\n        self.label_smoothing = 0.1\n        self.alpha = 0.5       # inversion strength for well-learned samples\n        self.sharpness = 5.0   # sigmoid transition sharpness\n        self.conf_weight = 0.5 # weight for max-confidence penalty\n        self.conf_cap = 0.85   # soft cap on max prediction probability\n\n    def compute_loss(self, logits, labels, epoch):\n        num_classes = logits.size(1)\n\n        # Class-adaptive target CE: scales with log(num_classes)\n        # 10 classes: ~1.0, 100 classes: ~1.92\n        target_ce = 0.4 * math.log(num_classes) + 0.08\n\n        # Per-sample cross-entropy with label smoothing\n        ce_per = F.cross_entropy(logits, labels, reduction='none',\n                                  label_smoothing=self.label_smoothing)\n\n        # Smooth blending factor (detached for stable gradients):\n        # ≈1 for well-learned samples (low CE), ≈0 for poorly-learned\n        with torch.no_grad():\n            blend = torch.sigmoid(self.sharpness * (target_ce - ce_per))\n\n        # Smoothly relaxed loss:\n        # For high-loss samples (blend≈0): standard CE gradient for learning\n        # For low-loss samples (blend≈1): inverted gradient pushes away from\n        #   memorization, preventing the train-test gap MIA exploits\n        relaxed = (1.0 - blend) * ce_per + blend * self.alpha * (target_ce - ce_per)\n        loss = relaxed.mean()\n\n        # Soft max-confidence penalty: directly limits prediction confidence\n        # which is the primary signal used by confidence-based MIA attacks\n        probs = torch.softmax(logits, dim=1)\n        max_prob = probs.max(dim=1)[0]\n        conf_penalty = F.softplus(max_prob - self.conf_cap, beta=10.0).mean()\n\n        return loss + self.conf_weight * conf_penalty\n# ============================================================\n# END EDITABLE\n# ============================================================\n"
7}