Override add_noise in DPOptimizerFastGradientClipping to use adjusted noise multiplier for adaptive clipping

Author: david-stan

opacus/optimizers/adaclipoptimizer.py

@@ -27,7 +27,6 @@
     _mark_as_processed,
 )
 
-
 logger = logging.getLogger(__name__)
 
 

opacus/optimizers/optimizer_fast_gradient_clipping.py

@@ -21,7 +21,12 @@
 import torch
 from torch.optim import Optimizer
 
-from .optimizer import DPOptimizer
+from .optimizer import (
+    DPOptimizer,
+    _check_processed_flag,
+    _generate_noise,
+    _mark_as_processed,
+)
 
 
 logger = logging.getLogger(__name__)
@@ -188,3 +193,32 @@ def clip_and_accumulate(self):
         Redefines a parent class' function to not do anything
         """
         pass
+
+    def add_noise(self):
+        """
+        Adds noise to gradients. Stores noised result in ``p.grad``.
+
+        Uses ``_adjusted_noise_multiplier`` if set (for adaptive clipping),
+        otherwise falls back to ``noise_multiplier``. This separation ensures
+        correct privacy accounting when using adaptive clipping, where the
+        original noise_multiplier is preserved for accounting while the
+        adjusted value is used only for noise generation.
+        """
+        # Use adjusted noise multiplier if set (for adaptive clipping),
+        # otherwise use the standard noise_multiplier
+        effective_noise_multiplier = getattr(
+            self, "_adjusted_noise_multiplier", self.noise_multiplier
+        )
+
+        for p in self.params:
+            _check_processed_flag(p.summed_grad)
+
+            noise = _generate_noise(
+                std=effective_noise_multiplier * self.max_grad_norm,
+                reference=p.summed_grad,
+                generator=self.generator,
+                secure_mode=self.secure_mode,
+            )
+            p.grad = (p.summed_grad + noise).view_as(p)
+
+            _mark_as_processed(p.summed_grad)

opacus/utils/adaptive_clipping/adaptive_clipping_utils.py

@@ -95,10 +95,13 @@ def backward(self):
             per_sample_norms
         )
 
-        # update max grad norm and noise multiplier
+        # update max grad norm and adjusted noise multiplier for noise generation
+        # Note: We set _adjusted_noise_multiplier instead of noise_multiplier to preserve
+        # the original noise_multiplier for correct privacy accounting (see Theorem 1 in
+        # https://arxiv.org/pdf/1905.03871.pdf)
         self.module.max_grad_norm = new_max_grad_norm
         self.optimizer.max_grad_norm = new_max_grad_norm
-        self.optimizer.noise_multiplier = new_noise_multiplier
+        self.optimizer._adjusted_noise_multiplier = new_noise_multiplier
 
         # get the loss rescaling coefficients using the updated max_grad_norm
         coeff = torch.where(