From ce8e6ecda7045ee247c8682143deea5ae9a301fd Mon Sep 17 00:00:00 2001 From: Ryan McKenna Date: Mon, 13 Jul 2026 12:54:27 -0700 Subject: [PATCH] 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. There are two separate ways to jitter depending on whether the attribute is int or float. For int, we expand each histogram bin into 4*m bins by randomly assigning each example to one of the equivalent bins (done in bulk via multinomial). For float, we move each example to one of the nearby bins symmetrically (done in bulk via multinomial). PiperOrigin-RevId: 947195276 --- dpsynth/local_mode/_quantiles.py | 107 ++++++++++++++--------- dpsynth/local_mode/initialization.py | 12 ++- dpsynth/local_mode/primitives.py | 16 ++-- tests/local_mode/_quantiles_test.py | 108 +++++++++++++++++++++--- tests/local_mode/initialization_test.py | 25 ++++-- 5 files changed, 204 insertions(+), 64 deletions(-) diff --git a/dpsynth/local_mode/_quantiles.py b/dpsynth/local_mode/_quantiles.py index a76c668..d754a14 100644 --- a/dpsynth/local_mode/_quantiles.py +++ b/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 diff --git a/dpsynth/local_mode/initialization.py b/dpsynth/local_mode/initialization.py index ea49c57..0cba735 100644 --- a/dpsynth/local_mode/initialization.py +++ b/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) diff --git a/dpsynth/local_mode/primitives.py b/dpsynth/local_mode/primitives.py index a5fd635..de44fde 100644 --- a/dpsynth/local_mode/primitives.py +++ b/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 diff --git a/tests/local_mode/_quantiles_test.py b/tests/local_mode/_quantiles_test.py index 13b08ca..e5dd7ad 100644 --- a/tests/local_mode/_quantiles_test.py +++ b/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__': diff --git a/tests/local_mode/initialization_test.py b/tests/local_mode/initialization_test.py index 5d070af..aa1d449 100644 --- a/tests/local_mode/initialization_test.py +++ b/tests/local_mode/initialization_test.py @@ -16,6 +16,7 @@ from absl.testing import parameterized import dp_accounting from dpsynth import domain +from dpsynth.local_mode import _quantiles from dpsynth.local_mode import initialization from dpsynth.local_mode import vectorized_transformations as vtx import numpy as np @@ -158,6 +159,19 @@ def test_max_grid_size_two_float(self): result = init(rng, np.linspace(0, 100, 100)) self.assertIsNotNone(result.categorical_attribute) + def test_int_grid_reserves_budget_for_jitter_refinement(self): + # A wide integer range with many partitions must not blow past + # max_grid_size once the m-fold jitter refinement is applied. + attr = domain.NumericalAttribute( + min_value=0, max_value=10_000_000, dtype='int' + ) + max_grid_size = 100_000 + init = initialization.NumericalInitializer( + name='x', num_partitions=64, attribute=attr, max_grid_size=max_grid_size + ) + m = _quantiles.jitter_factor(init.num_partitions) + self.assertLessEqual(init.grid_size * m, max_grid_size) + def test_numerical_initializer_measurement_with_estimated_total(self): attr = domain.NumericalAttribute(min_value=0, max_value=10) rng = np.random.default_rng(0) @@ -198,10 +212,12 @@ def test_integer_edges_at_max_value_absorbed_into_last_bin(self): initializer = initialization.NumericalInitializer( name='test', num_partitions=8, attribute=attr ) - # All data at max_value: all edges should land at or near max_value. - data = np.array([10] * 100) + # A spread of lower values carrying most of the mass, plus a moderate spike + # at max_value. The lower values form genuine interior bins, while the top + # quantile edges land at max_value and must be absorbed into the last bin. + data = np.concatenate([np.repeat(np.arange(0, 10), 20), np.full(100, 10)]) result = initializer.configure(zcdp_rho=100.0)( - rng, data, estimated_total=100.0 + rng, data, estimated_total=len(data) ) # No edge should equal max_value (they get absorbed). if len(result.bin_edges) > 0: @@ -210,8 +226,7 @@ def test_integer_edges_at_max_value_absorbed_into_last_bin(self): np.testing.assert_allclose( result.measurement.noisy_measurement.sum(), 1.0, atol=1e-10 ) - # The bin containing max_value=10 should get the most mass (either the - # last bin, or a bin that absorbed all degenerate edges). + # The last bin absorbs the max_value spike, so it should dominate. counts = result.measurement.noisy_measurement self.assertGreater(counts.max(), 1.0 / len(counts))