From 155b1e470cc53bf689c7ca3cd210ecefe8aad0ea Mon Sep 17 00:00:00 2001 From: Ryan McKenna Date: Tue, 22 Sep 2026 15:03:49 -0700 Subject: [PATCH] Expose numerical_epsilon_ratio on TabularConfig and update default to 1.0. PiperOrigin-RevId: 986258941 --- dpsynth/data_generation_v3.py | 16 +++++++++++++- dpsynth/local_mode/initialization.py | 13 ++++++++++-- tests/data_generation_v3_test.py | 7 +++++++ tests/local_mode/initialization_test.py | 28 +++++++++++++++++++++++++ 4 files changed, 61 insertions(+), 3 deletions(-) diff --git a/dpsynth/data_generation_v3.py b/dpsynth/data_generation_v3.py index 27ab85e..f1f5bb0 100644 --- a/dpsynth/data_generation_v3.py +++ b/dpsynth/data_generation_v3.py @@ -38,12 +38,17 @@ def create_initializers( domains: domain.Schema | Mapping[str, domain.AttributeType], numerical_bins: int, + numerical_epsilon_ratio: float = 1.0, ) -> dict[str, api.MechanismConfig]: """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. + numerical_epsilon_ratio: Ratio by which privacy budget epsilon increases at + each deeper level of recursive bisection. Defaults to 1.0 (uniform budget + split across levels). Setting to sqrt(2) approx 1.414 can provide minor + accuracy gains on smooth continuous data. Returns: A dictionary mapping column names to uncalibrated initializer configs. @@ -57,6 +62,7 @@ def create_initializers( if isinstance(attr, domain.NumericalAttribute): initializers[col] = initialization.NumericalInitializerConfig( num_partitions=numerical_bins, + epsilon_ratio=numerical_epsilon_ratio, ) elif isinstance(attr, domain.CategoricalAttribute): initializers[col] = initialization.CategoricalInitializerConfig() @@ -378,6 +384,11 @@ class TabularConfig(api.MechanismConfig): domains: Mapping from column names to attribute domain specifications. discrete_mechanism: The mechanism to run on the discretized data. numerical_bins: Number of bins for numerical attribute discretization. + numerical_epsilon_ratio: Ratio by which privacy budget epsilon increases at + each deeper level of recursive bisection for numerical attribute + discretization. Defaults to 1.0 (uniform budget split across levels). A + value of sqrt(2) approx 1.414 can provide minor accuracy gains on smooth + continuous distributions. init_budget_fraction: Fraction of total zCDP budget allocated to per-column initialization (the rest goes to the discrete mechanism). cross_attribute_constraints: Constraints to enforce on generated data. @@ -392,6 +403,7 @@ class TabularConfig(api.MechanismConfig): domains: Mapping[str, domain.AttributeType] | None = None discrete_mechanism: api.MechanismConfig = discrete_mechanisms.MSTConfig() numerical_bins: int = 32 + numerical_epsilon_ratio: float = 1.0 init_budget_fraction: float = 0.1 cross_attribute_constraints: Sequence[constraints.Constraint] = () compress_columns: bool = False @@ -486,7 +498,9 @@ def configure( api.validate_max_records_per_user(max_records_per_user) per_col_deltas = self._compute_per_col_deltas(schema, delta) - inits = create_initializers(schema, self.numerical_bins) + inits = create_initializers( + schema, self.numerical_bins, self.numerical_epsilon_ratio + ) init_rho = self.init_budget_fraction * zcdp_rho # +1 for the DPGaussianCount that always measures the total. per_col_rho = init_rho / (len(inits) + 1) diff --git a/dpsynth/local_mode/initialization.py b/dpsynth/local_mode/initialization.py index 9af3c8b..08c1324 100644 --- a/dpsynth/local_mode/initialization.py +++ b/dpsynth/local_mode/initialization.py @@ -101,11 +101,20 @@ def compute_grid_spec( @dataclasses.dataclass(frozen=True, kw_only=True) class NumericalInitializerConfig(api.MechanismConfig): - """Configuration for initializing numerical attributes.""" + """Configuration for initializing numerical attributes. + + Attributes: + num_partitions: Number of partitions (must be a power of 2). + max_grid_size: Maximum grid size for the histogram. + epsilon_ratio: Ratio by which privacy budget epsilon increases at each + deeper level of recursive bisection. Defaults to 1.0 (uniform budget split + across levels). Setting to sqrt(2) approx 1.414 can provide minor accuracy + gains on smooth continuous data. + """ num_partitions: int max_grid_size: int = 10_000_000 - epsilon_ratio: float = 2.0 + epsilon_ratio: float = 1.0 def __post_init__(self): if self.max_grid_size < 2: diff --git a/tests/data_generation_v3_test.py b/tests/data_generation_v3_test.py index e1f2c5c..599d45b 100644 --- a/tests/data_generation_v3_test.py +++ b/tests/data_generation_v3_test.py @@ -348,6 +348,13 @@ def test_empty_dataset(self): # The true count is 0, but DPSynth always outputs at least one row. self.assertLen(result.synthetic_data, 1) + def test_numerical_epsilon_ratio_plumbing(self): + domains = {'A': domain.NumericalAttribute(min_value=0, max_value=10)} + config = TabularConfig(numerical_epsilon_ratio=1.414) + calibrated = config.configure(domains, zcdp_rho=10.0) + init_mech = calibrated.initializers['A'] + self.assertEqual(init_mech.config.epsilon_ratio, 1.414) + class MaxRecordsPerUserTest(parameterized.TestCase): """Tests the experimental user-level DP knob end to end.""" diff --git a/tests/local_mode/initialization_test.py b/tests/local_mode/initialization_test.py index 0147a3f..240517b 100644 --- a/tests/local_mode/initialization_test.py +++ b/tests/local_mode/initialization_test.py @@ -167,6 +167,34 @@ def test_numerical_initializer_no_measurement_without_estimated_total(self): result = initializer.configure(attr, zcdp_rho=1.0)(rng, data) self.assertIsNone(result.noisy_counts) + def test_numerical_initializer_epsilon_ratio_default_and_custom(self): + attr = domain.NumericalAttribute(min_value=0, max_value=10) + # Default is 1.0 (uniform) + init_default = initialization.NumericalInitializerConfig(num_partitions=4) + self.assertEqual(init_default.epsilon_ratio, 1.0) + calibrated_default = init_default.configure(attr, zcdp_rho=2.0) + # For 4 partitions (2 levels), uniform budget splits zcdp_rho equally: + # rho/2 each. eps = sqrt(8 * rho_level) = sqrt(8 * 1.0) = sqrt(8). + np.testing.assert_allclose( + calibrated_default.epsilon_levels, + (np.sqrt(8.0), np.sqrt(8.0)), + ) + + # Custom ratio (e.g. sqrt(2)) + init_custom = initialization.NumericalInitializerConfig( + num_partitions=4, epsilon_ratio=np.sqrt(2) + ) + self.assertEqual(init_custom.epsilon_ratio, np.sqrt(2)) + calibrated_custom = init_custom.configure(attr, zcdp_rho=2.0) + # rho_ratio = 2.0. budget_weights = [2.0, 1.0]. + # Leaves get 2/3, root gets 1/3. + # rho_levels = [2.0 * 2/3, 2.0 * 1/3] = [4/3, 2/3] + # eps_levels = [sqrt(8 * 4/3), sqrt(8 * 2/3)] + np.testing.assert_allclose( + calibrated_custom.epsilon_levels, + (np.sqrt(8.0 * 4.0 / 3.0), np.sqrt(8.0 * 2.0 / 3.0)), + ) + def test_integer_edges_at_max_value_absorbed_into_last_bin(self): """Edges at max_value are removed; their count goes to the last bin.""" attr = domain.NumericalAttribute(min_value=0, max_value=10, dtype='int')