1{
2  "op": "replace",
3  "filename": "opacus/custom_dpsgd.py",
4  "start_line": 152,
5  "end_line": 233,
6  "content": "# =====================================================================\n# EDITABLE SECTION START (lines 152-233)\n# =====================================================================\n# Centered Adaptive DP-SGD: Control variate centering + adaptive clipping.\n# Novel contribution: subtract EMA of past privatized gradients as a\n# control variate before clipping. This reduces residual norms, enabling\n# tighter adaptive clipping and proportionally less absolute noise while\n# maintaining identical privacy guarantees (same noise multiplier sigma).\n\nclass DPMechanism:\n    \"\"\"DP mechanism with control variate centering and adaptive clipping.\n\n    The EMA of past privatized gradients is public (post-processing of\n    already-released data), so subtracting it doesn't affect sensitivity.\n    Smaller residuals -> smaller adaptive clip -> less noise -> better SNR.\n    \"\"\"\n\n    def __init__(self, max_grad_norm, noise_multiplier, n_params,\n                 dataset_size, batch_size, epochs, target_epsilon, target_delta):\n        self.max_grad_norm = max_grad_norm\n        self.noise_multiplier = noise_multiplier\n        self.n_params = n_params\n        self.dataset_size = dataset_size\n        self.batch_size = batch_size\n        self.epochs = epochs\n        self.target_epsilon = target_epsilon\n        self.target_delta = target_delta\n        # Control variate: EMA of past privatized gradients\n        self.grad_ema = None\n        self.ema_beta = 0.9\n        # Adaptive clipping for residuals (targets median norm)\n        self.clip_norm = max_grad_norm\n        self.target_quantile = 0.5\n        self.clip_lr = 0.2\n\n    def clip_and_noise(self, per_sample_grads, step, epoch):\n        batch_size = per_sample_grads[0].shape[0]\n        # Initialize EMA on first call\n        if self.grad_ema is None:\n            self.grad_ema = [torch.zeros_like(g[0]) for g in per_sample_grads]\n        # Center per-sample grads by subtracting EMA (free post-processing)\n        residuals = [g - e.unsqueeze(0)\n                     for g, e in zip(per_sample_grads, self.grad_ema)]\n        # Per-sample residual norms\n        flat = torch.cat([r.reshape(batch_size, -1) for r in residuals], dim=1)\n        norms = flat.norm(2, dim=1)\n        # Adapt clip threshold to target quantile of residual norms\n        frac = (norms > self.clip_norm).float().mean().item()\n        self.clip_norm *= math.exp(self.clip_lr * (frac - self.target_quantile))\n        self.clip_norm = max(0.005, min(50.0, self.clip_norm))\n        # Clip residuals (sensitivity = clip_norm)\n        clip_factor = (self.clip_norm / norms.clamp(min=1e-8)).clamp(max=1.0)\n        noised_grads = []\n        for r, e in zip(residuals, self.grad_ema):\n            shape = [batch_size] + [1] * (r.dim() - 1)\n            clipped = r * clip_factor.reshape(shape)\n            avg = clipped.mean(dim=0)\n            # Noise calibrated to residual clip norm (same sigma ratio)\n            noise = torch.randn_like(avg) * (\n                self.noise_multiplier * self.clip_norm / batch_size)\n            # Reconstruct full gradient: EMA + noised residual\n            noised_grads.append(e + avg + noise)\n        # Update EMA from privatized outputs (post-processing, free)\n        with torch.no_grad():\n            for i in range(len(noised_grads)):\n                self.grad_ema[i] = (self.ema_beta * self.grad_ema[i]\n                    + (1 - self.ema_beta) * noised_grads[i].detach())\n        return noised_grads\n\n    def get_effective_sigma(self, step, epoch):\n        return self.noise_multiplier\n\n# =====================================================================\n# EDITABLE SECTION END\n# ====================================================================="
7}

