Add PLD support for Quantiles (#607)

Author: dvadym

pipeline_dp/combiners.py

@@ -673,11 +673,13 @@ def compute_metrics(self, accumulator: AccumulatorType) -> AccumulatorType:
             tree = self._create_empty_quantile_tree()
             tree.merge(pydp._pydp.bytes_to_summary(accumulator))
 
-        quantiles = tree.compute_quantiles(
-            self._params.eps, self._params.delta,
+        quantiles = dp_computations.compute_dp_quantiles(
+            tree,
+            self._params.mechanism_spec,
+            self._quantiles_to_compute,
             self._params.aggregate_params.max_partitions_contributed,
             self._params.aggregate_params.max_contributions_per_partition,
-            self._quantiles_to_compute, self._noise_type())
+        )
 
         return dict([(name, value)
                      for name, value in zip(self.metrics_names(), quantiles)])
@@ -707,14 +709,6 @@ def _create_empty_quantile_tree(self):
             self._params.aggregate_params.max_value, DEFAULT_TREE_HEIGHT,
             DEFAULT_BRANCHING_FACTOR)
 
-    def _noise_type(self) -> str:
-        noise_kind = self._params.aggregate_params.noise_kind
-        if noise_kind == pipeline_dp.NoiseKind.LAPLACE:
-            return "laplace"
-        if noise_kind == pipeline_dp.NoiseKind.GAUSSIAN:
-            return "gaussian"
-        assert False, f"{noise_kind} is not support by PyDP quantile tree."
-
     def mechanism_spec(self) -> budget_accounting.MechanismSpec:
         return self._params.mechanism_spec
 

pipeline_dp/dp_computations.py

@@ -25,6 +25,7 @@
 from pipeline_dp import budget_accounting
 from pipeline_dp import partition_selection
 from pydp.algorithms import numerical_mechanisms as dp_mechanisms
+from pydp.algorithms import quantile_tree
 
 
 @dataclass
@@ -883,3 +884,56 @@ def create_thresholding_mechanism(
                                  strategy=strategy,
                                  l0_sensitivity=sensitivities.l0,
                                  pre_threshold=pre_threshold)
+
+
+def compute_dp_quantiles(
+    quantile_tree: quantile_tree.QuantileTree,
+    mechanism_spec: budget_accounting.MechanismSpec,
+    quantiles_to_compute: list[float],
+    max_partitions_contributed: int,
+    max_contributions_per_partition: int,
+) -> list[float]:
+    """Computes DP quantiles for the given quantile tree and quantiles to compute.
+
+    Args:
+      quantile_tree: The quantile tree to compute quantiles for.
+      mechanism_spec: The mechanism spec to use for the computation.
+      quantiles_to_compute: The quantiles to compute.
+      max_partitions_contributed: The maximum number of partitions contributed
+        per privacy unit.
+      max_contributions_per_partition: The maximum number of contributions per
+        partition.
+
+     Returns:
+        The computed DP quantiles, which are the same length as
+        quantiles_to_compute.
+    """
+
+    noise_kind = mechanism_spec.mechanism_type.to_noise_kind()
+
+    if mechanism_spec.standard_deviation_is_set:
+        # QuantileTree supports setting DP guarantees only via (epsilon, delta).
+        # But in PLD accounting we find the standard deviation of the equivalent
+        # Laplace or Gaussian mechanism. Let us convert the standard deviation
+        # to (epsilon, delta) here such that the corresponding Laplace or Gaussian
+        # mechanism (which is used internally by QuantileTree) has the same
+        # standard deviation.
+        sigma = mechanism_spec.noise_standard_deviation
+        if noise_kind == aggregate_params.NoiseKind.LAPLACE:
+            epsilon = np.sqrt(2) / sigma
+            delta = 0.0
+        else:
+            epsilon = 1.0
+            delta = gaussian_delta(sigma, epsilon)
+    else:
+        epsilon = mechanism_spec.eps
+        delta = mechanism_spec.delta
+
+    return quantile_tree.compute_quantiles(
+        epsilon=epsilon,
+        delta=delta,
+        max_partitions_contributed=max_partitions_contributed,
+        max_contributions_per_partition=max_contributions_per_partition,
+        quantiles=quantiles_to_compute,
+        noise_type=noise_kind.value,
+    )

pipeline_dp/dp_engine.py

@@ -556,15 +556,12 @@ def _check_budget_accountant_compatibility(self, is_public_partition: bool,
                       budget_accounting.NaiveBudgetAccountant):
             # All aggregations support NaiveBudgetAccountant.
             return
-        supported_metrics = [
-            aggregate_params.Metrics.COUNT,
-            aggregate_params.Metrics.PRIVACY_ID_COUNT,
-            aggregate_params.Metrics.SUM, aggregate_params.Metrics.MEAN
-        ]
-        non_supported_metrics = set(metrics) - set(supported_metrics)
-        if non_supported_metrics:
-            raise NotImplementedError(f"Metrics {non_supported_metrics} do not "
-                                      f"support PLD budget accounting")
+        if aggregate_params.Metrics.VECTOR_SUM in metrics:
+            raise NotImplementedError(
+                "Vector Sum is not supported by PLD budget accounting")
+        if aggregate_params.Metrics.VARIANCE in metrics:
+            raise NotImplementedError(
+                "Variance is not supported by PLD budget accounting")
         if custom_combiner:
             raise ValueError(f"PLD budget accounting does not support custom "
                              f"combiners")

