Fix biased DP quantiles utility bug introduced from recent quantiles refactor. The old quantiles path introduced jitter to handle elements with high mass (e.g., 40 hours per week on adult dataset). This change adds that back into the histogram-based quantiles codepaht, where the jitter occurs directly in the histogram domain rather than on a per-example level, preserving the nice properties of this implementation while hardening it to duplicate values.

dpsynth/local_mode/_quantiles.py

@@ -15,17 +15,32 @@
 """DP quantiles from dense histograms via recursive median bisection.
 
 This module computes differentially private quantile edges from a dense
-histogram of counts, using the discrete exponential mechanism.  The primary
-use case is a two-pass pipeline: a first pass computes a dense histogram
-over a fine-grained grid, then ``quantiles_from_histogram`` finds DP
-quantiles from that histogram without touching individual records.
-
-Public API:
-  - ``quantiles_from_histogram``: DP quantiles via recursive median splits.
+histogram of counts, using the discrete exponential mechanism.  It works purely
+in index space -- ``quantiles_from_histogram`` returns cell indices into the
+histogram, and the caller maps those indices to domain values.  The primary use
+case is a two-pass pipeline: a first pass computes a dense histogram over a
+fine-grained grid, then ``quantiles_from_histogram`` finds DP quantile indices
+from that histogram without touching individual records.
+
+Tie handling via jitter
+-----------------------
+Recursive median bisection needs each record assigned to one side of every
+split independently.  A "spike" of records tied on one grid cell breaks this: a
+whole-cell split sends all that mass to one side, biasing the quantiles and
+collapsing sub-ranges (dropping edges).  We fix this by breaking ties directly
+in the histogram domain rather than over the raw data values -- each cell's
+count is redistributed to nearby cells as ``Multinomial(count, kernel)`` (one
+draw per non-empty cell), which is distributionally identical to independently
+perturbing each record and so needs no extra privacy budget.  The ``refine``
+strategy uses a strictly-positive kernel over refined sub-cells (value-
+preserving); the ``symmetric`` strategy uses a symmetric kernel over neighboring
+grid cells.
 """
 
 from __future__ import annotations
 
+from typing import Literal
+
 import numpy as np
 import scipy.special
 