1{
2  "op": "replace",
3  "filename": "opacus/custom_dpsgd.py",
4  "start_line": 152,
5  "end_line": 233,
6  "content": "# =====================================================================\n# EDITABLE SECTION START (lines 152-233)\n# =====================================================================\n# Adaptive Control Variate DP-SGD (ACV-DPSGD): Uses EMA of noised outputs\n# as control variates to center gradients before clipping, enabling\n# smaller adaptive thresholds and proportionally less noise.\n\nclass DPMechanism:\n    \"\"\"Adaptive Control Variate DP-SGD.\n\n    Centers per-sample gradients around an EMA control variate before\n    clipping, yielding smaller residual norms. Adaptively tracks median\n    residual norm as clipping threshold. Noise is proportional to this\n    smaller threshold. Privacy guarantee is unchanged (sigma-GDP per step\n    since noise/sensitivity ratio = sigma, independent of clipping norm).\n    \"\"\"\n\n    def __init__(self, max_grad_norm, noise_multiplier, n_params,\n                 dataset_size, batch_size, epochs, target_epsilon, target_delta):\n        self.max_grad_norm = max_grad_norm\n        self.noise_multiplier = noise_multiplier\n        self.n_params = n_params\n        self.dataset_size = dataset_size\n        self.batch_size = batch_size\n        self.epochs = epochs\n        self.target_epsilon = target_epsilon\n        self.target_delta = target_delta\n        self.control_variate = None  # Per-param EMA tensors, init on first call\n        self.cv_beta = 0.9  # EMA decay rate for control variate\n        self.clip_norm = max_grad_norm  # Current adaptive clip threshold\n        self.target_quantile = 0.5  # Target: median of residual norms\n        self.clip_lr = 0.2  # Geometric adaptation rate\n        self.clip_min = 0.001  # Floor for clip threshold\n        self.clip_max = max_grad_norm  # Ceiling caps noise at standard level\n\n    def clip_and_noise(self, per_sample_grads, step, epoch):\n        batch_size = per_sample_grads[0].shape[0]\n        if self.control_variate is None:\n            self.control_variate = [\n                torch.zeros_like(g[0]) for g in per_sample_grads]\n        # Residuals: r_i = g_i - mu (center with control variate)\n        residuals = [g - cv.unsqueeze(0)\n                     for g, cv in zip(per_sample_grads, self.control_variate)]\n        # Per-sample residual norms (global across all layers)\n        flat = torch.cat([r.reshape(batch_size, -1) for r in residuals], dim=1)\n        norms = flat.norm(2, dim=1)\n        # Adapt clip threshold to target quantile of residual norms\n        frac_above = (norms > self.clip_norm).float().mean().item()\n        self.clip_norm *= math.exp(\n            self.clip_lr * (frac_above - self.target_quantile))\n        self.clip_norm = max(self.clip_min, min(self.clip_max, self.clip_norm))\n        # Clip residuals to adaptive threshold\n        clip_factor = (self.clip_norm / norms.clamp(min=1e-8)).clamp(max=1.0)\n        noised_grads = []\n        new_cv = []\n        for r, cv in zip(residuals, self.control_variate):\n            shape = [batch_size] + [1] * (r.dim() - 1)\n            clipped = r * clip_factor.reshape(shape)\n            avg = clipped.mean(dim=0)\n            # Noise proportional to adaptive (smaller) clip threshold\n            noise = torch.randn_like(avg) * (\n                self.noise_multiplier * self.clip_norm / batch_size)\n            # Reconstruct: noised residual + control variate\n            output = avg + noise + cv\n            noised_grads.append(output)\n            # Update control variate via EMA (post-processing of noised output)\n            new_cv.append(\n                self.cv_beta * cv + (1 - self.cv_beta) * output.detach())\n        self.control_variate = new_cv\n        return noised_grads\n\n    def get_effective_sigma(self, step, epoch):\n        \"\"\"Sigma is constant, independent of adaptive clip norm.\"\"\"\n        return self.noise_multiplier\n\n# =====================================================================\n# EDITABLE SECTION END\n# ====================================================================="
7}

