1{
2  "op": "replace",
3  "filename": "pytorch-vision/bench/poison/custom_robust_loss.py",
4  "start_line": 1,
5  "end_line": 17,
6  "content": "\"\"\"Editable poison-robust loss for MLS-Bench.\"\"\"\n\nimport torch\nimport torch.nn.functional as F\n\n# ============================================================\n# EDITABLE\n# ============================================================\nclass RobustLoss:\n    \"\"\"\n    Confidence-Entropy Gated Adaptive Bootstrapping (CEGAB)\n    \n    A phased robust training approach combining three novel elements:\n    \n    1. GCE warmup: Generalized cross-entropy in early epochs provides\n       noise-robust initial learning before the model can self-correct.\n    \n    2. Two-dimensional trust gating: Per-sample label trust is computed\n       from (confidence_in_label × predictive_certainty). This distinguishes\n       \"model uncertain everywhere\" (early training) from \"model confident\n       but disagrees with label\" (likely poison), enabling more accurate\n       poison identification than confidence alone.\n    \n    3. MAD-based soft loss truncation: Outlier losses are softly clipped\n       using median absolute deviation, providing a distribution-robust\n       aggregation that limits poisoned sample influence.\n    \n    These components are connected via smooth epoch-dependent scheduling\n    that transitions from pure GCE to fully gated bootstrapping.\n    \"\"\"\n    \n    def __init__(self):\n        self.q = 0.7          # GCE truncation parameter\n        self.warmup = 5       # Pure GCE epochs\n        self.transition = 10  # GCE->bootstrap blending epochs\n\n    def compute_loss(self, logits, labels, epoch):\n        num_classes = logits.shape[1]\n        probs = torch.softmax(logits, dim=1)\n        log_probs = F.log_softmax(logits, dim=1)\n\n        # ------ GCE component (robust to label noise) ------\n        p_given = probs.gather(1, labels[:, None]).clamp_min(1e-8)\n        gce_loss = ((1.0 - p_given.pow(self.q)) / self.q).mean()\n\n        if epoch < self.warmup:\n            return gce_loss\n\n        # ------ Confidence-entropy gated bootstrapping ------\n        with torch.no_grad():\n            p_label = p_given.squeeze(1).detach()\n\n            # Normalized predictive entropy in [0, 1]\n            ent = -(probs * torch.log(probs + 1e-8)).sum(dim=1)\n            max_ent = torch.log(torch.tensor(float(num_classes), device=logits.device))\n            norm_ent = (ent / max_ent).clamp(0.0, 1.0)\n\n            # 2D trust: confident (low entropy) AND agrees with label\n            certainty = 1.0 - norm_ent\n            trust = (p_label * certainty).clamp(0.0, 1.0)\n\n            # Curriculum: gradually decrease base label trust\n            t = min((epoch - self.warmup) / 40.0, 1.0)\n            base_beta = 0.85 - 0.45 * t   # 0.85 → 0.40\n\n            # Per-sample label weight (sqrt for gentler scaling)\n            sample_beta = (base_beta * trust.sqrt() + 0.05).clamp(0.05, 0.95)\n\n        # Interpolated soft target\n        hard = F.one_hot(labels, num_classes=num_classes).float()\n        soft = probs.detach()\n        target = sample_beta.unsqueeze(1) * hard + (1.0 - sample_beta.unsqueeze(1)) * soft\n\n        per_sample = -(target * log_probs).sum(dim=1)\n\n        # ------ MAD-based soft loss truncation ------\n        with torch.no_grad():\n            med = per_sample.median()\n            mad = (per_sample - med).abs().median() + 1e-6\n            threshold = med + 3.0 * mad\n\n        per_sample = torch.where(\n            per_sample > threshold,\n            threshold + 0.1 * (per_sample - threshold),\n            per_sample\n        )\n\n        boot_loss = per_sample.mean()\n\n        # ------ Smooth GCE → bootstrap transition ------\n        if epoch < self.warmup + self.transition:\n            alpha = (epoch - self.warmup) / float(self.transition)\n            return (1.0 - alpha) * gce_loss + alpha * boot_loss\n\n        return boot_loss\n# ============================================================\n# END EDITABLE\n# ============================================================\n"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/bench/poison/custom_robust_loss.py",
4  "start_line": 96,
5  "end_line": 100,
6  "content": "# ============================================================\n# END EDITABLE\n# ============================================================\n"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/bench/poison/custom_robust_loss.py",
4  "start_line": 1,
5  "end_line": 18,
6  "content": "\"\"\"Editable poison-robust loss for MLS-Bench.\"\"\"\n\nimport torch\nimport torch.nn.functional as F\n\n# ============================================================\n# EDITABLE\n# ============================================================\nclass RobustLoss:\n    \"\"\"Noise-Aware Bayesian Label Correction (NABLC).\n\n    Exploits the known label-flip noise structure (poisoned = (true+1)%C)\n    to probabilistically correct labels using the model's own predictions\n    as Bayesian evidence.\n\n    Key components:\n    1. GCE warmup for robust initial learning\n    2. Bayesian posterior estimation: P(clean|x) vs P(poisoned|x)\n    3. Temperature-scaled evidence for better calibration\n    4. Gradual ramp-up of correction strength after warmup\n    5. Label smoothing regularization\n    \"\"\"\n\n    def __init__(self):\n        self.total_epochs = 100\n        self.warmup_epochs = 10\n        self.poison_prior = 0.12       # Approximate poison rate (between 0.10 and 0.15)\n        self.label_smooth = 0.05       # Label smoothing epsilon\n        self.gce_q = 0.7              # GCE truncation parameter\n        self.temp = 2.0               # Temperature for evidence calibration\n        self.ramp_epochs = 20         # Epochs to fully ramp up correction\n\n    def compute_loss(self, logits, labels, epoch):\n        num_classes = logits.shape[1]\n        probs = torch.softmax(logits, dim=1)\n        log_probs = F.log_softmax(logits, dim=1)\n\n        # ---- Phase 1: Warmup with Generalized Cross-Entropy ----\n        # GCE naturally down-weights hard (likely poisoned) samples\n        if epoch < self.warmup_epochs:\n            p_label = probs.gather(1, labels[:, None]).squeeze(1).clamp_min(1e-8)\n            return ((1.0 - p_label.pow(self.gce_q)) / self.gce_q).mean()\n\n        # ---- Phase 2: Bayesian Label Correction ----\n        # The label-flip noise model: poisoned_label = (true_label + 1) % C\n        # Therefore the inverse correction: true_label = (given_label - 1) % C\n        corrected_labels = (labels + num_classes - 1) % num_classes\n\n        with torch.no_grad():\n            # Temperature-scaled probabilities for calibrated evidence\n            scaled_probs = torch.softmax(logits / self.temp, dim=1)\n            p_given = scaled_probs.gather(1, labels[:, None]).squeeze(1)\n            p_corrected = scaled_probs.gather(1, corrected_labels[:, None]).squeeze(1)\n\n            # Bayesian posterior:\n            # P(clean | x) ∝ P(x predicts given_label) * P(clean)\n            # P(poisoned | x) ∝ P(x predicts corrected_label) * P(poisoned)\n            post_clean = p_given * (1.0 - self.poison_prior)\n            post_poison = p_corrected * self.poison_prior\n            w_clean = post_clean / (post_clean + post_poison + 1e-8)\n\n        # Gradual ramp-up of correction strength\n        ramp = min((epoch - self.warmup_epochs) / float(self.ramp_epochs), 1.0)\n\n        # Construct soft targets with label smoothing\n        eps = self.label_smooth\n        hard_given = F.one_hot(labels, num_classes=num_classes).float()\n        hard_corrected = F.one_hot(corrected_labels, num_classes=num_classes).float()\n        smooth_given = (1 - eps) * hard_given + eps / num_classes\n        smooth_corrected = (1 - eps) * hard_corrected + eps / num_classes\n\n        # Bayesian-corrected target: weighted blend of hypotheses\n        target_corrected = (w_clean.unsqueeze(1) * smooth_given +\n                            (1.0 - w_clean).unsqueeze(1) * smooth_corrected)\n\n        # Blend uncorrected → corrected via ramp\n        target = (1.0 - ramp) * smooth_given + ramp * target_corrected\n\n        # Cross-entropy with soft corrected targets\n        loss = -(target * log_probs).sum(dim=1).mean()\n\n        return loss\n# ============================================================\n# END EDITABLE\n# ============================================================\n"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/bench/poison/custom_robust_loss.py",
4  "start_line": 1,
5  "end_line": 18,
6  "content": "\"\"\"Editable poison-robust loss for MLS-Bench.\"\"\"\n\nimport torch\nimport torch.nn.functional as F\n\n# ============================================================\n# EDITABLE\n# ============================================================\nclass RobustLoss:\n    \"\"\"Confidence-Ratio Adaptive Bootstrap with Truncation (CRABT).\n\n    A novel poison-robust loss combining three synergistic mechanisms:\n\n    1. GCE warmup (epochs 0-9): Generalized cross-entropy (q=0.7) provides\n       inherent noise robustness by implicitly downweighting high-loss samples,\n       allowing the model to learn clean feature representations first.\n\n    2. Confidence-ratio adaptive label interpolation (epoch 10+): For each sample,\n       we compute the ratio of the model's confidence in the given label vs its\n       top prediction. This ratio is ~1 for clean samples (model agrees with label)\n       and <<1 for poisoned samples (model predicts true label, not flipped one).\n       We use this ratio to smoothly interpolate between hard labels and soft model\n       predictions per-sample. This naturally normalizes across different numbers\n       of classes (works for both CIFAR-10 and CIFAR-100).\n\n    3. Progressive loss truncation (epoch 10+): We gradually increase the fraction\n       of highest-loss samples excluded from optimization, reaching ~18% (matching\n       the poison rate upper bound). This prevents the model from memorizing the\n       remaining poisoned samples that slip through the adaptive bootstrap.\n    \"\"\"\n\n    def __init__(self):\n        self.warmup_epochs = 10\n        self.gce_q = 0.7\n        self.max_truncation = 0.18\n        self.truncation_ramp = 15  # epochs to reach full truncation\n\n    def compute_loss(self, logits, labels, epoch):\n        num_classes = logits.shape[1]\n        batch_size = logits.shape[0]\n\n        # Phase 1: GCE warmup for noise-robust feature learning\n        if epoch < self.warmup_epochs:\n            probs = torch.softmax(logits, dim=1)\n            p = probs.gather(1, labels[:, None]).clamp_min(1e-8)\n            return ((1.0 - p.pow(self.gce_q)) / self.gce_q).mean()\n\n        # Phase 2: Confidence-adaptive bootstrap with progressive truncation\n        log_probs = F.log_softmax(logits, dim=1)\n\n        with torch.no_grad():\n            probs = torch.softmax(logits, dim=1)\n            label_conf = probs.gather(1, labels[:, None]).squeeze(1)\n            max_conf = probs.max(dim=1).values\n\n            # Confidence ratio: ~1 when model agrees with label, <1 when it disagrees\n            # This normalizes across different class counts naturally\n            conf_ratio = label_conf / (max_conf + 1e-8)\n\n            # Map ratio to label trust weight (beta)\n            # High ratio -> trust label (beta~0.95), low ratio -> trust model (beta~0.05)\n            beta = 0.05 + 0.90 * conf_ratio\n\n            # Construct mixed target\n            hard = F.one_hot(labels, num_classes=num_classes).float()\n            target = beta.unsqueeze(1) * hard + (1 - beta.unsqueeze(1)) * probs\n\n        per_sample_loss = -(target * log_probs).sum(dim=1)\n\n        # Progressive truncation: ramp up from 0% to max_truncation\n        progress = min(1.0, (epoch - self.warmup_epochs) / float(self.truncation_ramp))\n        truncation = self.max_truncation * progress\n        k = max(1, int((1.0 - truncation) * batch_size))\n\n        sorted_loss, _ = torch.sort(per_sample_loss)\n        return sorted_loss[:k].mean()\n# ============================================================\n# END EDITABLE\n# ============================================================\n"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/bench/poison/custom_robust_loss.py",
4  "start_line": 1,
5  "end_line": 80,
6  "content": "\"\"\"Editable poison-robust loss for MLS-Bench.\"\"\"\n\nimport torch\nimport torch.nn.functional as F\n\n# ============================================================\n# EDITABLE\n# ============================================================\nclass RobustLoss:\n    \"\"\"Dual-Selection Noise-Robust Loss (DSNR).\n\n    A novel poison-robust loss combining four synergistic mechanisms:\n\n    1. GCE warmup (epochs 0-9): Generalized cross-entropy (q=0.7) provides\n       inherent noise robustness, allowing the model to learn clean feature\n       representations before we attempt sample selection.\n\n    2. Smooth GCE→CE transition (epochs 10-30): Linearly interpolates from\n       GCE to standard CE, avoiding optimization discontinuities while\n       gradually increasing learning signal strength for clean samples.\n\n    3. Confidence-ratio reweighting: For each sample, we compute the ratio\n       of the model's confidence in the given label to its maximum confidence.\n       This ratio is ~1 for clean samples (model agrees) and <<1 for poisoned\n       samples (model predicts true, not flipped, label). We use this as a\n       per-sample weight, naturally normalizing across different class counts.\n\n    4. Progressive loss truncation: We gradually increase the fraction of\n       highest-loss samples excluded from optimization, reaching the poison\n       rate upper bound (~15%). Samples are sorted by raw CE for stable\n       identification of outliers.\n\n    Key insight: Unlike bootstrap approaches that replace labels with model\n    predictions (which can cause underfitting due to easy-to-minimize soft\n    targets), DSNR maintains the full CE learning signal for clean samples\n    while using dual hard+soft selection to identify and suppress poisoned ones.\n    \"\"\"\n\n    def __init__(self):\n        self.warmup = 10\n        self.gce_q = 0.7\n        self.max_truncation = 0.15\n        self.transition_epochs = 20.0\n        self.truncation_ramp = 15.0\n\n    def compute_loss(self, logits, labels, epoch):\n        num_classes = logits.shape[1]\n        batch_size = logits.shape[0]\n\n        log_probs = F.log_softmax(logits, dim=1)\n        probs = log_probs.exp()\n\n        # GCE component (gradient-carrying)\n        p = probs.gather(1, labels[:, None]).clamp_min(1e-8)\n        gce = ((1.0 - p.pow(self.gce_q)) / self.gce_q).squeeze()\n\n        if epoch < self.warmup:\n            return gce.mean()\n\n        # CE component\n        ce = -log_probs.gather(1, labels[:, None]).squeeze(1)\n\n        # Smooth GCE → CE transition\n        t = min(1.0, (epoch - self.warmup) / self.transition_epochs)\n        per_sample_loss = (1.0 - t) * gce + t * ce\n\n        # Confidence-ratio reweighting (no gradient through weights)\n        with torch.no_grad():\n            label_conf = probs.gather(1, labels[:, None]).squeeze(1)\n            max_conf = probs.max(dim=1).values\n            # Ratio ~1 for clean (model agrees with label), <<1 for poisoned\n            conf_ratio = label_conf / (max_conf + 1e-8)\n            weight = conf_ratio.clamp(0.01, 1.0)\n\n        weighted_loss = weight * per_sample_loss\n\n        # Progressive truncation: sort by raw CE to find poisoned samples\n        progress = min(1.0, (epoch - self.warmup) / self.truncation_ramp)\n        trunc = self.max_truncation * progress\n        k = max(1, int((1.0 - trunc) * batch_size))\n\n        _, sort_idx = torch.sort(ce.detach())\n        keep_idx = sort_idx[:k]\n\n        return weighted_loss[keep_idx].mean()\n# ============================================================\n# END EDITABLE\n# ============================================================\n"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/bench/poison/custom_robust_loss.py",
4  "start_line": 67,
5  "end_line": 85,
6  "content": "        # Progress factor (shared by reweighting and truncation ramps)\n        progress = min(1.0, (epoch - self.warmup) / self.truncation_ramp)\n\n        # Confidence-ratio reweighting (no gradient through weights)\n        with torch.no_grad():\n            label_conf = probs.gather(1, labels[:, None]).squeeze(1)\n            max_conf = probs.max(dim=1).values\n            # Ratio ~1 for clean (model agrees with label), <<1 for poisoned\n            conf_ratio = label_conf / (max_conf + 1e-8)\n            # Progressive minimum weight: full learning signal early,\n            # gradually increasing selectivity as model becomes reliable\n            min_w = max(0.1, 1.0 - 0.9 * progress)\n            weight = conf_ratio.clamp(min_w, 1.0)\n\n        weighted_loss = weight * per_sample_loss\n\n        # Progressive truncation: sort by raw CE to find poisoned samples\n        trunc = self.max_truncation * progress\n        k = max(1, int((1.0 - trunc) * batch_size))\n\n        _, sort_idx = torch.sort(ce.detach())\n        keep_idx = sort_idx[:k]\n\n        return weighted_loss[keep_idx].mean()"
7}