1{
2  "op": "replace",
3  "filename": "pytorch-vision/bench/poison/custom_robust_loss.py",
4  "start_line": 1,
5  "end_line": 19,
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    \"\"\"Adaptive label smoothing with intra-class consistency regularization.\"\"\"\n\n    def __init__(self):\n        # Hyperparameters\n        self.ema_decay = 0.9  # decay for per-class prototype moving average\n        self.lambda_consistency = 0.5  # weight for consistency loss\n        self.smooth_eps = 0.1  # base label smoothing epsilon\n        self.adaptive_factor = 0.5  # scaling factor for adaptive smoothing\n        \n        # State variables\n        self.num_classes = None\n        self.class_prototypes = None  # shape (num_classes, num_classes)\n        self.initialized = False\n        \n    def _initialize(self, num_classes):\n        self.num_classes = num_classes\n        # Initialize prototypes as uniform distributions\n        self.class_prototypes = torch.ones(num_classes, num_classes) / num_classes\n        self.initialized = True\n\n    def compute_loss(self, logits, labels, epoch):\n        batch_size, num_classes = logits.shape\n        if not self.initialized:\n            self._initialize(num_classes)\n        \n        device = logits.device\n        self.class_prototypes = self.class_prototypes.to(device)\n        \n        # Convert labels to one-hot\n        one_hot = F.one_hot(labels, num_classes=num_classes).float()\n        \n        # Compute softmax probabilities\n        probs = torch.softmax(logits, dim=1)\n        \n        # Update per-class prototypes using exponential moving average\n        with torch.no_grad():\n            for c in range(num_classes):\n                mask = (labels == c)\n                if mask.any():\n                    class_probs = probs[mask]\n                    mean_prob = class_probs.mean(dim=0)\n                    # Update EMA\n                    self.class_prototypes[c] = (\n                        self.ema_decay * self.class_prototypes[c] +\n                        (1 - self.ema_decay) * mean_prob\n                    )\n        \n        # Intra-class consistency loss: KL divergence between sample probs and its class prototype\n        # Gather prototype for each sample\n        prototypes = self.class_prototypes[labels]  # shape (batch_size, num_classes)\n        # Add small epsilon to avoid log(0)\n        prototypes = prototypes.clamp_min(1e-8)\n        probs_clamped = probs.clamp_min(1e-8)\n        consistency_loss = (probs_clamped * (torch.log(probs_clamped) - torch.log(prototypes))).sum(dim=1).mean()\n        \n        # Adaptive label smoothing\n        # Confidence: probability assigned to the given label\n        confidence = probs.gather(1, labels.unsqueeze(1)).squeeze(1)  # shape (batch_size,)\n        # Smoothing factor: lower confidence -> more smoothing\n        alpha = (1 - confidence) * self.adaptive_factor\n        alpha = alpha.clamp(max=self.smooth_eps)  # cap at base smoothing epsilon\n        \n        # Smoothed target: mix of one-hot and uniform distribution\n        uniform = torch.ones_like(one_hot) / num_classes\n        # Expand alpha to shape (batch_size, 1) for broadcasting\n        alpha_expanded = alpha.unsqueeze(1)\n        target = (1 - alpha_expanded) * one_hot + alpha_expanded * uniform\n        \n        # Cross-entropy with smoothed targets\n        log_probs = F.log_softmax(logits, dim=1)\n        ce_loss = -(target * log_probs).sum(dim=1).mean()\n        \n        # Total loss\n        total_loss = ce_loss + self.lambda_consistency * consistency_loss\n        \n        return total_loss\n# ============================================================\n# END EDITABLE\n# ============================================================"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/bench/poison/custom_robust_loss.py",
4  "start_line": 1,
5  "end_line": 88,
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    \"\"\"Adaptive Bootstrap with Confidence Weighting and Entropy Regularization.\"\"\"\n\n    def __init__(self):\n        # Hyperparameters\n        self.entropy_weight = 0.1  # weight for entropy regularization\n        self.confidence_threshold = 0.3  # threshold for high-confidence samples\n        self.temperature = 1.0  # temperature for soft targets\n        \n    def compute_loss(self, logits, labels, epoch):\n        batch_size, num_classes = logits.shape\n        \n        # Compute softmax probabilities with temperature\n        probs = torch.softmax(logits / self.temperature, dim=1)\n        # Compute confidence (probability of given label)\n        confidence = probs.gather(1, labels.unsqueeze(1)).squeeze(1)  # shape (batch_size,)\n        \n        # Adaptive bootstrap weight: beta = confidence, clipped to [0,1]\n        beta = confidence.detach()  # shape (batch_size,)\n        # Ensure beta is at least a small value to avoid zero weight on hard label\n        beta = beta.clamp(min=0.1, max=0.9)\n        \n        # Hard targets (one-hot)\n        hard = F.one_hot(labels, num_classes=num_classes).float()\n        # Soft targets (model predictions detached)\n        soft = probs.detach()\n        \n        # Interpolated targets per sample\n        beta_expanded = beta.unsqueeze(1)  # shape (batch_size, 1)\n        target = beta_expanded * hard + (1 - beta_expanded) * soft\n        \n        # Cross-entropy loss with interpolated targets\n        log_probs = F.log_softmax(logits, dim=1)\n        ce_loss = -(target * log_probs).sum(dim=1).mean()\n        \n        # Entropy regularization: encourage high confidence for samples with high given-label confidence\n        # Only apply to samples where confidence > threshold\n        mask = (confidence > self.confidence_threshold).float()\n        entropy = -(probs * torch.log(probs.clamp_min(1e-8))).sum(dim=1)\n        entropy_loss = (mask * entropy).mean()\n        \n        # Total loss\n        total_loss = ce_loss + self.entropy_weight * entropy_loss\n        \n        return total_loss\n# ============================================================\n# END EDITABLE\n# ============================================================"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/bench/poison/custom_robust_loss.py",
4  "start_line": 1,
5  "end_line": 56,
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    \"\"\"Dynamic Label Smoothing with Class Prototype Targets.\"\"\"\n\n    def __init__(self):\n        # Hyperparameters\n        self.ema_decay = 0.9  # for prototype updates\n        self.initial_alpha = 0.8  # initial weight on prototype\n        self.final_alpha = 0.0   # final weight on prototype\n        self.total_epochs = 5    # assumed total epochs (training runs 5 epochs)\n        self.entropy_weight = 0.05  # entropy regularization weight\n        \n        # State\n        self.num_classes = None\n        self.class_prototypes = None  # shape (num_classes, num_classes)\n        self.initialized = False\n        \n    def _initialize(self, num_classes):\n        self.num_classes = num_classes\n        # Initialize prototypes as uniform distributions\n        self.class_prototypes = torch.ones(num_classes, num_classes) / num_classes\n        self.initialized = True\n\n    def compute_loss(self, logits, labels, epoch):\n        batch_size, num_classes = logits.shape\n        if not self.initialized:\n            self._initialize(num_classes)\n        \n        device = logits.device\n        self.class_prototypes = self.class_prototypes.to(device)\n        \n        # Compute softmax probabilities\n        probs = torch.softmax(logits, dim=1)\n        \n        # Update per-class prototypes using exponential moving average\n        with torch.no_grad():\n            for c in range(num_classes):\n                mask = (labels == c)\n                if mask.any():\n                    class_probs = probs[mask]\n                    mean_prob = class_probs.mean(dim=0)\n                    # Update EMA\n                    self.class_prototypes[c] = (\n                        self.ema_decay * self.class_prototypes[c] +\n                        (1 - self.ema_decay) * mean_prob\n                    )\n        \n        # Schedule alpha: linear decay from initial_alpha to final_alpha over total_epochs\n        alpha = self.initial_alpha - (self.initial_alpha - self.final_alpha) * (epoch / (self.total_epochs - 1))\n        alpha = max(alpha, 0.0)\n        \n        # Prototype targets for each sample\n        prototypes = self.class_prototypes[labels]  # shape (batch_size, num_classes)\n        # Hard targets\n        hard = F.one_hot(labels, num_classes=num_classes).float()\n        \n        # Combined target\n        target = (1 - alpha) * hard + alpha * prototypes\n        \n        # Cross-entropy loss\n        log_probs = F.log_softmax(logits, dim=1)\n        ce_loss = -(target * log_probs).sum(dim=1).mean()\n        \n        # Entropy regularization: encourage confident predictions\n        entropy = -(probs * torch.log(probs.clamp_min(1e-8))).sum(dim=1).mean()\n        \n        # Total loss\n        total_loss = ce_loss + self.entropy_weight * entropy\n        \n        return total_loss\n# ============================================================\n# END EDITABLE\n# ============================================================"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/bench/poison/custom_robust_loss.py",
4  "start_line": 1,
5  "end_line": 19,
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    \"\"\"Adaptive Robust Loss (ARL): softmin-weighted cross-entropy.\n    \n    Weights samples by exp(-loss_i / temperature) where loss_i is per-sample CE.\n    Temperature decays with epoch to gradually focus on low-loss (clean) samples.\n    \"\"\"\n\n    def __init__(self):\n        # initial temperature (high = uniform weights)\n        self.t0 = 2.0\n        # final temperature after 100 epochs\n        self.t_min = 0.1\n        # decay factor per epoch\n        self.decay = (self.t0 / self.t_min) ** (1.0 / 100)\n\n    def compute_loss(self, logits, labels, epoch):\n        # per-sample cross-entropy\n        loss_per_sample = F.cross_entropy(logits, labels, reduction='none')\n        # temperature for current epoch\n        temperature = self.t0 * (self.decay ** (-epoch))\n        temperature = max(temperature, self.t_min)\n        # compute weights: softmin over losses\n        weights = F.softmax(-loss_per_sample.detach() / temperature, dim=0)\n        # weighted loss\n        loss = (weights * loss_per_sample).sum()\n        return loss\n# ============================================================\n# END EDITABLE\n# ============================================================"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/bench/poison/custom_robust_loss.py",
4  "start_line": 1,
5  "end_line": 19,
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    \"\"\"Adaptive Confidence-based Label Smoothing (ACLS).\"\"\"\n\n    def __init__(self):\n        self.num_classes = None\n        self.eps = 1e-8\n        # smoothing schedule: start with low smoothing, increase later\n        self.smoothing_start = 0.1\n        self.smoothing_end = 0.5\n        # confidence threshold for weighting\n        self.conf_threshold = 0.3\n\n    def compute_loss(self, logits, labels, epoch):\n        if self.num_classes is None:\n            self.num_classes = logits.shape[1]\n        \n        # compute confidence\n        probs = torch.softmax(logits, dim=1)\n        confidence = probs.gather(1, labels.unsqueeze(1)).squeeze(1)  # shape (batch,)\n        \n        # adaptive smoothing factor per sample based on confidence\n        # low confidence -> higher smoothing\n        smoothing_factor = (1.0 - confidence).clamp(min=self.smoothing_start, max=self.smoothing_end)\n        \n        # schedule: increase smoothing factor as training progresses\n        epoch_factor = min(1.0, epoch / 100.0)  # assuming 100 epochs total\n        smoothing_factor = smoothing_factor * epoch_factor\n        \n        # create smoothed labels\n        one_hot = F.one_hot(labels, num_classes=self.num_classes).float()\n        smoothing_matrix = smoothing_factor.unsqueeze(1) / self.num_classes\n        smoothed_labels = one_hot * (1.0 - smoothing_factor.unsqueeze(1)) + smoothing_matrix\n        \n        # compute cross-entropy with smoothed labels\n        log_probs = F.log_softmax(logits, dim=1)\n        loss = -(smoothed_labels * log_probs).sum(dim=1).mean()\n        return loss\n# ============================================================\n# END EDITABLE\n# ============================================================"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/bench/poison/custom_robust_loss.py",
4  "start_line": 1,
5  "end_line": 19,
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    \"\"\"Adaptive Bootstrap with Confidence Thresholding (ABCT).\"\"\"\n\n    def __init__(self):\n        # Base mixing parameter for bootstrap\n        self.beta_base = 0.8\n        # Confidence threshold starts low, increases over epochs\n        self.conf_thresh_start = 0.7\n        self.conf_thresh_end = 0.95\n        # Number of total epochs (assuming 100 as in description)\n        self.total_epochs = 100\n\n    def compute_loss(self, logits, labels, epoch):\n        # Current confidence threshold linear schedule\n        conf_thresh = self.conf_thresh_start + (self.conf_thresh_end - self.conf_thresh_start) * (epoch / self.total_epochs)\n        \n        probs = torch.softmax(logits, dim=1)\n        conf, preds = torch.max(probs, dim=1)\n        \n        # Determine per-sample beta: low if high confidence and prediction != label (suspicious)\n        beta = torch.ones_like(conf) * self.beta_base\n        suspicious = (conf > conf_thresh) & (preds != labels)\n        beta[suspicious] = 0.2  # trust predictions more for suspicious samples\n        \n        # Bootstrap target mixing\n        hard = F.one_hot(labels, num_classes=logits.shape[1]).float()\n        soft = torch.softmax(logits.detach(), dim=1)\n        # Expand beta to match dimensions\n        beta_expanded = beta.view(-1, 1)\n        target = beta_expanded * hard + (1.0 - beta_expanded) * soft\n        \n        log_probs = F.log_softmax(logits, dim=1)\n        return -(target * log_probs).sum(dim=1).mean()\n# ============================================================\n# END EDITABLE\n# ============================================================"
7}