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: Controls how per-sample gradients are clipped and noised.\n#\n# Interface contract:\n#   __init__(self, max_grad_norm, noise_multiplier, n_params, dataset_size,\n#            batch_size, epochs, target_epsilon, target_delta)\n#   clip_and_noise(self, per_sample_grads, step, epoch) -> noised_gradient\n#   get_effective_sigma(self, step, epoch) -> float\n#\n# The mechanism receives per-sample gradients (list of tensors, each [B, *param_shape])\n# and must return aggregated + noised gradients (list of tensors, each [*param_shape]).\n#\n# IMPORTANT:\n# - The total privacy budget (target_epsilon, target_delta) is FIXED.\n# - Your mechanism must not exceed it. The accounting is checked externally.\n# - You may adapt clipping thresholds, noise schedules, or gradient processing\n#   as long as privacy guarantees hold.\n\nclass DPMechanism:\n    \"\"\"ALDS: Adaptive Layer-wise Differentially Private SGD.\n\n    Features:\n    1. Per-layer adaptive clipping based on moving quantiles of gradient norms.\n    2. Noise scheduling: decay noise multiplier over training epochs.\n    3. Momentum-based variance reduction via gradient averaging.\n    \n    Privacy accounting: computes equivalent uniform sigma that matches total RDP.\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        # Per-layer adaptive clipping parameters\n        self.target_quantile = 0.5  # target median norm per layer\n        self.clip_lr = 0.1  # learning rate for clipping threshold adaptation\n        self.clip_min = 0.01\n        self.clip_max = 10.0\n        \n        # Layer clipping thresholds (will be initialized on first batch)\n        self.layer_clips = None\n        self.layer_momentum = None  # for tracking moving quantiles\n        self.momentum_beta = 0.9\n        \n        # Noise scheduling\n        self.noise_decay_factor = 0.85  # decay per epoch\n        self.clip_decay_factor = 0.9  # decay clipping thresholds over time\n        \n        # Compute equivalent uniform sigma for privacy accounting\n        # We decay sigma each epoch: sigma_t = sigma_0 * (decay_factor)^(t-1)\n        # Need to find sigma_0 such that harmonic mean equals calibrated noise_multiplier\n        total_steps = (dataset_size // batch_size) * epochs\n        inv_sq_sum = 0.0\n        for e in range(1, epochs + 1):\n            decay = self.noise_decay_factor ** (e - 1)\n            steps_per_epoch = dataset_size // batch_size\n            inv_sq_sum += steps_per_epoch / (decay * decay)\n        # Solve: noise_multiplier = sigma_0 * sqrt(total_steps / inv_sq_sum)\n        self.sigma_0 = noise_multiplier * (inv_sq_sum / total_steps) ** 0.5\n        \n        # Current values\n        self._current_sigma = self.sigma_0\n        self._current_global_clip = max_grad_norm  # global sensitivity bound\n        \n        # For privacy accounting: effective sigma that yields same total RDP\n        self._effective_sigma = noise_multiplier  # initially\n        \n        # Track steps for scheduling\n        self.total_steps = total_steps\n        self.steps_per_epoch = dataset_size // batch_size\n\n    def _init_layer_clips(self, n_layers):\n        \"\"\"Initialize per-layer clipping thresholds.\"\"\"\n        # Start with equal allocation: total sensitivity sqrt(sum C_l^2) = max_grad_norm\n        # If we set each C_l = max_grad_norm / sqrt(n_layers), then sqrt(sum C_l^2) = max_grad_norm\n        # But we want to adapt based on actual gradient scales\n        initial = self.max_grad_norm / math.sqrt(n_layers) if n_layers > 0 else self.max_grad_norm\n        self.layer_clips = [initial] * n_layers\n        self.layer_momentum = [initial] * n_layers  # for moving quantile estimate\n\n    def clip_and_noise(self, per_sample_grads, step, epoch):\n        \"\"\"Clip per-sample gradients (per-layer) and add noise.\"\"\"\n        batch_size = per_sample_grads[0].shape[0]\n        n_layers = len(per_sample_grads)\n        \n        # Initialize layer clips if needed\n        if self.layer_clips is None:\n            self._init_layer_clips(n_layers)\n        \n        # Update noise and global clip schedule\n        self._current_sigma = self.sigma_0 * (self.noise_decay_factor ** (epoch - 1))\n        self._current_global_clip = self.max_grad_norm * (self.clip_decay_factor ** (epoch - 1))\n        \n        # Compute per-sample per-layer gradient norms\n        layer_norms = []\n        for l, g in enumerate(per_sample_grads):\n            # Reshape to [B, -1] to compute norm per sample\n            flat = g.reshape(batch_size, -1)\n            norms = flat.norm(2, dim=1)  # [B]\n            layer_norms.append(norms)\n        \n        # Update per-layer clipping thresholds based on quantile of norms\n        for l in range(n_layers):\n            norms = layer_norms[l]\n            # Compute fraction exceeding current clip\n            frac_above = (norms > self.layer_clips[l]).float().mean().item()\n            # Geometric update toward target quantile\n            update = math.exp(self.clip_lr * (frac_above - self.target_quantile))\n            new_clip = self.layer_clips[l] * update\n            # Apply momentum smoothing\n            self.layer_momentum[l] = (self.momentum_beta * self.layer_momentum[l] + \n                                     (1 - self.momentum_beta) * new_clip)\n            # Bound and scale to maintain global sensitivity bound\n            # We want sqrt(sum layer_clips^2) <= current_global_clip\n            # We'll normalize after updating all layers\n            self.layer_clips[l] = max(self.clip_min, min(self.clip_max, new_clip))\n        \n        # Normalize layer clips to satisfy global sensitivity bound\n        # Compute current total sensitivity: sqrt(sum C_l^2)\n        total_sensitivity = math.sqrt(sum(c * c for c in self.layer_clips))\n        if total_sensitivity > self._current_global_clip:\n            # Scale down proportionally\n            scale = self._current_global_clip / total_sensitivity\n            self.layer_clips = [c * scale for c in self.layer_clips]\n            self.layer_momentum = [c * scale for c in self.layer_momentum]\n        \n        # Clip per-sample gradients using per-layer thresholds\n        noised_grads = []\n        for l, g in enumerate(per_sample_grads):\n            norms = layer_norms[l]\n            clip_factor = (self.layer_clips[l] / norms.clamp(min=1e-8)).clamp(max=1.0)\n            \n            # Apply clipping\n            shape = [batch_size] + [1] * (g.dim() - 1)\n            clipped = g * clip_factor.reshape(shape)\n            \n            # Average over batch\n            avg = clipped.mean(dim=0)\n            \n            # Add noise calibrated to this layer's contribution to sensitivity\n            # The sensitivity contribution of layer l is self.layer_clips[l]\n            # But total sensitivity is sqrt(sum C_l^2), noise should be scaled by total sensitivity\n            # Actually, standard DP-SGD: noise ~ sigma * C / B where C is total L2 sensitivity bound\n            # Since we're clipping each layer independently with C_l, the total sensitivity is\n            # sqrt(sum C_l^2). So noise should be sigma * sqrt(sum C_l^2) / B.\n            # However, we add noise independently per parameter. The noise variance per parameter\n            # should be (sigma * C_total / B)^2. So we can add noise with stddev = sigma * C_total / B\n            # to each parameter (including all layers).\n            # This is equivalent to standard DP-SGD with C = C_total.\n            total_C = math.sqrt(sum(c * c for c in self.layer_clips))\n            noise_std = self._current_sigma * total_C / batch_size\n            noise = torch.randn_like(avg) * noise_std\n            noised_grads.append(avg + noise)\n        \n        # Update effective sigma for privacy accounting\n        # The effective sigma that yields same total RDP as our schedule\n        # can be computed as harmonic mean: sigma_eff = sqrt(total_steps / sum_t 1/sigma_t^2)\n        # We compute this incrementally\n        current_step = (epoch - 1) * self.steps_per_epoch + (step % self.steps_per_epoch)\n        if current_step > 0:\n            # Contribution of this step: 1/sigma_t^2\n            inv_sq = 1.0 / (self._current_sigma * self._current_sigma)\n            # Approximate effective sigma (not used for privacy accounting, just for reporting)\n            # Actual accounting is done via get_effective_sigma\n            pass\n        \n        return noised_grads\n\n    def get_effective_sigma(self, step, epoch):\n        \"\"\"Return effective noise multiplier for privacy accounting.\n        \n        We need to return a sigma such that compute_epsilon(steps, sigma, q, delta)\n        gives the correct total privacy spend for our non-uniform schedule.\n        The external accounting uses uniform sigma assumption.\n        We pre-computed sigma_0 such that harmonic mean equals calibrated noise_multiplier.\n        But actually, we need to compute the RDP of the entire composition.\n        Since we already calibrated sigma_0 to make the harmonic mean equal to\n        the calibrated noise_multiplier, we can return noise_multiplier.\n        \n        However, this is approximate because RDP composition for varying sigma\n        is not exactly equivalent to harmonic mean. But it's a common approximation.\n        \"\"\"\n        # Return the calibrated uniform sigma that was provided\n        return self.noise_multiplier\n\n# =====================================================================\n# EDITA\n... (truncated)"
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: Controls how per-sample gradients are clipped and noised.\n#\n# Interface contract:\n#   __init__(self, max_grad_norm, noise_multiplier, n_params, dataset_size,\n#            batch_size, epochs, target_epsilon, target_delta)\n#   clip_and_noise(self, per_sample_grads, step, epoch) -> noised_gradient\n#   get_effective_sigma(self, step, epoch) -> float\n#\n# The mechanism receives per-sample gradients (list of tensors, each [B, *param_shape])\n# and must return aggregated + noised gradients (list of tensors, each [*param_shape]).\n#\n# IMPORTANT:\n# - The total privacy budget (target_epsilon, target_delta) is FIXED.\n# - Your mechanism must not exceed it. The accounting is checked externally.\n# - You may adapt clipping thresholds, noise schedules, or gradient processing\n#   as long as privacy guarantees hold.\n\nclass DPMechanism:\n    \"\"\"Adaptive Momentum-Aware Clipping with Noise Scheduling (AMC-NS).\n\n    Combines:\n    1. Per-layer adaptive clipping thresholds based on EMA of gradient norms.\n    2. Momentum-based gradient aggregation to reduce noise variance.\n    3. Noise decay schedule allocating more privacy budget to later steps.\n    4. Privacy-preserving momentum update.\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_ema_decay = 0.9  # EMA decay for gradient norm statistics\n        self.clip_init = max_grad_norm\n        self.clip_min = 0.01\n        self.clip_max = 10.0\n        # Per-layer clipping thresholds (will be initialized when first seeing gradients)\n        self.layer_clips = None\n        self.layer_norm_ema = None  # EMA of per-sample gradient norms per layer\n\n        # Momentum parameters\n        self.momentum = 0.9\n        self.momentum_buffer = None  # will be list of tensors\n        self.velocity = None  # momentum velocity for variance reduction\n\n        # Noise schedule: allocate more noise early, less later\n        self.noise_decay_start = noise_multiplier\n        self.noise_decay_end = noise_multiplier * 0.3  # 70% reduction by end\n        self.total_steps = (dataset_size // batch_size) * epochs\n        # Effective sigma for accounting: compute harmonic mean of per-step sigma\n        # We'll track cumulative RDP, but for simplicity we compute equivalent uniform sigma\n        # that yields same total privacy cost given our schedule.\n        self._compute_noise_schedule()\n\n        # Privacy accounting: we need to return effective sigma that matches our schedule\n        # We'll compute the uniform sigma that gives same epsilon as our varying sigma schedule\n        self.uniform_sigma = self._compute_uniform_sigma()\n\n    def _compute_noise_schedule(self):\n        \"\"\"Pre-compute per-step noise multipliers.\"\"\"\n        # Linear decay from start to end over total steps\n        self.sigma_schedule = []\n        for step in range(self.total_steps):\n            progress = step / max(self.total_steps - 1, 1)\n            sigma = self.noise_decay_start * (1 - progress) + self.noise_decay_end * progress\n            self.sigma_schedule.append(sigma)\n\n    def _compute_uniform_sigma(self):\n        \"\"\"Compute uniform sigma that yields same RDP as our varying schedule.\n        \n        Using the fact that RDP of Gaussian mechanism scales as ~1/sigma^2.\n        For composition, the effective uniform sigma is the harmonic mean.\n        \"\"\"\n        inv_sq_sum = 0.0\n        for sigma in self.sigma_schedule:\n            inv_sq_sum += 1.0 / (sigma * sigma)\n        harmonic_mean_sq = len(self.sigma_schedule) / inv_sq_sum\n        return math.sqrt(harmonic_mean_sq)\n\n    def clip_and_noise(self, per_sample_grads, step, epoch):\n        \"\"\"Clip per-sample gradients and add noise with momentum.\"\"\"\n        batch_size = per_sample_grads[0].shape[0]\n\n        # Initialize per-layer clipping thresholds if first call\n        if self.layer_clips is None:\n            self.layer_clips = [self.clip_init for _ in per_sample_grads]\n            self.layer_norm_ema = [self.clip_init for _ in per_sample_grads]\n            self.momentum_buffer = [torch.zeros_like(g[0]) for g in per_sample_grads]\n            self.velocity = [torch.zeros_like(g[0]) for g in per_sample_grads]\n\n        # Compute per-sample gradient norms per layer\n        layer_norms = []\n        for g in per_sample_grads:\n            # flatten per-sample dimensions\n            flat = g.reshape(batch_size, -1)\n            norms = flat.norm(2, dim=1)  # [B]\n            layer_norms.append(norms)\n\n        # Update per-layer EMA of gradient norms and adjust clipping thresholds\n        for i, norms in enumerate(layer_norms):\n            # Compute batch median norm as robust statistic\n            median_norm = torch.median(norms).item()\n            # Update EMA\n            self.layer_norm_ema[i] = (\n                self.clip_ema_decay * self.layer_norm_ema[i] +\n                (1 - self.clip_ema_decay) * median_norm\n            )\n            # Set clipping threshold to EMA value, bounded\n            self.layer_clips[i] = max(\n                self.clip_min,\n                min(self.clip_max, self.layer_norm_ema[i])\n            )\n\n        # Clip per-sample gradients using per-layer thresholds\n        clipped_grads = []\n        for i, g in enumerate(per_sample_grads):\n            norms = layer_norms[i]\n            clip_factor = (self.layer_clips[i] / norms.clamp(min=1e-8)).clamp(max=1.0)\n            shape = [batch_size] + [1] * (g.dim() - 1)\n            clipped = g * clip_factor.reshape(shape)\n            clipped_grads.append(clipped)\n\n        # Average clipped gradients over batch\n        avg_grads = [clipped.mean(dim=0) for clipped in clipped_grads]\n\n        # Get current noise multiplier from schedule\n        current_step = step % len(self.sigma_schedule) if self.sigma_schedule else 0\n        current_sigma = self.sigma_schedule[current_step] if current_step < len(self.sigma_schedule) else self.noise_multiplier\n\n        # Add calibrated Gaussian noise\n        noised_grads = []\n        for i, avg in enumerate(avg_grads):\n            # Noise scaled by per-layer clipping threshold\n            noise = torch.randn_like(avg) * (\n                current_sigma * self.layer_clips[i] / batch_size\n            )\n            noised = avg + noise\n            noised_grads.append(noised)\n\n        # Momentum update: combine current noised gradient with momentum buffer\n        # This reduces variance of the gradient estimate\n        if self.momentum > 0:\n            for i in range(len(noised_grads)):\n                # Update velocity: v = momentum * v + (1 - momentum) * g\n                self.velocity[i] = (\n                    self.momentum * self.velocity[i] +\n                    (1 - self.momentum) * noised_grads[i]\n                )\n                # Bias correction for momentum (like Adam)\n                corrected_velocity = self.velocity[i] / (1 - self.momentum ** (step + 1))\n                noised_grads[i] = corrected_velocity\n\n        return noised_grads\n\n    def get_effective_sigma(self, step, epoch):\n        \"\"\"Return effective uniform sigma for privacy accounting.\"\"\"\n        return self.uniform_sigma\n\n# =====================================================================\n# EDITABLE SECTION END\n# ====================================================================="
7}

