Expose numerical_epsilon_ratio on TabularConfig and update default to 1.0.

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)

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:

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."""

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')