1{
2  "start_line": 172,
3  "content": "class DPMechanism:\n    \"\"\"Centered Adaptive Clipping.\n    \n    Maintains an exponential moving average (m) of the noised gradients to center the \n    per-sample gradients before clipping. Because the variance (norm) of g_i - m is smaller\n    than g_i, we can use a smaller clipping threshold and thus add less noise.\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\n        # Adaptive clipping parameters\n        self.clip_norm = max_grad_norm\n        self.target_quantile = 0.5\n        self.clip_lr = 0.2\n        self.clip_min = 1e-4\n        self.clip_max = 100.0\n\n        # EMA for gradient centering\n        self.m = None\n        self.beta = 0.9  # momentum for the center\n        \n    def clip_and_noise(self, per_sample_grads, step, epoch):\n        batch_size = per_sample_grads[0].shape[0]\n\n        if self.m is None:\n            self.m = [torch.zeros_like(g[0]) for g in per_sample_grads]\n\n        # Compute deltas: g_i - m\n        deltas = []\n        for g, m_p in zip(per_sample_grads, self.m):\n            deltas.append(g - m_p.unsqueeze(0))\n\n        # Compute norms of deltas\n        flat = torch.cat([d.reshape(batch_size, -1) for d in deltas], dim=1)\n        norms = flat.norm(2, dim=1)  # [B]\n\n        # Adaptive clipping\n        frac_above = (norms > self.clip_norm).float().mean().item()\n        self.clip_norm = self.clip_norm * math.exp(\n            self.clip_lr * (frac_above - self.target_quantile)\n        )\n        self.clip_norm = max(self.clip_min, min(self.clip_max, self.clip_norm))\n\n        clip_factor = (self.clip_norm / norms.clamp(min=1e-8)).clamp(max=1.0)\n\n        noised_grads = []\n        for d, m_p in zip(deltas, self.m):\n            shape = [batch_size] + [1] * (d.dim() - 1)\n            clipped_d = d * clip_factor.reshape(shape)\n\n            avg_d = clipped_d.mean(dim=0)\n            noise = torch.randn_like(avg_d) * (\n                self.noise_multiplier * self.clip_norm / batch_size\n            )\n            \n            noised_d = avg_d + noise\n            noised_grads.append(m_p + noised_d)\n\n        # Update EMA center using the noised gradients\n        with torch.no_grad():\n            for i in range(len(self.m)):\n                self.m[i] = self.beta * self.m[i] + (1 - self.beta) * noised_grads[i]\n\n        return noised_grads\n\n    def get_effective_sigma(self, step, epoch):\n        return self.noise_multiplier",
4  "filename": "opacus/custom_dpsgd.py",
5  "after_line": 151,
6  "end_line": 230,
7  "op": "replace"
8}

1{
2  "end_line": 230,
3  "start_line": 172,
4  "content": "class DPMechanism:\n    \"\"\"Adaptive Clipping with Momentum-based Variance Reduction.\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\n        # Adaptive clipping parameters\n        self.clip_norm = max_grad_norm\n        self.target_quantile = 0.5\n        self.clip_lr = 0.2\n        self.clip_min = 1e-4\n        self.clip_max = 100.0\n\n        # Variance reduction\n        self.m = None\n        self.beta = 0.9\n\n    def clip_and_noise(self, per_sample_grads, step, epoch):\n        batch_size = per_sample_grads[0].shape[0]\n\n        if self.m is None:\n            self.m = [torch.zeros_like(g[0]) for g in per_sample_grads]\n\n        # Compute per-sample gradient differences\n        centered_grads = []\n        for g, m in zip(per_sample_grads, self.m):\n            centered_grads.append(g - m.unsqueeze(0))\n\n        # Compute norms\n        flat = torch.cat([g.reshape(batch_size, -1) for g in centered_grads], dim=1)\n        norms = flat.norm(2, dim=1)  # [B]\n\n        # Adaptive clipping\n        frac_above = (norms > self.clip_norm).float().mean().item()\n        self.clip_norm = self.clip_norm * math.exp(\n            self.clip_lr * (frac_above - self.target_quantile)\n        )\n        self.clip_norm = max(self.clip_min, min(self.clip_max, self.clip_norm))\n\n        clip_factor = (self.clip_norm / norms.clamp(min=1e-8)).clamp(max=1.0)\n\n        noised_grads = []\n        new_m = []\n        for cg, m in zip(centered_grads, self.m):\n            shape = [batch_size] + [1] * (cg.dim() - 1)\n            clipped = cg * clip_factor.reshape(shape)\n\n            avg = clipped.mean(dim=0)\n            noise = torch.randn_like(avg) * (\n                self.noise_multiplier * self.clip_norm / batch_size\n            )\n            \n            eff_grad = m + avg + noise\n            noised_grads.append(eff_grad)\n            new_m.append(self.beta * m + (1 - self.beta) * eff_grad)\n\n        self.m = new_m\n        return noised_grads\n\n    def get_effective_sigma(self, step, epoch):\n        return self.noise_multiplier\n",
5  "after_line": 152,
6  "filename": "opacus/custom_dpsgd.py",
7  "op": "replace"
8}