1{
2  "op": "replace",
3  "filename": "opacus/custom_dpsgd.py",
4  "start_line": 172,
5  "end_line": 230,
6  "content": "class DPMechanism:\n    \"\"\"Adaptive Clipping with Momentum Variance Reduction (ALC-MVR).\n\n    Features:\n    1. Adaptive clipping threshold per layer, targeting median of gradient norms.\n    2. Momentum buffer to reduce noise variance across steps.\n    3. Noise scaling proportional to clipping threshold (effective sigma adaptation).\n    4. Exponential decay of clipping threshold over epochs to focus on fine-tuning.\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.target_quantile = 0.5  # target median\n        self.clip_lr = 0.2  # geometric update learning rate\n        self.clip_min = 0.01\n        self.clip_max = 10.0\n        # Initialize clipping threshold per layer? We'll start with global.\n        self.clip_norm = max_grad_norm\n        \n        # Momentum parameters\n        self.momentum_beta = 0.9\n        self.momentum_buffers = None  # initialized after first gradient shape\n        \n        # Exponential decay of clipping threshold over epochs\n        self.clip_decay_factor = 0.95  # per epoch decay\n        self.base_clip_norm = max_grad_norm\n\n    def clip_and_noise(self, per_sample_grads, step, epoch):\n        batch_size = per_sample_grads[0].shape[0]\n\n        # Compute per-sample gradient norms (global across all parameters)\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        # Update clipping threshold using geometric update\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        # Apply exponential decay over epochs (additional)\n        epoch_decay = self.clip_decay_factor ** (epoch - 1)\n        effective_clip = self.clip_norm * epoch_decay\n        effective_clip = max(self.clip_min, min(self.clip_max, effective_clip))\n\n        # Clip per-sample gradients\n        clip_factor = (effective_clip / norms.clamp(min=1e-8)).clamp(max=1.0)\n\n        noised_grads = []\n        for i, g in enumerate(per_sample_grads):\n            shape = [batch_size] + [1] * (g.dim() - 1)\n            clipped = g * clip_factor.reshape(shape)\n\n            # Average over batch\n            avg = clipped.mean(dim=0)\n\n            # Add calibrated Gaussian noise\n            noise = torch.randn_like(avg) * (\n                self.noise_multiplier * effective_clip / batch_size\n            )\n            noised = avg + noise\n            \n            # Apply momentum if buffers initialized\n            if self.momentum_buffers is not None:\n                self.momentum_buffers[i] = self.momentum_beta * self.momentum_buffers[i] + (1 - self.momentum_beta) * noised\n                noised = self.momentum_buffers[i]\n            else:\n                # first call: initialize momentum buffers\n                self.momentum_buffers = [torch.zeros_like(noised) for noised in noised_grads + [noised]]\n                # need to redo? We'll just store for next iteration\n                # For now, skip momentum for first batch\n                pass\n            \n            noised_grads.append(noised)\n\n        # If first batch, initialize momentum buffers with current noised grads\n        if self.momentum_buffers is None:\n            self.momentum_buffers = [ng.clone() for ng in noised_grads]\n\n        return noised_grads\n\n    def get_effective_sigma(self, step, epoch):\n        \"\"\"Return effective noise multiplier for privacy accounting.\n        \n        The actual noise added is noise_multiplier * effective_clip / batch_size.\n        The baseline expects noise_multiplier * max_grad_norm / batch_size.\n        So effective sigma = noise_multiplier * (effective_clip / max_grad_norm).\n        \"\"\"\n        epoch_decay = self.clip_decay_factor ** (epoch - 1)\n        effective_clip = self.clip_norm * epoch_decay\n        effective_clip = max(self.clip_min, min(self.clip_max, effective_clip))\n        return self.noise_multiplier * effective_clip / self.max_grad_norm"
7}

