Add user-level DP support to the local-mode synthesizer.

dpsynth/api.py

@@ -241,3 +241,9 @@ def calibrate(
         target_delta=delta,
     )
     return self.configure(zcdp_rho=optimal_rho, delta=delta, **kwargs)
+
+
+def validate_max_records_per_user(value: int) -> None:
+  """Raises ValueError if the per-user record bound is not a positive int."""
+  if value < 1:
+    raise ValueError(f'max_records_per_user must be >= 1, got {value}.')

dpsynth/data_generation_v3.py

@@ -21,6 +21,7 @@
 
 from absl import logging
 import dp_accounting
+from dpsynth import api
 from dpsynth import constraints
 from dpsynth import discrete_mechanisms
 from dpsynth import domain
@@ -37,13 +38,16 @@ def _create_initializers(
     domains: Mapping[str, domain.AttributeType],
     numerical_bins: int,
     init_delta: float,
+    max_records_per_user: int = 1,
 ) -> dict[str, primitives.DPMechanism]:
   """Creates per-column initializers from the domain specification.
 
   Args:
     domains: Mapping from column names to attribute domain specifications.
     numerical_bins: Number of bins for numerical discretization.
     init_delta: Delta for open-set categorical partition selection.
+    max_records_per_user: Assumed upper bound on the number of records a single
+      user contributes; forwarded to each initializer.
 
   Returns:
     A dictionary mapping column names to uncalibrated initializer instances.
@@ -55,15 +59,21 @@ def _create_initializers(
   for col, attr in domains.items():
     if isinstance(attr, domain.NumericalAttribute):
       initializers[col] = initialization.NumericalInitializer(
-          name=col, num_partitions=numerical_bins, attribute=attr
+          name=col,
+          num_partitions=numerical_bins,
+          attribute=attr,
+          max_records_per_user=max_records_per_user,
       )
     elif isinstance(attr, domain.CategoricalAttribute):
       initializers[col] = initialization.CategoricalInitializer(
-          name=col, attribute=attr
+          name=col, attribute=attr, max_records_per_user=max_records_per_user
       )
     elif isinstance(attr, domain.OpenSetCategoricalAttribute):
       initializers[col] = initialization.OpenSetCategoricalInitializer(
-          name=col, attribute=attr, delta=init_delta
+          name=col,
+          attribute=attr,
+          delta=init_delta,
+          max_records_per_user=max_records_per_user,
       )
     else:
       raise ValueError(
@@ -172,10 +182,19 @@ def decode(
 
 @dataclasses.dataclass
 class DataGenerationResult:
-  """Result of end-to-end DP synthetic data generation."""
+  """Result of end-to-end DP synthetic data generation.
+
+  Attributes:
+    synthetic_data: The generated synthetic data in the original domain.
+    discrete_mechanism_result: The raw result of the discrete mechanism run on
+      the discretized data.
+    codec: The codec mapping columns between raw values and the discrete domain
+      used to encode/decode the data.
+  """
 
   synthetic_data: pd.DataFrame
   discrete_mechanism_result: dm_common.DiscreteMechanismResult
+  codec: TabularCodec
 
 
 @dataclasses.dataclass
@@ -204,6 +223,12 @@ class TabularSynthesizer(primitives.DPMechanism):
       automatically from ``domains`` during ``configure()``.
     skip_compression: Whether to skip domain compression.
     cross_attribute_constraints: Constraints to enforce on generated data.
+    max_records_per_user: Assumed upper bound on the number of records a single
+      user contributes. All added noise (and selection sensitivity) is scaled by
+      this factor so the mechanism provides user-level (rather than
+      record-level) DP under the stated ``dp_event``. Soundness relies on the
+      caller enforcing the bound; open-set categorical columns require this to
+      be 1.
   """
 
   domains: Mapping[str, domain.AttributeType]
@@ -215,6 +240,10 @@ class TabularSynthesizer(primitives.DPMechanism):
   initializers: dict[str, primitives.DPMechanism] | None = None
   total_count_mechanism: primitives.DPGaussianCount | None = None
   cross_attribute_constraints: Sequence[constraints.Constraint] = ()