tests/dp_computations_test.py

@@ -1045,5 +1045,110 @@ def is_monotonic(self) -> bool:
         self.assertSetEqual(chosen_values, {0, 1})
 
 
+class ComputeDpQuantilesTest(parameterized.TestCase):
+
+    def test_compute_dp_quantiles_std_dev_set_laplace(self):
+        # Arrange
+        mock_tree = MagicMock()
+        mock_tree.compute_quantiles.return_value = [10.0, 20.0]
+
+        spec = budget_accounting.MechanismSpec(
+            aggregate_params.MechanismType.LAPLACE)
+        spec.set_noise_standard_deviation(2.0)
+
+        quantiles = [0.5, 0.9]
+        max_partitions = 5
+        max_contributions = 2
+
+        # Act
+        result = dp_computations.compute_dp_quantiles(
+            quantile_tree=mock_tree,
+            mechanism_spec=spec,
+            quantiles_to_compute=quantiles,
+            max_partitions_contributed=max_partitions,
+            max_contributions_per_partition=max_contributions,
+        )
+
+        # Assert
+        expected_eps = np.sqrt(2) / 2.0
+        mock_tree.compute_quantiles.assert_called_once_with(
+            epsilon=expected_eps,
+            delta=0.0,
+            max_partitions_contributed=max_partitions,
+            max_contributions_per_partition=max_contributions,
+            quantiles=quantiles,
+            noise_type="laplace",
+        )
+        self.assertEqual(result, [10.0, 20.0])
+
+    def test_compute_dp_quantiles_std_dev_set_gaussian(self):
+        # Arrange
+        mock_tree = MagicMock()
+        mock_tree.compute_quantiles.return_value = [10.0, 20.0]
+
+        spec = budget_accounting.MechanismSpec(
+            aggregate_params.MechanismType.GAUSSIAN)
+        spec.set_noise_standard_deviation(2.0)
+
+        quantiles = [0.5, 0.9]
+        max_partitions = 5
+        max_contributions = 2
+
+        # Act
+        result = dp_computations.compute_dp_quantiles(
+            quantile_tree=mock_tree,
+            mechanism_spec=spec,
+            quantiles_to_compute=quantiles,
+            max_partitions_contributed=max_partitions,
+            max_contributions_per_partition=max_contributions,
+        )
+
+        # Assert
+        expected_eps = 1.0
+        expected_delta = dp_computations.gaussian_delta(2.0, 1.0)
+        mock_tree.compute_quantiles.assert_called_once_with(
+            epsilon=expected_eps,
+            delta=expected_delta,
+            max_partitions_contributed=max_partitions,
+            max_contributions_per_partition=max_contributions,
+            quantiles=quantiles,
+            noise_type="gaussian",
+        )
+        self.assertEqual(result, [10.0, 20.0])
+
+    def test_compute_dp_quantiles_std_dev_not_set_gaussian(self):
+        # Arrange
+        mock_tree = MagicMock()
+        mock_tree.compute_quantiles.return_value = [10.0, 20.0]
+
+        spec = budget_accounting.MechanismSpec(
+            aggregate_params.MechanismType.GAUSSIAN)
+        spec.set_eps_delta(1.0, 1e-5)
+
+        quantiles = [0.5, 0.9]
+        max_partitions = 5
+        max_contributions = 2
+
+        # Act
+        result = dp_computations.compute_dp_quantiles(
+            quantile_tree=mock_tree,
+            mechanism_spec=spec,
+            quantiles_to_compute=quantiles,
+            max_partitions_contributed=max_partitions,
+            max_contributions_per_partition=max_contributions,
+        )
+
+        # Assert
+        mock_tree.compute_quantiles.assert_called_once_with(
+            epsilon=1.0,
+            delta=1e-5,
+            max_partitions_contributed=max_partitions,
+            max_contributions_per_partition=max_contributions,
+            quantiles=quantiles,
+            noise_type="gaussian",
+        )
+        self.assertEqual(result, [10.0, 20.0])
+
+
 if __name__ == '__main__':
     absltest.main()

tests/dp_engine_test.py

@@ -1307,13 +1307,10 @@ def test_min_max_sum_per_partition(self):
         self.assertLen(output, 1)
         self.assertAlmostEqual(output[0][1].sum, -3, delta=0.1)
 
-    @unittest.skipIf(
-        sys.version_info.major == 3 and sys.version_info.minor <= 8,
-        "dp_accounting library only support python >=3.9")
     def test_pld_not_supported_metrics(self):
         with self.assertRaisesRegex(
                 NotImplementedError,
-                "Metrics {VARIANCE} do not support PLD budget accounting"):
+                "Variance is not supported by PLD budget accounting"):
             budget_accountant = pipeline_dp.PLDBudgetAccountant(
                 total_epsilon=1, total_delta=1e-10)
             engine = pipeline_dp.DPEngine(