1{
2  "after_line": 152,
3  "op": "replace",
4  "end_line": 230,
5  "content": "class DPMechanism:\n    \"\"\"Auto-S with Momentum-based Gradient Centering.\n    \n    Centers per-sample gradients using a running average of public noised gradients.\n    Then applies Auto-S normalization to the centered gradients (bounding sensitivity to 1).\n    Restores the mean and updates the running average.\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.noise_multiplier = noise_multiplier\n        self.gamma = 0.01\n        self.beta = 0.9  # Momentum parameter for the running average\n        self.mu = None\n\n    def clip_and_noise(self, per_sample_grads, step, epoch):\n        batch_size = per_sample_grads[0].shape[0]\n\n        if self.mu is None:\n            self.mu = [torch.zeros_like(g[0]) for g in per_sample_grads]\n\n        # Centering per-sample gradients\n        centered_grads = []\n        for g, m in zip(per_sample_grads, self.mu):\n            shape = [batch_size] + [1] * (g.dim() - 1)\n            centered_grads.append(g - m.unsqueeze(0))\n\n        # Compute per-sample gradient norms of the centered gradients\n        flat = torch.cat([c.reshape(batch_size, -1) for c in centered_grads], dim=1)\n        norms = flat.norm(2, dim=1)  # [B]\n\n        # AUTO-S normalization on centered gradients\n        scale = 1.0 / (norms + self.gamma)  # [B]\n\n        noised_grads = []\n        updated_mu = []\n        for c, m in zip(centered_grads, self.mu):\n            shape = [batch_size] + [1] * (c.dim() - 1)\n            normalized = c * scale.reshape(shape)\n\n            # Average over batch\n            avg = normalized.mean(dim=0)\n\n            # Add calibrated noise (sensitivity is bounded by 1)\n            noise = torch.randn_like(avg) * (self.noise_multiplier * 1.0 / batch_size)\n            noised_delta = avg + noise\n\n            # Restore the mean\n            final_grad = noised_delta + m\n            noised_grads.append(final_grad)\n\n            # Update running average (public state)\n            updated_mu.append(m * self.beta + final_grad * (1.0 - self.beta))\n            \n        self.mu = updated_mu\n\n        return noised_grads\n\n    def get_effective_sigma(self, step, epoch):\n        return self.noise_multiplier\n",
6  "start_line": 172,
7  "filename": "opacus/custom_dpsgd.py"
8}