1{
2  "op": "replace",
3  "filename": "opacus/custom_dpsgd.py",
4  "start_line": 172,
5  "end_line": 274,
6  "content": "class DPMechanism:\n    \"\"\"Adaptive Clipping with Momentum Variance Reduction (ALC-MVR).\n\n    Features:\n    1. Adaptive clipping threshold targeting median of gradient norms.\n    2. Momentum buffer to reduce noise variance across steps.\n    3. Noise scaling proportional to clipping threshold (effective sigma adaptation).\n    4. Exponential decay of clipping threshold over epochs to focus on fine-tuning.\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.target_quantile = 0.5  # target median\n        self.clip_lr = 0.2  # geometric update learning rate\n        self.clip_min = 0.01\n        self.clip_max = 10.0\n        self.clip_norm = max_grad_norm\n        \n        # Momentum parameters\n        self.momentum_beta = 0.9\n        self.momentum_buffers = None  # will be list of tensors\n        \n        # Exponential decay of clipping threshold over epochs\n        self.clip_decay_factor = 0.95  # per epoch decay\n\n    def clip_and_noise(self, per_sample_grads, step, epoch):\n        batch_size = per_sample_grads[0].shape[0]\n\n        # Compute per-sample gradient norms (global across all parameters)\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        # Update clipping threshold using geometric update\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        # Apply exponential decay over epochs (additional)\n        epoch_decay = self.clip_decay_factor ** (epoch - 1)\n        effective_clip = self.clip_norm * epoch_decay\n        effective_clip = max(self.clip_min, min(self.clip_max, effective_clip))\n\n        # Clip per-sample gradients\n        clip_factor = (effective_clip / norms.clamp(min=1e-8)).clamp(max=1.0)\n\n        # Initialize momentum buffers if first call\n        if self.momentum_buffers is None:\n            self.momentum_buffers = [torch.zeros_like(g[0]) for g in per_sample_grads]\n\n        noised_grads = []\n        for i, g in enumerate(per_sample_grads):\n            shape = [batch_size] + [1] * (g.dim() - 1)\n            clipped = g * clip_factor.reshape(shape)\n\n            # Average over batch\n            avg = clipped.mean(dim=0)\n\n            # Add calibrated Gaussian noise\n            noise = torch.randn_like(avg) * (\n                self.noise_multiplier * effective_clip / batch_size\n            )\n            noised = avg + noise\n            \n            # Apply momentum\n            self.momentum_buffers[i] = self.momentum_beta * self.momentum_buffers[i] + (1 - self.momentum_beta) * noised\n            noised_grads.append(self.momentum_buffers[i])\n\n        return noised_grads\n\n    def get_effective_sigma(self, step, epoch):\n        \"\"\"Return effective noise multiplier for privacy accounting.\n        \n        The actual noise added is noise_multiplier * effective_clip / batch_size.\n        The baseline expects noise_multiplier * max_grad_norm / batch_size.\n        So effective sigma = noise_multiplier * (effective_clip / max_grad_norm).\n        \"\"\"\n        epoch_decay = self.clip_decay_factor ** (epoch - 1)\n        effective_clip = self.clip_norm * epoch_decay\n        effective_clip = max(self.clip_min, min(self.clip_max, effective_clip))\n        return self.noise_multiplier * effective_clip / self.max_grad_norm"
7}