1{
2  "op": "replace",
3  "filename": "opacus/custom_dpsgd.py",
4  "start_line": 152,
5  "end_line": 233,
6  "content": "# =====================================================================\n# EDITABLE SECTION START (lines 152-233)\n# =====================================================================\n# DPMechanism: Cosine-Scheduled AUTO-S with Post-Noise EMA Denoising.\n#\n# Algorithmic innovations over standard DP-SGD:\n# 1. AUTO-S normalization: per-sample g/(||g||+gamma), sensitivity=1\n# 2. Cosine noise schedule: more privacy budget for later, useful steps\n# 3. EMA post-processing: denoises gradients (privacy-free post-processing)\n\nclass DPMechanism:\n    \"\"\"Cosine-Scheduled AUTO-S with EMA Denoising (CS-AUTO-EMA).\"\"\"\n\n    def __init__(self, max_grad_norm, noise_multiplier, n_params,\n                 dataset_size, batch_size, epochs, target_epsilon, target_delta):\n        self.max_grad_norm = max_grad_norm\n        self.noise_multiplier = noise_multiplier\n        self.n_params = n_params\n        self.dataset_size = dataset_size\n        self.batch_size = batch_size\n        self.epochs = epochs\n        self.target_epsilon = target_epsilon\n        self.target_delta = target_delta\n        self.gamma = 0.01  # AUTO-S regularization\n        self.steps_per_epoch = dataset_size // batch_size\n        self.total_steps = self.steps_per_epoch * epochs\n        self.noise_end_ratio = 0.65  # Cosine decays noise to 65% of initial\n        # Calibrate sigma_0 so total RDP matches uniform-sigma budget\n        inv_sq_sum = 0.0\n        for t in range(self.total_steps):\n            f = self._sched(t)\n            inv_sq_sum += 1.0 / (f * f)\n        self.sigma_0 = noise_multiplier * math.sqrt(inv_sq_sum / self.total_steps)\n        self.ema_beta = 0.5  # EMA smoothing factor\n        self.ema_grads = None\n\n    def _sched(self, step):\n        \"\"\"Cosine noise schedule factor: 1.0 -> noise_end_ratio.\"\"\"\n        p = min(step / max(self.total_steps - 1, 1), 1.0)\n        r = self.noise_end_ratio\n        return r + (1.0 - r) * 0.5 * (1.0 + math.cos(math.pi * p))\n\n    def clip_and_noise(self, per_sample_grads, step, epoch):\n        batch_size = per_sample_grads[0].shape[0]\n        flat = torch.cat([g.reshape(batch_size, -1) for g in per_sample_grads], dim=1)\n        norms = flat.norm(2, dim=1)  # [B]\n        # AUTO-S normalization: g / (||g|| + gamma), sensitivity bounded by 1\n        scale = 1.0 / (norms + self.gamma)  # [B]\n        # Cosine-scheduled noise level\n        f = self._sched(step)\n        noise_std = self.sigma_0 * f / batch_size\n        noised = []\n        for g in per_sample_grads:\n            shape = [batch_size] + [1] * (g.dim() - 1)\n            avg = (g * scale.reshape(shape)).mean(dim=0)\n            noised.append(avg + torch.randn_like(avg) * noise_std)\n        # EMA denoising (post-processing, privacy-free)\n        if self.ema_grads is None:\n            self.ema_grads = [g.clone() for g in noised]\n        else:\n            b = self.ema_beta\n            for i in range(len(noised)):\n                self.ema_grads[i].mul_(b).add_(noised[i], alpha=1.0 - b)\n        return [g.clone() for g in self.ema_grads]\n\n    def get_effective_sigma(self, step, epoch):\n        \"\"\"Return calibrated sigma for external privacy accounting.\"\"\"\n        return self.noise_multiplier\n\n# =====================================================================\n# EDITABLE SECTION END\n# =====================================================================\n"
7}

