1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_membership_defense.py",
4  "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-capped adaptive smoothing for membership privacy.\n\n    Idea:\n    - keep standard ERM behavior on hard / misclassified examples;\n    - only regularize examples that are already correct and becoming\n      overconfident, since those dominate confidence-based MIA signals;\n    - use an EMA of correct-class confidence to set an adaptive target,\n      making the defense robust across datasets with very different\n      difficulty and class count.\n    \"\"\"\n\n    def __init__(self):\n        # Delay privacy regularization until the model has learned a useful\n        # classifier, then smoothly ramp the defense on.\n        self.warmup_epochs = 8\n        self.ramp_tau = 12.0\n\n        # EMA target for correctly-classified confidence.\n        self.ema_momentum = 0.90\n        self.correct_conf_ema = None\n\n        # Confidence target schedule. During warmup the target is near 1.0,\n        # then it anneals toward an EMA-based cap.\n        self.init_target = 0.98\n        self.min_target = 0.66\n        self.max_target = 0.82\n        self.conf_gap = 0.08\n\n        # Regularization strengths.\n        self.max_smoothing = 0.16\n        self.excess_weight = 0.70\n        self.mean_weight = 0.20\n        self.entropy_weight = 0.01\n\n    def _update_conf_ema(self, value):\n        value = value.detach()\n        if self.correct_conf_ema is None:\n            self.correct_conf_ema = value\n        else:\n            self.correct_conf_ema = (\n                self.ema_momentum * self.correct_conf_ema.to(value.device)\n                + (1.0 - self.ema_momentum) * value\n            )\n\n    def compute_loss(self, logits, labels, epoch):\n        log_probs = F.log_softmax(logits, dim=1)\n        probs = log_probs.exp()\n\n        true_log_probs = log_probs.gather(1, labels.unsqueeze(1)).squeeze(1)\n        true_probs = true_log_probs.exp()\n        nll = -true_log_probs\n        uniform_ce = -log_probs.mean(dim=1)\n\n        preds = probs.argmax(dim=1)\n        correct_mask = preds.eq(labels)\n        correct_conf = true_probs[correct_mask]\n\n        with torch.no_grad():\n            ema_source = correct_conf if correct_conf.numel() > 0 else true_probs\n            self._update_conf_ema(ema_source.mean())\n\n        if epoch < self.warmup_epochs:\n            ramp = logits.new_tensor(0.0)\n            target_conf = logits.new_tensor(self.init_target)\n        else:\n            steps = float(epoch - self.warmup_epochs + 1)\n            decay = torch.exp(logits.new_tensor(-steps / self.ramp_tau))\n            ramp = 1.0 - decay\n\n            ema_target = torch.clamp(\n                self.correct_conf_ema.to(logits.device) - self.conf_gap,\n                min=self.min_target,\n                max=self.max_target,\n            )\n            target_conf = ema_target + (self.init_target - ema_target) * decay\n\n        # Confidence-adaptive label smoothing: only examples whose true-class\n        # confidence exceeds the target receive extra smoothing.\n        excess_ratio = (\n            (true_probs.detach() - target_conf) / (1.0 - target_conf + 1e-6)\n        ).clamp(min=0.0, max=1.0)\n        adaptive_eps = self.max_smoothing * ramp * excess_ratio\n        adaptive_ce = ((1.0 - adaptive_eps) * nll + adaptive_eps * uniform_ce).mean()\n\n        # Additional cap on overconfident correct examples to directly shrink\n        # the train-confidence tail exploited by MIA.\n        over_mask = correct_mask & (true_probs.detach() > target_conf)\n        over_conf = true_probs[over_mask]\n        if over_conf.numel() > 0:\n            excess_pen = (over_conf - target_conf).pow(2).mean()\n            mean_pen = (over_conf.mean() - target_conf).clamp_min(0.0).pow(2)\n            over_entropy = -(probs[over_mask] * log_probs[over_mask]).sum(dim=1).mean()\n        else:\n            zero = logits.new_tensor(0.0)\n            excess_pen = zero\n            mean_pen = zero\n            over_entropy = zero\n\n        return adaptive_ce + ramp * (\n            self.excess_weight * excess_pen\n            + self.mean_weight * mean_pen\n            - self.entropy_weight * over_entropy\n        )\n# ============================================================\n# END EDITABLE\n# ============================================================\n",
5  "after_line": 0,
6  "start_line": 1,
7  "end_line": 32
8}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_membership_defense.py",
4  "content": "\"\"\"Editable membership-inference defense for MLS-Bench.\"\"\"\n\nimport math\n\nimport torch\nimport torch.nn.functional as F\n\n# ============================================================\n# EDITABLE\n# ============================================================\nclass MembershipDefense:\n    \"\"\"Adaptive confidence-capped self-distillation.\n\n    Core idea:\n    1) keep standard supervised learning on hard examples;\n    2) on already-easy / high-confidence examples, replace one-hot targets\n       with a tempered self-target so the model preserves class ranking while\n       reducing overconfident train-only predictions;\n    3) explicitly cap the upper tail of true-class confidence and shrink its\n       spread, which directly targets confidence-based membership inference.\n    \"\"\"\n\n    def __init__(self):\n        # Mild global smoothing keeps optimization stable.\n        self.base_smoothing = 0.02\n\n        # Adaptive self-distillation parameters.\n        self.alpha_max = 0.35\n        self.temperature = 2.2\n\n        # Confidence gate is applied on normalized true-class confidence so the\n        # same thresholds transfer across datasets with different class counts.\n        self.gate_center = 0.55\n        self.gate_width = 0.08\n\n        # Upper-tail confidence control.\n        self.cap_weight = 0.8\n        self.var_weight = 0.04\n\n    def compute_loss(self, logits, labels, epoch):\n        num_classes = logits.size(1)\n        dtype = logits.dtype\n\n        log_probs = F.log_softmax(logits, dim=1)\n        probs = log_probs.exp()\n\n        true_probs = probs.gather(1, labels.unsqueeze(1)).squeeze(1)\n        uniform_prob = 1.0 / float(num_classes)\n        norm_true = ((true_probs - uniform_prob) / max(1.0 - uniform_prob, 1e-6)).clamp(0.0, 1.0)\n\n        # Start close to ERM, then progressively regularize confident examples.\n        ramp = 1.0 - math.exp(-float(epoch + 1) / 20.0)\n        ramp_t = logits.new_tensor(ramp)\n\n        gate_center = self.gate_center + 0.10 * (1.0 - ramp)\n        gate = torch.sigmoid((norm_true.detach() - gate_center) / self.gate_width)\n\n        with torch.no_grad():\n            tempered = F.softmax(logits / self.temperature, dim=1)\n\n        one_hot = F.one_hot(labels, num_classes=num_classes).to(dtype=dtype)\n        uniform = torch.full_like(tempered, 1.0 / float(num_classes))\n        alpha = (self.alpha_max * ramp_t * gate).unsqueeze(1)\n\n        # Interpolate between one-hot supervision and a softer, model-aware\n        # target that preserves class similarities better than uniform label\n        # smoothing.\n        targets = (1.0 - self.base_smoothing) * ((1.0 - alpha) * one_hot + alpha * tempered) + self.base_smoothing * uniform\n        soft_ce = -(targets * log_probs).sum(dim=1)\n\n        # Confidence-based MIA relies on the train-only upper tail. Penalize\n        # only excessive confidence, not all confidence.\n        target_norm = 0.74 + 0.08 * ramp\n        cap_penalty = gate * F.relu(norm_true - target_norm).pow(2)\n\n        # Compress the spread of easy-example confidence to further reduce the\n        # separability exploited by confidence attacks.\n        weight_sum = gate.sum().clamp_min(1e-6)\n        mean_conf = (gate * norm_true.detach()).sum() / weight_sum\n        var_penalty = (gate * (norm_true - mean_conf).pow(2)).sum() / weight_sum\n\n        return (\n            soft_ce.mean()\n            + self.cap_weight * ramp_t * cap_penalty.mean()\n            + self.var_weight * ramp_t * var_penalty\n        )\n# ============================================================\n# END EDITABLE\n# ============================================================\n",
5  "after_line": 0,
6  "start_line": 1,
7  "end_line": 32
8}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_membership_defense.py",
4  "content": "\"\"\"Editable membership-inference defense for MLS-Bench.\"\"\"\n\nimport math\n\nimport torch\nimport torch.nn.functional as F\n\n# ============================================================\n# EDITABLE\n# ============================================================\nclass MembershipDefense:\n    \"\"\"Adaptive tempered flooding for privacy-utility tradeoffs.\n\n    The loss keeps standard supervised learning, but only once an example has\n    become \"easy\" (low loss relative to the random-guess loss log(K)) it adds:\n      1) tempered self-distillation to soften train-only overconfidence while\n         preserving class similarity structure; and\n      2) a smooth flooding barrier that prevents the loss from collapsing far\n         below a class-count-aware target, directly shrinking the signal used by\n         confidence/loss-based membership inference.\n    \"\"\"\n\n    def __init__(self):\n        self.base_smoothing = 0.01\n        self.temperature = 2.4\n        self.distill_weight = 0.80\n        self.floor_weight = 0.90\n        self.spread_weight = 0.02\n\n        # Gate over easy examples using loss normalized by log(num_classes).\n        self.gate_center = 0.58\n        self.gate_width = 0.10\n\n    def compute_loss(self, logits, labels, epoch):\n        num_classes = logits.size(1)\n        log_k = math.log(float(num_classes))\n\n        log_probs = F.log_softmax(logits, dim=1)\n        hard_ce = F.cross_entropy(logits, labels, reduction=\"none\")\n        base_ce = F.cross_entropy(\n            logits,\n            labels,\n            reduction=\"none\",\n            label_smoothing=self.base_smoothing,\n        )\n\n        # Start near ERM and gradually turn on privacy regularization.\n        ramp = 1.0 - math.exp(-float(epoch + 1) / 15.0)\n        ramp_t = logits.new_tensor(ramp)\n\n        # Normalized skill = improvement over random-guess loss, in [0, 1].\n        skill = (1.0 - hard_ce.detach() / log_k).clamp(0.0, 1.0)\n        gate_center = self.gate_center + 0.08 * (1.0 - ramp)\n        gate = torch.sigmoid((skill - gate_center) / self.gate_width)\n\n        with torch.no_grad():\n            tempered = F.softmax(logits / self.temperature, dim=1)\n\n        distill = F.kl_div(log_probs, tempered, reduction=\"none\").sum(dim=1)\n\n        # Smooth, class-count-aware flooding barrier: allow learning on hard\n        # samples, but discourage memorized train samples from driving the loss\n        # much below a moderate target.\n        target_ce = log_k * (0.24 + 0.05 * ramp)\n        flood = gate * F.relu(target_ce - hard_ce).pow(2)\n\n        # Compress the easy-sample loss spread; confidence attacks exploit the\n        # particularly low-loss tail of the train distribution.\n        weight_sum = gate.sum().clamp_min(1e-6)\n        mean_ce = (gate * hard_ce.detach()).sum() / weight_sum\n        spread = (gate * (hard_ce - mean_ce).pow(2)).sum() / weight_sum\n\n        return (\n            base_ce.mean()\n            + self.distill_weight * ramp_t * (gate * distill).mean()\n            + self.floor_weight * ramp_t * flood.mean()\n            + self.spread_weight * ramp_t * spread\n        )\n# ============================================================\n# END EDITABLE\n# ============================================================\n",
5  "after_line": 0,
6  "start_line": 1,
7  "end_line": 89
8}

1{
2  "n": 1
3}