1{
2  "op": "replace",
3  "filename": "opacus/custom_dpsgd.py",
4  "start_line": 172,
5  "end_line": 263,
6  "content": "class DPMechanism:\n    \"\"\"Adaptive Clipping with Variance-Aware Noise (AC-VAN).\n\n    Features:\n    1. Adaptive clipping threshold targeting median of gradient norms.\n    2. Per-layer gradient normalization for outliers.\n    3. Noise scaling inversely proportional to gradient stability (variance).\n    4. Momentum for variance reduction (post-processing).\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.target_quantile = 0.5  # target median\n        self.clip_lr = 0.2  # geometric update learning rate\n        self.clip_min = 0.1\n        self.clip_max = 10.0\n        self.clip_norm = max_grad_norm\n        \n        # Momentum parameters\n        self.momentum_beta = 0.9\n        self.momentum_buffers = None\n        \n        # Track gradient norm history for stability\n        self.gnorm_history = []\n        self.gnorm_decay = 0.9\n        self.gnorm_ema = None\n\n    def clip_and_noise(self, per_sample_grads, step, epoch):\n        batch_size = per_sample_grads[0].shape[0]\n\n        # Compute per-sample gradient norms (global across all parameters)\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        # Update clipping threshold using geometric update\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        # Update EMA of gradient norms for stability estimation\n        batch_mean_norm = norms.mean().item()\n        if self.gnorm_ema is None:\n            self.gnorm_ema = batch_mean_norm\n        else:\n            self.gnorm_ema = self.gnorm_decay * self.gnorm_ema + (1 - self.gnorm_decay) * batch_mean_norm\n        self.gnorm_history.append(self.gnorm_ema)\n\n        # Compute stability factor: inverse of coefficient of variation\n        if len(self.gnorm_history) > 10:\n            recent = self.gnorm_history[-10:]\n            mean_gnorm = sum(recent) / len(recent)\n            var_gnorm = sum((x - mean_gnorm) ** 2 for x in recent) / len(recent)\n            if var_gnorm > 0:\n                stability = mean_gnorm / (math.sqrt(var_gnorm) + 1e-8)\n                # normalize stability factor between 0.5 and 2.0\n                stability = max(0.5, min(2.0, stability / 5.0))\n            else:\n                stability = 1.0\n        else:\n            stability = 1.0\n\n        # Effective clipping threshold adjusted by stability\n        effective_clip = self.clip_norm * stability\n\n        # Clip per-sample gradients\n        clip_factor = (effective_clip / norms.clamp(min=1e-8)).clamp(max=1.0)\n\n        # Initialize momentum buffers if first call\n        if self.momentum_buffers is None:\n            self.momentum_buffers = [torch.zeros_like(g[0]) for g in per_sample_grads]\n\n        noised_grads = []\n        for i, g in enumerate(per_sample_grads):\n            shape = [batch_size] + [1] * (g.dim() - 1)\n            clipped = g * clip_factor.reshape(shape)\n\n            # Average over batch\n            avg = clipped.mean(dim=0)\n\n            # Add calibrated Gaussian noise\n            noise = torch.randn_like(avg) * (\n                self.noise_multiplier * effective_clip / batch_size\n            )\n            noised = avg + noise\n            \n            # Apply momentum (post-processing)\n            self.momentum_buffers[i] = self.momentum_beta * self.momentum_buffers[i] + (1 - self.momentum_beta) * noised\n            noised_grads.append(self.momentum_buffers[i])\n\n        return noised_grads\n\n    def get_effective_sigma(self, step, epoch):\n        \"\"\"Return the noise multiplier for privacy accounting.\n        \n        Privacy depends only on sigma, not clipping threshold.\n        \"\"\"\n        return self.noise_multiplier"
7}