+  max_records_per_user: int = 1
+
+  def __post_init__(self):
+    api.validate_max_records_per_user(self.max_records_per_user)
 
   def configure(  # pyrefly: ignore[bad-override]
       self,
@@ -269,8 +298,16 @@ def configure(  # pyrefly: ignore[bad-override]
         thresholding_delta / num_open_set if num_open_set > 0 else 0.0
     )
 
+    if self.initializers is not None and self.max_records_per_user > 1:
+      raise ValueError(
+          'max_records_per_user > 1 requires auto-created initializers; set '
+          'max_records_per_user on each custom initializer directly instead.'
+      )
     inits = self.initializers or _create_initializers(
-        self.domains, self.numerical_bins, per_col_delta
+        self.domains,
+        self.numerical_bins,
+        per_col_delta,
+        self.max_records_per_user,
     )
     init_rho = self.init_budget_fraction * zcdp_rho
     # +1 for the DPGaussianCount that always measures the total.
@@ -280,12 +317,12 @@ def configure(  # pyrefly: ignore[bad-override]
     calibrated_inits = {
         col: init.configure(zcdp_rho=per_col_rho) for col, init in inits.items()
     }
-    calibrated_total = primitives.DPGaussianCount().configure(
-        zcdp_rho=per_col_rho
-    )
-    calibrated_discrete = self.discrete_mechanism.configure(
-        zcdp_rho=discrete_rho
-    )
+    calibrated_total = primitives.DPGaussianCount(
+        max_records_per_user=self.max_records_per_user
+    ).configure(zcdp_rho=per_col_rho)
+    calibrated_discrete = dataclasses.replace(
+        self.discrete_mechanism, max_records_per_user=self.max_records_per_user
+    ).configure(zcdp_rho=discrete_rho)
     return dataclasses.replace(
         self,
         initializers=calibrated_inits,
@@ -344,7 +381,9 @@ def __call__(
     any_col = next(iter(self.domains))
     total = max(1.0, self.total_count_mechanism(rng, data[any_col].values))
     total_measurement = mbi.LinearMeasurement(
-        np.array([total]), (), stddev=self.total_count_mechanism.sigma  # pyrefly: ignore[bad-argument-type]
+        np.array([total]),
+        (),
+        stddev=(self.max_records_per_user * self.total_count_mechanism.sigma),  # pyrefly: ignore[bad-argument-type, unsupported-operation]
     )
 
     results: dict[str, initialization.ColumnMeasurement] = {}
@@ -383,4 +422,5 @@ def __call__(
     return DataGenerationResult(
         synthetic_data=synthetic_data,
         discrete_mechanism_result=mechanism_result,
+        codec=codec,
     )

dpsynth/discrete_mechanisms/base.py

@@ -62,6 +62,10 @@ class DiscreteMechanism(api.DPMechanism):
     zcdp_rho: Total zCDP budget (set by configure).
     one_way_rho: zCDP budget for one-way measurements (set by configure).
     measurement_rho: zCDP budget for selected marginal measurements.
+    max_records_per_user: Assumed upper bound on the number of records a single
+      user contributes. The noise added to counts (and the exponential-mechanism
+      selection sensitivity) is scaled by this factor; the reported dp_event is
+      unchanged. Soundness relies on the caller enforcing this bound.
   """
 
   marginal_oracle: mbi.MarginalOracle | None = None
@@ -71,6 +75,10 @@ class DiscreteMechanism(api.DPMechanism):
   zcdp_rho: float | None = None
   one_way_rho: float | None = dataclasses.field(default=None, repr=False)
   measurement_rho: float | None = dataclasses.field(default=None, repr=False)
+  max_records_per_user: int = 1
+
+  def __post_init__(self):
+    api.validate_max_records_per_user(self.max_records_per_user)
 
   @abc.abstractmethod
   def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]:
@@ -162,7 +170,9 @@ def _measure_one_way(
     if self.one_way_rho is None:
       return []
     with common.timed(phase_times, 'measurement'):
-      sigma = accounting.zcdp_gaussian_sigma(self.one_way_rho)
+      sigma = self.max_records_per_user * accounting.zcdp_gaussian_sigma(
+          self.one_way_rho
+      )
       cliques = self._one_way_cliques(data)
       return common.measure_marginals_with_noise(rng, data, cliques, sigma)
 
@@ -236,7 +246,9 @@ def _run(self, rng, data, measurements, constraints, phase_times):
 
     if selected:
       with common.timed(phase_times, 'measurement'):
-        sigma = accounting.zcdp_gaussian_sigma(self.measurement_rho)  # pyrefly: ignore[bad-argument-type]
+        sigma = self.max_records_per_user * accounting.zcdp_gaussian_sigma(
+            self.measurement_rho  # pyrefly: ignore[bad-argument-type]
+        )
         measurements = measurements + common.measure_marginals_with_noise(
             rng, data, selected, sigma
         )

dpsynth/discrete_mechanisms/mst.py

@@ -37,6 +37,7 @@ def dp_maximum_spanning_tree(
     zcdp_rho: float | None = None,
     exponential_mechanism_epsilon: float | None = None,
     initial_marginal_queries: Sequence[tuple[str, str]] = (),
+    sensitivity: float = 1.0,
 ) -> list[tuple[str, str]]:
   """Computes an approximate maximum spanning tree with differential privacy.
 
@@ -48,8 +49,9 @@ def dp_maximum_spanning_tree(
     2. otherwise, it has the same privacy guarantees as the len(weights)-1
      Exponential Mechanism with parameter exponential_mechanism_epsilon.
 
-  It is assumed that the weights are obtained from sensitivity 1 functions of
-  the data (i.e., L1 norm between true and estimated marginal).
+  It is assumed that the weights are obtained from sensitivity ``sensitivity``
+  functions of the data (i.e., L1 norm between true and estimated marginal,
+  scaled by the maximum number of records a single user contributes).
 
   Args:
     rng: A numpy random number generator.
@@ -60,6 +62,7 @@ def dp_maximum_spanning_tree(
       mechanism. If None, the value is computed from zcdp_rho.
     initial_marginal_queries: The list of initial attribute pairs to include in
       the tree.
+    sensitivity: The sensitivity of the quality scores in ``weights``.
 
   Returns:
     A list of attribute pairs that constitute an approximate maximum spanning
@@ -89,7 +92,7 @@ def dp_maximum_spanning_tree(
     candidates = [e for e in candidates if not ds.connected(*e)]
     wgts = np.array([weights[e] for e in candidates])
     idx = common.exponential_mechanism(
-        wgts, exponential_mechanism_epsilon, sensitivity=1.0, rng=rng
+        wgts, exponential_mechanism_epsilon, sensitivity=sensitivity, rng=rng
     )
     e = candidates[idx]
     tree.add_edge(*e)
@@ -105,6 +108,7 @@ def _select_two_way_marginal_queries(
     one_way_measurements: list[mbi.LinearMeasurement],
     initial_marginal_queries: Sequence[tuple[str, ...]] = (),
     maximum_marginal_size: int = 10_000_000,
+    sensitivity: float = 1.0,
 ) -> list[tuple[str, ...]]:
   """Selects a set of two-way marginal queries with DP to form a spanning tree.
 
@@ -118,6 +122,7 @@ def _select_two_way_marginal_queries(
     one_way_measurements: The initial one-way measurements already made.
     initial_marginal_queries: The list of cliques to start with.
     maximum_marginal_size: The maximum size of a marginal query.
+    sensitivity: The sensitivity of the correlation quality scores.
 
   Returns:
     A list of two-way marginal queries over highly correlated attributes.
@@ -145,6 +150,7 @@ def _select_two_way_marginal_queries(
       weights,  # pyrefly: ignore[bad-argument-type]
       zcdp_rho=zcdp_rho,
       initial_marginal_queries=initial_marginal_queries,  # pyrefly: ignore[bad-argument-type]
+      sensitivity=sensitivity,
   )
 
 
@@ -199,4 +205,5 @@ def _select(self, rng, data, measurements, phase_times):
           self._select_rho,  # pyrefly: ignore[bad-argument-type]
           measurements,
           maximum_marginal_size=self.maximum_marginal_size,
+          sensitivity=self.max_records_per_user,
       )

dpsynth/local_mode/beam_adapter.py

@@ -433,6 +433,7 @@ def generate_from_marginals(
   return data_generation_v3.DataGenerationResult(
       synthetic_data=synthetic_data,
       discrete_mechanism_result=mechanism_result,
+      codec=codec,
   )