Accounting for 2C sensitivity when doing microbatches

PiperOrigin-RevId: 494792995

Author: nataliaponomareva

tensorflow_privacy/privacy/keras_models/dp_keras_model.py

@@ -82,6 +82,9 @@ def __init__(
       super().__init__(*args, **kwargs)
       self._l2_norm_clip = l2_norm_clip
       self._noise_multiplier = noise_multiplier
+      # For microbatching version, the sensitivity is 2*l2_norm_clip.
+      self._sensitivity_multiplier = 2.0 if (num_microbatches is not None and
+                                             num_microbatches > 1) else 1.0
 
       # Given that `num_microbatches` was added as an argument after the fact,
       # this check helps detect unintended calls to the earlier API.
@@ -109,7 +112,7 @@ def _process_per_example_grads(self, grads):
 
     def _reduce_per_example_grads(self, stacked_grads):
       summed_grads = tf.reduce_sum(input_tensor=stacked_grads, axis=0)
-      noise_stddev = self._l2_norm_clip * self._noise_multiplier
+      noise_stddev = self._l2_norm_clip * self._sensitivity_multiplier * self._noise_multiplier
       noise = tf.random.normal(
           tf.shape(input=summed_grads), stddev=noise_stddev)
       noised_grads = summed_grads + noise

tensorflow_privacy/privacy/optimizers/dp_optimizer.py

@@ -340,8 +340,13 @@ def __init__(
       self._num_microbatches = num_microbatches
       self._base_optimizer_class = cls
 
+      # For microbatching version, the sensitivity is 2*l2_norm_clip.
+      sensitivity_multiplier = 2.0 if (num_microbatches is not None and
+                                       num_microbatches > 1) else 1.0
+
       dp_sum_query = gaussian_query.GaussianSumQuery(
-          l2_norm_clip, l2_norm_clip * noise_multiplier)
+          l2_norm_clip,
+          sensitivity_multiplier * l2_norm_clip * noise_multiplier)
 
       super(DPGaussianOptimizerClass,
             self).__init__(dp_sum_query, num_microbatches, unroll_microbatches,

tensorflow_privacy/privacy/optimizers/dp_optimizer_keras.py

@@ -459,8 +459,12 @@ def return_gaussian_query_optimizer(
       *args: These will be passed on to the base class `__init__` method.
       **kwargs: These will be passed on to the base class `__init__` method.
     """
+    # For microbatching version, the sensitivity is 2*l2_norm_clip.
+    sensitivity_multiplier = 2.0 if (num_microbatches is not None and
+                                     num_microbatches > 1) else 1.0
+
     dp_sum_query = gaussian_query.GaussianSumQuery(
-        l2_norm_clip, l2_norm_clip * noise_multiplier)
+        l2_norm_clip, sensitivity_multiplier * l2_norm_clip * noise_multiplier)
     return cls(
         dp_sum_query=dp_sum_query,
         num_microbatches=num_microbatches,

tensorflow_privacy/privacy/optimizers/dp_optimizer_keras_sparse.py

@@ -185,13 +185,18 @@ def __init__(
       self._num_microbatches = num_microbatches
       self._was_dp_gradients_called = False
       self._noise_stddev = None
+      # For microbatching version, the sensitivity is 2*l2_norm_clip.
+      self._sensitivity_multiplier = 2.0 if (num_microbatches is not None and
+                                             num_microbatches > 1) else 1.0
+
       if self._num_microbatches is not None:
         # The loss/gradients is the mean over the microbatches so we
         # divide the noise by num_microbatches too to obtain the correct
         # normalized noise.  If _num_microbatches is not set, the noise stddev
         # will be set later when the loss is given.
-        self._noise_stddev = (self._l2_norm_clip * self._noise_multiplier /
-                              self._num_microbatches)
+        self._noise_stddev = (
+            self._l2_norm_clip * self._noise_multiplier *
+            self._sensitivity_multiplier / self._num_microbatches)
 
     def _generate_noise(self, g):
       """Returns noise to be added to `g`."""

tensorflow_privacy/privacy/optimizers/dp_optimizer_vectorized.py

@@ -104,6 +104,9 @@ def __init__(
       self._noise_multiplier = noise_multiplier
       self._num_microbatches = num_microbatches
       self._was_compute_gradients_called = False
+      # For microbatching version, the sensitivity is 2*l2_norm_clip.
+      self._sensitivity_multiplier = 2.0 if (num_microbatches is not None and
+                                             num_microbatches > 1) else 1.0
 
     def compute_gradients(self,
                           loss,
@@ -166,7 +169,7 @@ def process_microbatch(microbatch_loss):
 
         def reduce_noise_normalize_batch(stacked_grads):
           summed_grads = tf.reduce_sum(input_tensor=stacked_grads, axis=0)
-          noise_stddev = self._l2_norm_clip * self._noise_multiplier
+          noise_stddev = self._l2_norm_clip * self._noise_multiplier * self._sensitivity_multiplier
           noise = tf.random.normal(
               tf.shape(input=summed_grads), stddev=noise_stddev)
           noised_grads = summed_grads + noise