@@ -64,51 +79,67 @@ def _median_from_histogram(
   return int(rng.choice(total_points, p=probs))
 
 
+def jitter_factor(num_partitions):
+  """Returns a data-independent jitter resolution m from num_partitions."""
+  # m >= num_partitions keeps each jittered cell below one partition's mass;
+  # the 4x absorbs multinomial fluctuation.
+  return max(1, 4 * num_partitions)
+
+
 def quantiles_from_histogram(
     rng: np.random.Generator,
     counts: np.ndarray,
-    lower: float,
-    upper: float,
     epsilon_levels: np.ndarray,
-    grid_size: int = 10_000_000,
-) -> list[float]:
-  """Computes DP quantiles from a dense histogram via recursive median splits.
+    jitter_strategy: Literal['symmetric', 'refine'],
+) -> list[int]:
+  """DP quantile edge indices into ``counts`` via jittered median bisection.
 
-  Uses the discrete exponential mechanism to recursively find medians,
-  splitting the histogram at each level to produce ``num_buckets - 1``
-  quantile edges.  The number of buckets is ``2 ** len(epsilon_levels)``.
+  Operates purely in index space: it returns cell indices into ``counts`` and
+  leaves the mapping from index to domain value to the caller.
 
   Args:
     rng: A numpy random number generator.
-    counts: Dense 1D histogram of shape ``(grid_size,)``.
-    lower: Lower bound of the data domain.
-    upper: Upper bound of the data domain (exclusive).
+    counts: Dense 1D histogram counts.
     epsilon_levels: Per-level exponential mechanism epsilons, ordered from the
       deepest (finest) level to the shallowest (coarsest).
-    grid_size: Number of uniformly spaced grid points spanning ``[lower,
-      upper]``.
+    jitter_strategy: Specifies the pre-processing jitter strategy, -
+      'symmetric': jitter mass to +/- m//2 neighbors on the same grid. -
+      'refine': jitter mass to m equivalent sub-cells.
 
   Returns:
-    A sorted list of ``2 ** len(epsilon_levels) - 1`` quantile edge values.
+    A sorted list of ``2 ** len(epsilon_levels) - 1`` cell indices.
   """
-  levels = len(epsilon_levels)
-  if levels == 0:
-    return []
-
-  # Uniform grid: counts[i] corresponds to value lower + i * delta.
-  delta = (upper - lower) / (grid_size - 1)
+  counts = np.asarray(counts)
+  m = jitter_factor(2 ** len(epsilon_levels))
+
+  if jitter_strategy == 'refine':
+    stride, offsets = m, np.arange(m)
+  else:
+    half = m // 2
+    stride, offsets = 1, np.arange(-half, half + 1)
+
+  # Scatter each cell's mass over its jittered targets: same law as perturbing
+  # each record, so it breaks ties without spending extra privacy budget.
+  num_cells = counts.size * stride
+  nz = np.flatnonzero(counts)
+  probas = np.full(offsets.size, 1.0 / offsets.size)
+  split = rng.multinomial(counts[nz].astype(np.int64), probas)
+  targets = np.clip(nz[:, None] * stride + offsets, 0, num_cells - 1)
+  jittered = np.bincount(
+      targets.flatten(), weights=split.flatten(), minlength=num_cells
+  )
 
   def _rec(lo_idx, hi_idx, depth):
-    if depth == 0 or lo_idx >= hi_idx:
+    if depth == 0:
       return []
-    sub_counts = counts[lo_idx:hi_idx]
-    median_local = _median_from_histogram(
-        rng, sub_counts, epsilon_levels[depth - 1]
+    median_idx = lo_idx + _median_from_histogram(
+        rng, jittered[lo_idx:hi_idx], epsilon_levels[depth - 1]
     )
-    median_global_idx = lo_idx + median_local
-    median_value = lower + median_global_idx * delta
-    left = _rec(lo_idx, median_global_idx, depth - 1)
-    right = _rec(median_global_idx, hi_idx, depth - 1)
-    return left + [median_value] + right
-
-  return _rec(0, len(counts), levels)
+    left = _rec(lo_idx, median_idx, depth - 1)
+    right = _rec(median_idx, hi_idx, depth - 1)
+    return left + [median_idx] + right
+
+  result = _rec(0, jittered.size, len(epsilon_levels))
+  if jitter_strategy == 'refine':
+    result = [idx // m for idx in result]
+  return result

dpsynth/local_mode/initialization.py

@@ -22,6 +22,7 @@
 
 import dp_accounting
 from dpsynth import domain
+from dpsynth.local_mode import _quantiles
 from dpsynth.local_mode import primitives
 from dpsynth.local_mode import vectorized_transformations as vtx
 import mbi
@@ -95,8 +96,11 @@ def _grid_spec(self) -> tuple[float, float, int]:
     """Returns (lower, upper, grid_size) for the quantile candidate grid."""
     attr = self.attribute
     if attr.dtype == 'int':
+      # Reserve budget for the m-fold refinement so the refined grid fits.
+      m = _quantiles.jitter_factor(self.num_partitions)
+      budget = max(2, self.max_grid_size // m)
       int_range = int(attr.max_value - attr.min_value + 1)
-      step = max(1, math.ceil(int_range / self.max_grid_size))
+      step = max(1, math.ceil(int_range / budget))
       gs = math.ceil(int_range / step)
       return (attr.min_value, attr.min_value + (gs - 1) * step, gs)
     return (attr.min_value, attr.exclusive_max_value, self.max_grid_size)
@@ -110,12 +114,14 @@ def configure(
       self, *, zcdp_rho: float, delta: float = 0.0, epsilon_ratio: float = 2.0
   ) -> NumericalInitializer:
     """Returns a copy calibrated to the given zCDP budget."""
-    lower, upper, gs = self._grid_spec
+    lower, upper, _ = self._grid_spec
     mechanism = primitives.DPQuantiles(
         num_partitions=self.num_partitions,
         lower=lower,
         upper=upper,
-        grid_size=gs,
+        jitter_strategy=(
+            'refine' if self.attribute.dtype == 'int' else 'symmetric'
+        ),
     ).configure(zcdp_rho=zcdp_rho, epsilon_ratio=epsilon_ratio)
     return dataclasses.replace(self, mechanism=mechanism)
 

dpsynth/local_mode/primitives.py

@@ -22,6 +22,7 @@
 
 import dataclasses
 import math
+from typing import Literal
 
 import dp_accounting
 from dpsynth import api
@@ -274,13 +275,14 @@ class DPQuantiles(DPMechanism):
     num_partitions: Number of quantile partitions (must be a power of 2).
     lower: Lower bound of the data domain.
     upper: Upper bound of the data domain (exclusive).
-    grid_size: Number of uniformly spaced grid points.
+    jitter_strategy: Tie-breaking jitter passed to ``quantiles_from_histogram``:
+      ``'refine'`` for integer attributes, ``'symmetric'`` for continuous ones.
   """
 
   num_partitions: int
   lower: float
   upper: float
-  grid_size: int = 10_000_000
+  jitter_strategy: Literal['symmetric', 'refine'] = 'symmetric'
   _epsilon_levels: tuple[float, ...] | None = dataclasses.field(
       default=None, repr=False
   )
@@ -336,14 +338,16 @@ def __call__(
     """Returns quantile edges from a dense histogram of counts."""
     if self._epsilon_levels is None:
       raise ValueError(_UNCALIBRATED_MSG.format(param='_epsilon_levels'))
-    return _quantiles.quantiles_from_histogram(
+    indices = _quantiles.quantiles_from_histogram(
         rng,
         counts,
-        self.lower,
-        self.upper,
         epsilon_levels=np.asarray(self._epsilon_levels),
-        grid_size=self.grid_size,
+        jitter_strategy=self.jitter_strategy,
     )
+    # Map cell indices back to domain values; delta is the grid step, which
+    # equals the integer step for integer attributes so edges stay integer.
+    delta = (self.upper - self.lower) / max(1, np.asarray(counts).size - 1)
+    return [self.lower + i * delta for i in indices]
 
 
 @dataclasses.dataclass

tests/local_mode/_quantiles_test.py

@@ -23,29 +23,113 @@ class QuantilesFromHistogramTest(parameterized.TestCase):
   def test_no_levels_returns_empty(self):
     rng = np.random.default_rng(0)
     counts = np.array([10])
+    for jitter_strategy in ('symmetric', 'refine'):
+      edges = _quantiles.quantiles_from_histogram(
+          rng, counts, np.array([]), jitter_strategy
+      )
+      self.assertEmpty(edges)
+
+  @parameterized.product(
+      levels=(1, 2, 3, 4),
+      jitter_strategy=('symmetric', 'refine'),
+  )
+  def test_edge_count_matches_levels(self, levels, jitter_strategy):
+    rng = np.random.default_rng(0)
+    grid_size = 10001
+    counts = rng.integers(0, 20, size=grid_size)
     edges = _quantiles.quantiles_from_histogram(
         rng,
         counts,
-        0.0,
-        10.0,
-        epsilon_levels=np.array([]),
+        epsilon_levels=np.ones(levels),
+        jitter_strategy=jitter_strategy,
     )
-    self.assertEmpty(edges)
+    self.assertLen(edges, 2**levels - 1)
 
   @parameterized.parameters(1, 2, 3, 4)
-  def test_edge_count_matches_levels(self, levels):
+  def test_edge_count_matches_levels_with_spike(self, levels):
+    # A large tied spike must not collapse split ranges and drop edges. With
+    # whole-cell splits this dropped edges; jitter breaks up the spike so the
+    # recursion always emits the full 2**levels - 1 edges.
+    counts = np.zeros(101, dtype=np.int64)
+    counts[:40] = 1
+    counts[40] = 1000
+    counts[41:80] = 1
+    edges = _quantiles.quantiles_from_histogram(
+        np.random.default_rng(0),
+        counts,
+        epsilon_levels=np.array([np.inf] * levels),
+        jitter_strategy='refine',
+    )
+    self.assertLen(edges, 2**levels - 1)
+
+  def test_integer_edges_are_integer_indices(self):
+    # Integer attributes must return integer cell indices into ``counts``.
+    counts = np.zeros(101, dtype=np.int64)
+    counts[40] = 5000
+    counts[:40] = 50
+    counts[41:] = 50
+    edges = _quantiles.quantiles_from_histogram(
+        np.random.default_rng(0),
+        counts,
+        epsilon_levels=np.array([np.inf] * 3),
+        jitter_strategy='refine',
+    )
+    for edge in edges:
+      self.assertEqual(edge, int(edge))
+      self.assertBetween(edge, 0, counts.size - 1)
+
+  def test_exact_budget_matches_numpy_smooth(self):
     rng = np.random.default_rng(0)
-    grid_size = 10001
-    counts = rng.integers(0, 20, size=grid_size)
+    data = rng.integers(0, 100, size=50000)
+    counts = np.bincount(data, minlength=101)
     edges = _quantiles.quantiles_from_histogram(
         rng,
         counts,
-        0.0,
-        10.0,
-        epsilon_levels=np.ones(levels),
-        grid_size=grid_size,
+        epsilon_levels=np.array([np.inf] * 3),
+        jitter_strategy='refine',
     )
-    self.assertLen(edges, 2**levels - 1)
+    # Cell indices map 1:1 to values here (delta == 1), so compare directly.
+    expected = np.quantile(data, [0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875])
+    np.testing.assert_allclose(edges, expected, atol=1.0)
+
+  def test_exact_budget_matches_numpy_with_spike(self):
+    # Reproduces the failure mode from the 'hours-per-week' column of the adult
+    # census dataset: a dominant spike at 40 carrying ~45% of the mass, with
+    # lighter integer support on either side. The pre-jitter whole-cell split
+    # lumped all tied mass into one subtree, biasing the low quantiles (e.g. the
+    # 0.25 edge came out as 18 instead of 40) and dropping upper edges. Jitter
+    # breaks up the spike so the recursive medians match numpy.
+    below = np.arange(1, 40).repeat(230)
+    spike = np.full(13500, 40)
+    above = np.arange(41, 80).repeat(190)
+    data = np.concatenate([below, spike, above])
+    counts = np.bincount(data, minlength=101)
+    edges = _quantiles.quantiles_from_histogram(
+        np.random.default_rng(0),
+        counts,
+        epsilon_levels=np.array([np.inf] * 3),
+        jitter_strategy='refine',
+    )
+    # Cell indices map 1:1 to values here (delta == 1), so compare directly.
+    expected = np.quantile(data, [0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875])
+    np.testing.assert_allclose(edges, expected, atol=1.0)
+
+  def test_spike_owns_consecutive_edges(self):
+    # When a single value holds a majority of the mass, the quantiles on both
+    # sides of the median should collapse onto that value. This is the key
+    # correctness property the deterministic whole-cell split got wrong.
+    counts = np.zeros(101, dtype=np.int64)
+    counts[:40] = 20
+    counts[40] = 20000  # ~96% of the mass.
+    counts[41:80] = 20
+    edges = _quantiles.quantiles_from_histogram(
+        np.random.default_rng(0),
+        counts,
+        epsilon_levels=np.array([np.inf] * 3),
+        jitter_strategy='refine',
+    )
+    # The interior quantiles (0.25 through 0.75) must all be cell 40.
+    self.assertEqual(edges[1:6], [40, 40, 40, 40, 40])
 
 
 if __name__ == '__main__':