1{
2  "op": "replace",
3  "filename": "opacus/custom_dpsgd.py",
4  "start_line": 172,
5  "end_line": 280,
6  "content": "class DPMechanism:\n    \"\"\"Adaptive Layer-wise Clipping with Noise Decay (ALC-ND).\n\n    Combines per-layer adaptive clipping with noise decay scheduling.\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        # Noise decay schedule (similar to noise_decay baseline)\n        self.decay_interval = max(1, epochs // 4)  # 4 decay stages\n        self.noise_decay_factor = 0.8  # Reduce noise by 20% at each stage\n        self.clip_decay_factor = 0.85  # Reduce clip norm by 15% at each stage\n\n        self.steps_per_epoch = dataset_size // batch_size\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        self.sigma_0 = noise_multiplier * (inv_sq_sum / total_steps) ** 0.5\n        self.clip_0 = max_grad_norm\n\n        # Per-layer adaptive clipping thresholds\n        self.layer_clip_norms = None  # will be list initialized in first call\n        self.target_quantile = 0.5\n        self.clip_lr = 0.1\n        self.clip_min = 0.01\n        self.clip_max = 10.0\n\n        # Current stage values\n        self._current_sigma = self.sigma_0\n        self._current_clip_base = self.clip_0\n\n    def clip_and_noise(self, per_sample_grads, step, epoch):\n        batch_size = per_sample_grads[0].shape[0]\n\n        # Update noise and clip schedule based on epoch\n        stage = (epoch - 1) // self.decay_interval\n        self._current_sigma = self.sigma_0 * (self.noise_decay_factor ** stage)\n        self._current_clip_base = self.clip_0 * (self.clip_decay_factor ** stage)\n\n        # Initialize per-layer clipping thresholds if first call\n        if self.layer_clip_norms is None:\n            self.layer_clip_norms = [self._current_clip_base for _ in per_sample_grads]\n\n        # Compute per-sample gradient norms per layer\n        layer_norms = []  # list of tensors [B] per layer\n        for g in per_sample_grads:\n            # flatten per-sample gradients for this layer\n            flat = g.reshape(batch_size, -1)\n            norms = flat.norm(2, dim=1)  # [B]\n            layer_norms.append(norms)\n\n        # Update per-layer clipping thresholds\n        for i, norms in enumerate(layer_norms):\n            frac_above = (norms > self.layer_clip_norms[i]).float().mean().item()\n            self.layer_clip_norms[i] = self.layer_clip_norms[i] * math.exp(\n                self.clip_lr * (frac_above - self.target_quantile)\n            )\n            self.layer_clip_norms[i] = max(self.clip_min, min(self.clip_max, self.layer_clip_norms[i]))\n\n        # Combine per-layer thresholds with global decay\n        effective_layer_clips = [self.layer_clip_norms[i] * self._current_clip_base / self.clip_0\n                                 for i in range(len(per_sample_grads))]\n        # Ensure each effective clip is within bounds\n        effective_layer_clips = [max(self.clip_min, min(self.clip_max, c)) for c in effective_layer_clips]\n\n        # Clip per-sample gradients per layer\n        noised_grads = []\n        for i, g in enumerate(per_sample_grads):\n            clip_factor = (effective_layer_clips[i] / layer_norms[i].clamp(min=1e-8)).clamp(max=1.0)\n            shape = [batch_size] + [1] * (g.dim() - 1)\n            clipped = g * clip_factor.reshape(shape)\n\n            # Average over batch\n            avg = clipped.mean(dim=0)\n\n            # Add noise calibrated to current layer clip and sigma\n            noise = torch.randn_like(avg) * (\n                self._current_sigma * effective_layer_clips[i] / batch_size\n            )\n            noised_grads.append(avg + noise)\n\n        return noised_grads\n\n    def get_effective_sigma(self, step, epoch):\n        \"\"\"Return effective noise multiplier for privacy accounting.\n        \n        The noise multiplier varies per step; return the current sigma.\n        Privacy accounting uses this per-step sigma.\n        \"\"\"\n        stage = (epoch - 1) // self.decay_interval\n        return self.sigma_0 * (self.noise_decay_factor ** stage)"
7}