1{
2  "filename": "opacus/custom_dpsgd.py",
3  "end_line": 232,
4  "after_line": 152,
5  "start_line": 172,
6  "op": "replace",
7  "content": "class DPMechanism:\n    \"\"\"Adaptive Clipping with Momentum-based Gradient Centering (DP-GC-Adapt).\n    \n    Uses an exponential moving average (EMA) of noised gradients to center the\n    per-sample gradients, reducing their variance. It then applies adaptive \n    quantile clipping to the centered gradients and adds calibrated noise,\n    finally restoring the mean. The reduced variance allows the adaptive clipping\n    threshold to drop significantly, meaning much less noise is added overall.\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.noise_multiplier = noise_multiplier\n        self.clip_norm = max_grad_norm\n        self.target_quantile = 0.5\n        self.clip_lr = 0.2\n        self.clip_min = 1e-4  # lowered to accommodate smaller centered gradients\n        self.clip_max = 100.0\n        \n        self.beta = 0.9\n        self.mu = None\n\n    def clip_and_noise(self, per_sample_grads, step, epoch):\n        batch_size = per_sample_grads[0].shape[0]\n\n        if self.mu is None:\n            self.mu = [torch.zeros_like(g[0]) for g in per_sample_grads]\n\n        # Bias correction for the EMA to track the gradient properly early on\n        bias_corr = 1.0 - self.beta ** (step + 1)\n\n        centered_grads = []\n        m_hats = []\n        for g, m in zip(per_sample_grads, self.mu):\n            m_hat = m / bias_corr\n            m_hats.append(m_hat)\n            shape = [batch_size] + [1] * (g.dim() - 1)\n            centered_grads.append(g - m_hat.unsqueeze(0))\n\n        # Compute per-sample gradient norms of centering difference\n        flat = torch.cat([c.reshape(batch_size, -1) for c in centered_grads], dim=1)\n        norms = flat.norm(2, dim=1)  # [B]\n\n        # Adaptive clipping update on the centered gradients\n        frac_above = (norms > self.clip_norm).float().mean().item()\n        self.clip_norm = self.clip_norm * math.exp(self.clip_lr * (frac_above - self.target_quantile))\n        self.clip_norm = max(self.clip_min, min(self.clip_max, self.clip_norm))\n\n        # Apply clipping\n        clip_factor = (self.clip_norm / norms.clamp(min=1e-8)).clamp(max=1.0)\n\n        noised_grads = []\n        updated_mu = []\n        for c, m, m_hat in zip(centered_grads, self.mu, m_hats):\n            shape = [batch_size] + [1] * (c.dim() - 1)\n            clipped = c * clip_factor.reshape(shape)\n\n            # Average over batch\n            avg = clipped.mean(dim=0)\n\n            # Add calibrated noise proportional to adaptive clip_norm\n            noise = torch.randn_like(avg) * (self.noise_multiplier * self.clip_norm / batch_size)\n\n            noised_delta = avg + noise\n            final_grad = noised_delta + m_hat  # Restore the mean\n\n            # Update public state (EMA)\n            updated_mu.append(m * self.beta + final_grad * (1.0 - self.beta))\n            noised_grads.append(final_grad)\n\n        self.mu = updated_mu\n        return noised_grads\n\n    def get_effective_sigma(self, step, epoch):\n        return self.noise_multiplier"
8}

1{
2  "content": "class DPMechanism:\n    \"\"\"Auto-S with Noise Decay (AUTO-S-ND).\n    \n    Combines the automatic clipping (per-sample gradient normalization) to bound\n    sensitivity to 1, with a step-decay schedule for the noise multiplier.\n    This gives more privacy budget to later stages of training where gradients\n    are smaller and SNR is lower.\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.noise_multiplier = noise_multiplier\n        self.gamma = 0.01\n\n        self.decay_interval = max(1, epochs // 4)\n        self.noise_decay_factor = 0.8\n        self.steps_per_epoch = dataset_size // batch_size\n\n        total_steps = self.steps_per_epoch * epochs\n        inv_sq_sum = 0.0\n        for e in range(1, epochs + 1):\n            stage = (e - 1) // self.decay_interval\n            factor = self.noise_decay_factor ** stage\n            inv_sq_sum += self.steps_per_epoch / (factor * factor)\n            \n        self.sigma_0 = noise_multiplier * (inv_sq_sum / total_steps) ** 0.5\n        self._current_sigma = self.sigma_0\n\n    def clip_and_noise(self, per_sample_grads, step, epoch):\n        batch_size = per_sample_grads[0].shape[0]\n\n        stage = (epoch - 1) // self.decay_interval\n        self._current_sigma = self.sigma_0 * (self.noise_decay_factor ** stage)\n\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\n        scale = 1.0 / (norms + self.gamma)  # [B]\n\n        noised_grads = []\n        for g in per_sample_grads:\n            shape = [batch_size] + [1] * (g.dim() - 1)\n            normalized = g * scale.reshape(shape)\n\n            avg = normalized.mean(dim=0)\n\n            noise = torch.randn_like(avg) * (self._current_sigma * 1.0 / batch_size)\n            noised_grads.append(avg + noise)\n\n        return noised_grads\n\n    def get_effective_sigma(self, step, epoch):\n        return self.noise_multiplier",
3  "op": "replace",
4  "filename": "opacus/custom_dpsgd.py",
5  "after_line": 152,
6  "start_line": 172,
7  "end_line": 247
8}