1{
2  "op": "replace",
3  "filename": "opacus/custom_dpsgd.py",
4  "start_line": 152,
5  "end_line": 223,
6  "content": "# =====================================================================\n# EDITABLE SECTION START (lines 152-233)\n# =====================================================================\n# DPMechanism: AUTO-S with LR-Adaptive Noise Budget Allocation.\n#\n# Key insight: For cosine-annealed LR, the optimal noise allocation that\n# minimizes sum_t (lr_t * sigma_t)^2 subject to fixed total RDP is\n# sigma_t proportional to 1/sqrt(lr_t). This gives LESS noise when LR\n# is high (maximizing early learning), MORE when LR is low (minimal\n# impact since updates are dampened). Combined with AUTO-S normalization.\n\nclass DPMechanism:\n    \"\"\"AUTO-S with LR-Adaptive Noise Schedule (LR-ANS).\n\n    Optimal noise allocation: sigma_t ∝ 1/sqrt(lr_t) minimizes total\n    effective noise in parameter updates under fixed privacy budget.\n    \"\"\"\n\n    def __init__(self, max_grad_norm, noise_multiplier, n_params,\n                 dataset_size, batch_size, epochs, target_epsilon, target_delta):\n        self.max_grad_norm = max_grad_norm\n        self.noise_multiplier = noise_multiplier\n        self.n_params = n_params\n        self.dataset_size = dataset_size\n        self.batch_size = batch_size\n        self.epochs = epochs\n        self.target_epsilon = target_epsilon\n        self.target_delta = target_delta\n        self.gamma = 0.01  # AUTO-S regularization\n        self.steps_per_epoch = dataset_size // batch_size\n        self.total_steps = self.steps_per_epoch * epochs\n        self.max_noise_factor = 4.0  # Cap noise increase at end\n        self.lr_floor = 0.05  # Floor for LR ratio to avoid inf\n        # Calibrate sigma_0 so total RDP matches uniform-sigma budget\n        inv_sq_sum = 0.0\n        for t in range(self.total_steps):\n            f = self._noise_factor(t)\n            inv_sq_sum += 1.0 / (f * f)\n        self.sigma_0 = noise_multiplier * math.sqrt(inv_sq_sum / self.total_steps)\n\n    def _noise_factor(self, step):\n        \"\"\"Noise factor ∝ 1/sqrt(lr_t) for cosine LR schedule.\"\"\"\n        progress = min(step / max(self.total_steps - 1, 1), 1.0)\n        # Approximate cosine LR: lr_ratio = (1 + cos(pi * progress)) / 2\n        lr_ratio = 0.5 * (1.0 + math.cos(math.pi * progress))\n        lr_ratio = max(lr_ratio, self.lr_floor)\n        return min(1.0 / math.sqrt(lr_ratio), self.max_noise_factor)\n\n    def clip_and_noise(self, per_sample_grads, step, epoch):\n        batch_size = per_sample_grads[0].shape[0]\n        flat = torch.cat([g.reshape(batch_size, -1) for g in per_sample_grads], dim=1)\n        norms = flat.norm(2, dim=1)  # [B]\n        # AUTO-S normalization: g / (||g|| + gamma), sensitivity <= 1\n        scale = 1.0 / (norms + self.gamma)  # [B]\n        # LR-adaptive noise level\n        f = self._noise_factor(step)\n        noise_std = self.sigma_0 * f / batch_size\n        noised_grads = []\n        for g in per_sample_grads:\n            shape = [batch_size] + [1] * (g.dim() - 1)\n            avg = (g * scale.reshape(shape)).mean(dim=0)\n            noised_grads.append(avg + torch.randn_like(avg) * noise_std)\n        return noised_grads\n\n    def get_effective_sigma(self, step, epoch):\n        \"\"\"Return calibrated sigma for external privacy accounting.\"\"\"\n        return self.noise_multiplier\n\n# =====================================================================\n# EDITABLE SECTION END\n# =====================================================================\n"
7}

1{
2  "n": 1
3}