From f7351edd7d32ca887019f785c17c5c17a4986355 Mon Sep 17 00:00:00 2001 From: Ryan McKenna Date: Thu, 30 Jul 2026 11:35:49 -0700 Subject: [PATCH] Add user-level DP support to the local-mode synthesizer. Introduce a `max_records_per_user` (k) parameter across the discrete mechanisms, local-mode primitives/initializers, and `TabularSynthesizer`, so a single user may contribute up to k rows; the noise for one-way, measurement, and total-count queries is scaled accordingly. The default k=1 preserves existing behavior. Also surface the `TabularCodec` on `DataGenerationResult` so callers can map model axes back to labels. PiperOrigin-RevId: 956644098 --- dpsynth/api.py | 6 +++ dpsynth/data_generation_v3.py | 64 ++++++++++++++++++++++------ dpsynth/discrete_mechanisms/base.py | 16 ++++++- dpsynth/discrete_mechanisms/mst.py | 13 ++++-- dpsynth/local_mode/beam_adapter.py | 1 + dpsynth/local_mode/initialization.py | 39 ++++++++++++++++- dpsynth/local_mode/primitives.py | 33 ++++++++++++-- 7 files changed, 150 insertions(+), 22 deletions(-) diff --git a/dpsynth/api.py b/dpsynth/api.py index 0b9383ac..d1e7ddca 100644 --- a/dpsynth/api.py +++ b/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}.') diff --git a/dpsynth/data_generation_v3.py b/dpsynth/data_generation_v3.py index e40d0d8a..9736a89e 100644 --- a/dpsynth/data_generation_v3.py +++ b/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,6 +38,7 @@ 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. @@ -44,6 +46,8 @@ def _create_initializers( 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, ) diff --git a/dpsynth/discrete_mechanisms/base.py b/dpsynth/discrete_mechanisms/base.py index 9695d4b6..cb4ceff7 100644 --- a/dpsynth/discrete_mechanisms/base.py +++ b/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 ) diff --git a/dpsynth/discrete_mechanisms/mst.py b/dpsynth/discrete_mechanisms/mst.py index 1c86072e..946cece2 100644 --- a/dpsynth/discrete_mechanisms/mst.py +++ b/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, ) diff --git a/dpsynth/local_mode/beam_adapter.py b/dpsynth/local_mode/beam_adapter.py index d794c7d2..44e288fe 100644 --- a/dpsynth/local_mode/beam_adapter.py +++ b/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, ) diff --git a/dpsynth/local_mode/initialization.py b/dpsynth/local_mode/initialization.py index a5ac8f93..ca0ef9cb 100644 --- a/dpsynth/local_mode/initialization.py +++ b/dpsynth/local_mode/initialization.py @@ -21,6 +21,7 @@ from typing import TypeVar import dp_accounting +from dpsynth import api from dpsynth import domain from dpsynth.local_mode import _quantiles from dpsynth.local_mode import primitives @@ -87,17 +88,21 @@ class NumericalInitializer(primitives.DPMechanism): name: Attribute name used as the clique key in the measurement. num_partitions: Number of quantile partitions (must be a power of 2). attribute: The NumericalAttribute defining the data domain. + max_records_per_user: Assumed upper bound on the number of records a single + user contributes; forwarded to the underlying ``DPQuantiles`` mechanism. """ name: str num_partitions: int attribute: domain.NumericalAttribute max_grid_size: int = 10_000_000 + max_records_per_user: int = 1 mechanism: primitives.DPQuantiles | None = dataclasses.field( default=None, repr=False ) def __post_init__(self): + api.validate_max_records_per_user(self.max_records_per_user) if self.max_grid_size < 2: raise ValueError(f'max_grid_size must be >= 2, got {self.max_grid_size}.') @@ -132,6 +137,7 @@ def configure( # pyrefly: ignore[bad-override] jitter_strategy=( 'refine' if self.attribute.dtype == 'int' else 'symmetric' ), + max_records_per_user=self.max_records_per_user, ).configure(zcdp_rho=zcdp_rho, epsilon_ratio=epsilon_ratio) return dataclasses.replace(self, mechanism=mechanism) @@ -195,6 +201,7 @@ def from_summary( name=self.name, zcdp_rho=mechanism.zcdp_rho, estimated_total=estimated_total, + max_records_per_user=self.max_records_per_user, ) @@ -204,6 +211,7 @@ def edges_to_column_measurement( name, zcdp_rho, estimated_total=None, + max_records_per_user=1, ): """Converts raw quantile edges into a ColumnMeasurement. @@ -218,6 +226,8 @@ def edges_to_column_measurement( name: Attribute name used as the clique key in any measurement. zcdp_rho: Total zCDP rho consumed by the quantile mechanism. estimated_total: If provided, a heuristic one-way measurement is included. + max_records_per_user: Assumed upper bound on the number of records a single + user contributes; inflates the heuristic measurement's stddev. Returns: A ``ColumnMeasurement`` with bin edges and optionally a measurement. @@ -242,7 +252,7 @@ def edges_to_column_measurement( # Prepend zero weight for the OUT_OF_DOMAIN slot at index 0. bin_weights = np.r_[0, bin_weights] counts = estimated_total * bin_weights / bin_weights.sum() - stddev = 1.0 / np.sqrt(zcdp_rho) + stddev = max_records_per_user / np.sqrt(zcdp_rho) measurement = mbi.LinearMeasurement( counts, (name,), @@ -263,20 +273,27 @@ class CategoricalInitializer(primitives.DPMechanism): Attributes: name: Attribute name used as the clique key in the measurement. attribute: The CategoricalAttribute defining the closed domain. + max_records_per_user: Assumed upper bound on the number of records a single + user contributes; forwarded to the underlying ``DPGaussianHistogram``. """ name: str attribute: domain.CategoricalAttribute + max_records_per_user: int = 1 mechanism: primitives.DPGaussianHistogram | None = dataclasses.field( default=None, repr=False ) + def __post_init__(self): + api.validate_max_records_per_user(self.max_records_per_user) + def configure( # pyrefly: ignore[bad-override] self, *, zcdp_rho: float, delta: float = 0.0 ) -> CategoricalInitializer: """Returns a copy calibrated to the given zCDP budget.""" mechanism = primitives.DPGaussianHistogram( domain_size=self.attribute.size, + max_records_per_user=self.max_records_per_user, ).configure(zcdp_rho=zcdp_rho) return dataclasses.replace(self, mechanism=mechanism) @@ -300,7 +317,9 @@ def from_summary( mechanism = _validate_mechanism(self.mechanism) result = mechanism(rng, counts) measurement = mbi.LinearMeasurement( - result.counts, (self.name,), stddev=mechanism.sigma # pyrefly: ignore[bad-argument-type] + result.counts, + (self.name,), + stddev=mechanism.max_records_per_user * mechanism.sigma, # pyrefly: ignore[bad-argument-type, unsupported-operation] ) return ColumnMeasurement(self.attribute, measurement=measurement) @@ -320,16 +339,32 @@ class OpenSetCategoricalInitializer(primitives.DPMechanism): attribute: The OpenSetCategoricalAttribute specifying the default value. delta: Failure probability for the partition selection threshold. min_count: Minimum true count for a partition to be discovered. + max_records_per_user: Assumed upper bound on the number of records a single + user contributes. Values greater than 1 are not yet supported for open-set + columns (see __post_init__). """ name: str attribute: domain.OpenSetCategoricalAttribute delta: float min_count: int = 1 + max_records_per_user: int = 1 mechanism: primitives.DPPartitionSelection | None = dataclasses.field( default=None, repr=False ) + def __post_init__(self): + api.validate_max_records_per_user(self.max_records_per_user) + if self.max_records_per_user > 1: + # Bounding user contribution for open-set partition selection without a + # user->record mapping would require a conservative worst-case threshold + # with poor utility. Proper support is deferred to the user-id-aware path. + raise ValueError( + 'max_records_per_user > 1 is not yet supported for open-set ' + 'categorical attributes; it requires a record-to-user mapping for ' + 'sound, good-utility partition selection.' + ) + def configure( # pyrefly: ignore[bad-override] self, *, zcdp_rho: float, delta: float = 0.0 ) -> OpenSetCategoricalInitializer: diff --git a/dpsynth/local_mode/primitives.py b/dpsynth/local_mode/primitives.py index 88792c51..10bd1f86 100644 --- a/dpsynth/local_mode/primitives.py +++ b/dpsynth/local_mode/primitives.py @@ -277,16 +277,25 @@ class DPQuantiles(DPMechanism): upper: Upper bound of the data domain (exclusive). jitter_strategy: Tie-breaking jitter passed to ``quantiles_from_histogram``: ``'refine'`` for integer attributes, ``'symmetric'`` for continuous ones. + max_records_per_user: Assumed upper bound on the number of records a single + user contributes. The per-level exponential-mechanism epsilon is divided + by this factor (the quality-score sensitivity grows linearly in it); the + reported dp_event is unchanged. Soundness relies on the caller enforcing + this bound. """ num_partitions: int lower: float upper: float jitter_strategy: Literal['symmetric', 'refine'] = 'symmetric' + max_records_per_user: int = 1 _epsilon_levels: tuple[float, ...] | None = dataclasses.field( default=None, repr=False ) + def __post_init__(self): + api.validate_max_records_per_user(self.max_records_per_user) + @property def _num_levels(self) -> int: result = int(np.log2(self.num_partitions)) @@ -341,7 +350,9 @@ def __call__( indices = _quantiles.quantiles_from_histogram( rng, counts, - epsilon_levels=np.asarray(self._epsilon_levels), + epsilon_levels=( + np.asarray(self._epsilon_levels) / self.max_records_per_user + ), jitter_strategy=self.jitter_strategy, ) # Map cell indices back to domain values; delta is the grid step, which @@ -360,10 +371,18 @@ class DPGaussianHistogram(DPMechanism): Attributes: domain_size: Number of categories in the histogram domain. sigma: Gaussian noise standard deviation. Set directly or via ``configure``. + max_records_per_user: Assumed upper bound on the number of records a single + user contributes. The noise standard deviation is scaled by this factor + (the count-vector sensitivity grows linearly in it); the reported dp_event + is unchanged. Soundness relies on the caller enforcing this bound. """ domain_size: int sigma: float | None = None + 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, *, zcdp_rho: float, delta: float = 0.0 @@ -384,7 +403,9 @@ def __call__( """Adds Gaussian noise to the given counts.""" if self.sigma is None: raise ValueError(_UNCALIBRATED_MSG.format(param='sigma')) - noise = rng.normal(scale=self.sigma, size=self.domain_size) + noise = rng.normal( + scale=self.max_records_per_user * self.sigma, size=self.domain_size + ) return HistogramResult(counts=counts.astype(float) + noise) @@ -393,6 +414,10 @@ class DPGaussianCount(DPMechanism): """Differentially private count via the Gaussian mechanism.""" sigma: float | None = None + 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, *, zcdp_rho: float, delta: float = 0.0 @@ -411,7 +436,9 @@ def noisy_count(self, rng: np.random.Generator, true_count: int) -> float: """Returns ``true_count`` plus calibrated Gaussian noise.""" if self.sigma is None: raise ValueError(_UNCALIBRATED_MSG.format(param='sigma')) - return float(true_count + rng.normal(scale=self.sigma)) + return float( + true_count + rng.normal(scale=self.max_records_per_user * self.sigma) + ) def __call__(self, rng: np.random.Generator, data: np.ndarray) -> float: """Returns a noisy count of len(data) + Gaussian noise."""