diff --git a/bin/main.py b/bin/main.py index 9d98ae2..6631513 100644 --- a/bin/main.py +++ b/bin/main.py @@ -32,7 +32,6 @@ import numpy as np import pandas as pd - _DATASET_PATH = flags.DEFINE_string( 'dataset', None, @@ -63,7 +62,7 @@ _MECHANISM = flags.DEFINE_enum( 'mechanism', 'mst', - ['mst', 'aim', 'independent', 'aim_gdp'], + ['mst', 'aim', 'independent'], 'Mechanism to use.', ) @@ -100,8 +99,6 @@ def main(_): mechanism_config = dpsynth.discrete_mechanisms.AIMConfig() case 'independent': mechanism_config = dpsynth.discrete_mechanisms.IndependentConfig() - case 'aim_gdp': - mechanism_config = dpsynth.discrete_mechanisms.AIMGDPConfig() case _: raise ValueError(f'Unknown mechanism: {_MECHANISM.value}') diff --git a/docs/api_reference.rst b/docs/api_reference.rst index 412ffc9..751f23d 100644 --- a/docs/api_reference.rst +++ b/docs/api_reference.rst @@ -133,7 +133,6 @@ marginal measurement and domain compression. IndependentConfig DirectConfig SWIFTConfig - AIMGDPConfig DiscreteConfig and DiscreteMechanism ------------------------------------ diff --git a/docs/in_memory_api.md b/docs/in_memory_api.md index 5a16d3e..0bb9e6a 100644 --- a/docs/in_memory_api.md +++ b/docs/in_memory_api.md @@ -162,8 +162,7 @@ python3 bin/main.py \ arguments via `--read_csv_args`). * `--domain`: Path to the YAML domain specification file. * `--epsilon`, `--delta`: Total DP privacy budget. -* `--mechanism`: Supported options are `mst`, `aim`, `independent`, and - `aim_gdp`. +* `--mechanism`: Supported options are `mst`, `aim`, and `independent`. * `--seed`: Integer seed for reproducible randomness across DP sampling and PGM inference. * `--output_path`: Destination filepath where the synthetic CSV will be diff --git a/dpsynth/discrete_mechanisms/README.md b/dpsynth/discrete_mechanisms/README.md index 061a695..c9380bc 100644 --- a/dpsynth/discrete_mechanisms/README.md +++ b/dpsynth/discrete_mechanisms/README.md @@ -88,18 +88,6 @@ workload; `_allocate_budget()` reserves rho for the adaptive loop; `_run()` replaces the standard base execution path. Helper functions filter valid candidates and privately choose the worst-approximated marginal. -## `aim_gdp.py` — AIM with GDP-Oriented Allocation - -Implements `AIMGDPConfig`, a variant of AIM with the same adaptive workflow -but GDP units for its internal loop budgeting. It is useful when its alternative -privacy-accounting behavior is preferred. - -**Public API:** `AIMGDPConfig(workload=...)` - -**Internal behavior:** Like `aim.py`, it overrides `_one_way_cliques()`, -`_allocate_budget()`, and `_run()`. Its internal helpers compute GDP-aware error -scores and select the next workload marginal. - ## `swift.py` — Workload and Clique-Tree Mechanism Implements `SWIFT`, a workload-informed mechanism that selects diff --git a/dpsynth/discrete_mechanisms/__init__.py b/dpsynth/discrete_mechanisms/__init__.py index 7de542d..42633a9 100644 --- a/dpsynth/discrete_mechanisms/__init__.py +++ b/dpsynth/discrete_mechanisms/__init__.py @@ -26,8 +26,6 @@ from dpsynth.api import MechanismConfig from dpsynth.discrete_mechanisms.aim import AIM from dpsynth.discrete_mechanisms.aim import AIMConfig -from dpsynth.discrete_mechanisms.aim_gdp import AIMGDP -from dpsynth.discrete_mechanisms.aim_gdp import AIMGDPConfig from dpsynth.discrete_mechanisms.common import DiscreteMechanismResult from dpsynth.discrete_mechanisms.common import MechanismDiagnostics from dpsynth.discrete_mechanisms.direct import Direct @@ -43,7 +41,6 @@ # Backwards-compatible aliases. AIMMechanism = AIMConfig -AIMGDPMechanism = AIMGDPConfig DirectMechanism = DirectConfig IndependentMechanism = IndependentConfig MSTMechanism = MSTConfig diff --git a/dpsynth/discrete_mechanisms/aim_gdp.py b/dpsynth/discrete_mechanisms/aim_gdp.py deleted file mode 100644 index 44bc6f7..0000000 --- a/dpsynth/discrete_mechanisms/aim_gdp.py +++ /dev/null @@ -1,368 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Variant of the Adaptive+Iterative Mechanism (AIM) that satisfies Gaussian DP.""" - -from collections.abc import Iterable, Mapping -from collections.abc import Sequence -import dataclasses -import typing - -from absl import logging -import dp_accounting -from dpsynth import api -from dpsynth.discrete_mechanisms import accounting -from dpsynth.discrete_mechanisms import common -import jax.numpy as jnp -import mbi -import mbi.junction_tree -import numpy as np - -MarginalQuery: typing.TypeAlias = tuple[str, ...] - - -def _filter_candidates( - candidates: Mapping[mbi.Clique, float], - model: mbi.MarkovRandomField, - size_limit: float, -) -> Mapping[mbi.Clique, float]: - """Filters the given candidates that lead to tractable graphical models. - - Args: - candidates: The candidate marginal queries. - model: The current graphical model. - size_limit: The size limit in megabytes for the new graphical model, if a - given candidate is selected. - - Returns: - A collection of new candidates that pass the size_limit filter. - """ - - def expected_size(cl): - return mbi.junction_tree.hypothetical_model_size( - model.domain, [*model.cliques, cl] - ) - - ans = {} - free_cliques = common.downward_closure(model.cliques) - for cl in candidates: - if expected_size(cl) <= size_limit or cl in free_cliques: - ans[cl] = candidates[cl] - return ans - - -def _compute_dp_errors( - rng: np.random.Generator, - data: mbi.Dataset | mbi.CliqueVector, - estimates: mbi.CliqueVector, - gdp_budget: float, - subset: Iterable[mbi.Clique], - max_records_per_user: int = 1, -) -> dict[mbi.Clique, float]: - """Compute L1 error between the model answers and the true answers with DP.""" - clique_list = list(subset) - # The L1 error of a marginal changes by at most ``max_records_per_user`` when - # a single user (contributing up to that many records) is added or removed. - per_candidate_sigma = max_records_per_user * accounting.gdp_gaussian_sigma( - gdp_budget / len(clique_list) # pyrefly: ignore[bad-argument-type] - ) - result = {} - for cl in clique_list: - actual = data.project(cl).datavector(flatten=True) - estimate = estimates[cl].datavector(flatten=True) - error = jnp.linalg.norm(actual - estimate, ord=1) - noise = rng.normal(loc=0, scale=per_candidate_sigma) - result[cl] = error + noise - return result - - -def _worst_approximated( - rng: np.random.Generator, - candidates: Mapping[mbi.Clique, float], - errors: dict[mbi.Clique, float], # will be updated in-place. - data: mbi.Dataset | mbi.CliqueVector, # sensitive data. - model: mbi.MarkovRandomField, - select_budget: float, # satisfies select_budget-GDP. - measure_sigma: float, - max_new_evals: int, - max_records_per_user: int = 1, -) -> mbi.Clique: - """Returns the worst approximated candidate in the given candidates.""" - current_score_estimates = {} - for cl in candidates: - weight = candidates[cl] - bias = ( - (2 / np.pi) ** 0.5 - * max_records_per_user - * measure_sigma - * model.domain.size(cl) - ) - current_score_estimates[cl] = weight * (errors[cl] - bias) - - subset = sorted(current_score_estimates, key=current_score_estimates.get) # pyrefly: ignore[no-matching-overload] - subset = subset[-max_new_evals:] - - estimates = mbi.marginal_oracles.bulk_variable_elimination( - model.potentials, subset, model.total # pyrefly: ignore[bad-argument-type] - ) - # Only step that uses "data", satisfies DP. - current_errors = _compute_dp_errors( - rng, - data, - estimates, - select_budget, - subset, - max_records_per_user=max_records_per_user, - ) - errors.update(current_errors) - - current_scores = {} - for cl in subset: - weight = candidates[cl] - bias = ( - (2 / np.pi) ** 0.5 - * max_records_per_user - * measure_sigma - * model.domain.size(cl) - ) - current_scores[cl] = weight * (errors[cl] - bias) - - return max(current_scores, key=current_scores.get) # pyrefly: ignore[no-matching-overload] - - -# select loop, injecting the budgeting strategy (zCDP vs. GDP) as configuration. -@dataclasses.dataclass(frozen=True) -class AIMGDPConfig(api.MechanismConfig): - """Configuration for the AIM mechanism with Gaussian DP. - - Details are described in the paper: - [AIM: An Adaptive and Iterative Mechanism for Differentially Private Synthetic - Data](https://arxiv.org/abs/2201.12677). This mechanism is a competitive - algorithm within the broader SELECT-MEASURE-GENERATE paradigm. It is an - MWEM-style algorithm (Multiplicative Weights + Exponential Mechanism), that - iteratively improves the estimate of the data distribution by selecting - marginal queries that are poorly approximated by the current model. It is a - scalable algorithm that can handle high-dimensional datasets, but it can be - time consuming to run (hours). The runtime/utility trade-off can be controlled - by the max_model_size parameter. For quick experimentation, we recommend - setting max_model_size = 1, for production use cases, we recommend setting - max_model_size >= 80. - - Attributes: - workload: A collection of marginal queries (and weights) the synthetic data - should be tailored to. The weights determine the relative importance of - each marginal query. A default value of 1.0 will be assigned if the - workload is provided as a list. - max_rounds: The maximum number of rounds to run the mechanism. - max_model_size: The maximum size of the graphical model in megabytes. - Controls the utility/runtime trade-off. - max_marginal_size: The maximum size of a marginal query to consider. - max_candidates_per_round: The maximum number of candidates to consider per - round. This can improve privacy budget utilization as well as speed up the - "Select" step, which in some settings is the main bottelneck of the - mechanism. - anneal_factor: The factor by which to anneal the privacy budget. - select_budget_fraction: The fraction of the privacy budget to use for the - "Select" step. - """ - - workload: Mapping[mbi.Clique, float] | Iterable[mbi.Clique] | None = None - max_rounds: int | None = None - max_model_size: int = 80 - max_marginal_size: float = 1e6 - max_candidates_per_round: int = 16 - anneal_factor: float = 4.0 - select_budget_fraction: float = 0.1 - pgm_iters: int = 1000 - marginal_oracle: mbi.MarginalOracle | None = None - - def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: - """Returns the workload cliques filtered by max_marginal_size.""" - return common.supporting_cliques( - domain, self.workload, self.max_marginal_size - ) - - def configure(self, _=None, *, zcdp_rho, delta=0, max_records_per_user=1): - api.validate_max_records_per_user(max_records_per_user) - return AIMGDP( - config=self, - gdp_budget=accounting.zcdp_to_gdp(zcdp_rho), - max_records_per_user=max_records_per_user, - ) - - -@dataclasses.dataclass(frozen=True, kw_only=True) -class AIMGDP(api.CalibratedMechanism): - """Calibrated AIMGDP instance.""" - - config: AIMGDPConfig - gdp_budget: float - max_records_per_user: int = 1 - - @property - def dp_event(self) -> dp_accounting.DpEvent: - """Returns the DP event for the AIM-GDP mechanism.""" - return dp_accounting.GaussianDpEvent( - accounting.gdp_gaussian_sigma(self.gdp_budget) - ) - - def __call__( - self, - rng: np.random.Generator, - data: mbi.Dataset | mbi.CliqueVector, - *, - initial_measurements: Sequence[mbi.LinearMeasurement] | None = None, - constraints: Sequence[mbi.Constraint] = (), - ) -> common.DiscreteMechanismResult: - common.validate_initial_measurements(initial_measurements) - measurements = list(initial_measurements) if initial_measurements else [] - phase_times = {} - logging.info('[AIM] Starting Mechanism.') - - gdp_budget = self.gdp_budget - - terminate = False - budget_remaining = gdp_budget - max_rounds = self.config.max_rounds or 16 * len(data.domain) - budget_per_round = budget_remaining / max_rounds - - ######################################################################### - # Compile workload into candidate measurements. # - ######################################################################### - candidates = common.compiled_workload( - data.domain, self.config.workload, self.config.max_marginal_size - ) - domain = data.domain - - estimator = mbi.estimation.MirrorDescent(self.config.marginal_oracle) - model = estimator.estimate( - domain, - measurements, - iters=self.config.pgm_iters, - constraints=constraints, - ) - assert isinstance(model, mbi.MarkovRandomField) - logging.info('[AIM] Estimated initial model.') - - # The initial model is fitted from 1-way measurements only, so it IS the - # independence model. compute_independence_errors is much faster than - # bulk_variable_elimination for this case (pure numpy, no XLA compilation). - budget_remaining -= 0.5 * budget_per_round - per_candidate_sigma = int( - self.max_records_per_user - ) * accounting.gdp_gaussian_sigma(0.5 * budget_per_round / len(candidates)) - errors = common.compute_independence_errors(data, model, list(candidates)) # pyrefly: ignore[bad-argument-type] - for cl in errors: - errors[cl] += rng.normal(loc=0.0, scale=per_candidate_sigma) - logging.info('[AIM] Computed initial errors.') - - t = 0 - while not terminate: - t += 1 - if budget_remaining < 2 * budget_per_round: - logging.info('[AIM] Final round, Using all remaining privacy budget.') - budget_per_round = budget_remaining - terminate = True - - ######################################################################## - # Select a marginal query worst approximated by the current model. # - ######################################################################## - with common.timed(phase_times, 'selection'): - budget_remaining -= budget_per_round - measure_budget = budget_per_round * ( - 1 - self.config.select_budget_fraction - ) - select_budget = budget_per_round * self.config.select_budget_fraction - measure_sigma = accounting.gdp_gaussian_sigma(measure_budget) - percent_used = (gdp_budget - budget_remaining) / gdp_budget - size_limit = self.config.max_model_size * percent_used - small_candidates = _filter_candidates(candidates, model, size_limit) - - marginal_query = _worst_approximated( - rng, - candidates=small_candidates, - errors=errors, - data=data, - model=model, - select_budget=select_budget, - measure_sigma=measure_sigma, - max_new_evals=self.config.max_candidates_per_round, - max_records_per_user=self.max_records_per_user, - ) - - summary = mbi.summarize( - domain, [m.clique for m in measurements] + [marginal_query] - ) - logging.info( - '[AIM-GDP] Round %d, Budget used: %.4f, Measuring: %s,' - ' Candidates: %d, cliques: %d, treewidth: %d, memory: %d bytes', - t, - percent_used, - marginal_query, - len(small_candidates), - summary.num_cliques, - summary.treewidth, - summary.memory_bytes, - ) - - ###################################################################### - # Measure the marginal query privately using the Gaussian mechanism. # - ###################################################################### - with common.timed(phase_times, 'measurement'): - measurement = common.measure_marginals_with_noise( - rng, data, [marginal_query], measure_sigma # pyrefly: ignore[bad-argument-type] - )[0] - measurements.append(measurement) - old_estimate = model.project(marginal_query).datavector() - - ##################################################### - # Estimate the data distribution using Private-PGM. # - ##################################################### - with common.timed(phase_times, 'estimation'): - callback_fn = mbi.callbacks.default(measurements, domain) - model = estimator.estimate( - domain, - measurements, - warm_start=model, - iters=self.config.pgm_iters, - callback_fn=callback_fn, - constraints=constraints, - ) - model = typing.cast(mbi.MarkovRandomField, model) - - new_estimate = model.project(marginal_query).datavector() - - ########################################## - # Anneal epsilon and sigma if necessary. # - ########################################## - # See Alg 4 of https://arxiv.org/pdf/2201.12677. - # of just the largest error candidate), we can maybe simplify this logic. - threshold = ( - self.max_records_per_user - * measure_sigma - * (2 / np.pi) ** 0.5 - * domain.size(marginal_query) - ) - if np.linalg.norm(new_estimate - old_estimate, ord=1) <= threshold: - # No useful information at this noise level, increase budget per round. - budget_per_round *= self.config.anneal_factor - logging.info( - '[AIM] Increasing budget per round: %.5f', budget_per_round - ) - - return common.DiscreteMechanismResult( - measurements=measurements, - model=model, - diagnostics=common.clique_stats(model), - ) diff --git a/tests/data_generation_v3_test.py b/tests/data_generation_v3_test.py index 1eb09c5..758d637 100644 --- a/tests/data_generation_v3_test.py +++ b/tests/data_generation_v3_test.py @@ -24,7 +24,6 @@ from dpsynth import discrete_mechanisms from dpsynth import domain from dpsynth.discrete_mechanisms import aim -from dpsynth.discrete_mechanisms import aim_gdp from dpsynth.discrete_mechanisms.independent import IndependentConfig import mbi import numpy as np @@ -318,19 +317,6 @@ def test_discrete_workload_regression_with_aim(self): ) self.assertLess(mechanism_error, 0.05 * baseline_error) - def test_discrete_workload_regression_with_aim_gdp(self): - workload = [('a',), ('b',), ('c',), ('a', 'b'), ('a', 'c'), ('b', 'c')] - config = aim_gdp.AIMGDPConfig( - workload=workload, max_rounds=4, pgm_iters=500 - ) - baseline_config = IndependentConfig(pgm_iters=500) - mechanism_error, baseline_error = ( - _discrete_workload_mechanism_baseline_errors( - config, baseline_config, workload - ) - ) - self.assertLess(mechanism_error, 0.05 * baseline_error) - def test_mixed_workload_regression_with_aim(self): workload = [('a',), ('b',), ('c',), ('a', 'b'), ('a', 'c'), ('b', 'c')] config = aim.AIMConfig(workload=workload, max_rounds=4, pgm_iters=500) @@ -340,17 +326,6 @@ def test_mixed_workload_regression_with_aim(self): ) self.assertLess(mechanism_error, 0.05 * baseline_error) - def test_mixed_workload_regression_with_aim_gdp(self): - workload = [('a',), ('b',), ('c',), ('a', 'b'), ('a', 'c'), ('b', 'c')] - config = aim_gdp.AIMGDPConfig( - workload=workload, max_rounds=4, pgm_iters=500 - ) - baseline_config = IndependentConfig(pgm_iters=500) - mechanism_error, baseline_error = _mixed_workload_mechanism_baseline_errors( - config, baseline_config, workload - ) - self.assertLess(mechanism_error, 0.05 * baseline_error) - def test_empty_dataset(self): """Tests that DPSynth works without crashing on empty datasets, and outputs noisy rows.""" domains = { diff --git a/tests/discrete_mechanisms/aim_test.py b/tests/discrete_mechanisms/aim_test.py index 7370d85..70f353e 100644 --- a/tests/discrete_mechanisms/aim_test.py +++ b/tests/discrete_mechanisms/aim_test.py @@ -14,9 +14,7 @@ from absl.testing import absltest from dpsynth.discrete_mechanisms import aim -from dpsynth.discrete_mechanisms import aim_gdp from dpsynth.discrete_mechanisms import common -from dpsynth.discrete_mechanisms import independent import mbi import numpy as np @@ -73,54 +71,10 @@ def test_fits_one_way_marginals_with_aim(self): actual = result.model.project([col]).datavector() np.testing.assert_allclose(actual, expected, atol=1) - def test_fits_one_way_marginals_with_aim_gdp(self): - data = mbi.Dataset.synthetic(mbi.Domain(["a", "b", "c"], [3, 4, 5]), N=1000) - workload = [("a",), ("b",), ("c",)] - - config = aim_gdp.AIMGDPConfig( - workload=workload, max_rounds=4, pgm_iters=500 - ) - calibrated = config.configure(zcdp_rho=10000) - result = calibrated(np.random.default_rng(0), data) - - self.assertIsInstance(result, common.DiscreteMechanismResult) - self.assertNotEmpty(result.measurements) - for col in data.domain: - expected = data.project([col]).datavector() - actual = result.model.project([col]).datavector() - np.testing.assert_allclose(actual, expected, atol=1) - - def test_correlated_workload_regression_with_aim(self): - workload = [("a",), ("b",), ("c",), ("a", "b"), ("a", "c"), ("b", "c")] - config = aim.AIMConfig(workload=workload, max_rounds=4, pgm_iters=500) - baseline_config = independent.IndependentConfig() - mechanism_error, baseline_error = ( - _correlated_workload_mechanism_baseline_errors( - config, baseline_config, workload - ) - ) - self.assertLess(mechanism_error, 0.05 * baseline_error) - - def test_correlated_workload_regression_with_aim_gdp(self): - workload = [("a",), ("b",), ("c",), ("a", "b"), ("a", "c"), ("b", "c")] - config = aim_gdp.AIMGDPConfig( - workload=workload, max_rounds=4, pgm_iters=500 - ) - baseline_config = independent.IndependentConfig() - mechanism_error, baseline_error = ( - _correlated_workload_mechanism_baseline_errors( - config, baseline_config, workload - ) - ) - self.assertLess(mechanism_error, 0.05 * baseline_error) - def test_default_configuration_values(self): config = aim.AIMConfig() self.assertEqual(config.pgm_iters, 1000) - gdp_config = aim_gdp.AIMGDPConfig() - self.assertEqual(gdp_config.pgm_iters, 1000) - if __name__ == "__main__": absltest.main() diff --git a/tests/discrete_mechanisms/discrete_mechanisms_test.py b/tests/discrete_mechanisms/discrete_mechanisms_test.py index 3eeffd4..06f624c 100644 --- a/tests/discrete_mechanisms/discrete_mechanisms_test.py +++ b/tests/discrete_mechanisms/discrete_mechanisms_test.py @@ -23,7 +23,6 @@ from absl.testing import absltest from absl.testing import parameterized from dpsynth.discrete_mechanisms import aim -from dpsynth.discrete_mechanisms import aim_gdp from dpsynth.discrete_mechanisms import common from dpsynth.discrete_mechanisms import direct from dpsynth.discrete_mechanisms import discrete @@ -38,9 +37,6 @@ _MECHANISMS = { 'AIM': aim.AIMConfig(workload=_WORKLOAD, max_rounds=4, pgm_iters=500), - 'AIM_GDP': aim_gdp.AIMGDPConfig( - workload=_WORKLOAD, max_rounds=4, pgm_iters=500 - ), 'MST': mst.MSTConfig(pgm_iters=500), 'SWIFT': swift.SWIFTConfig(workload=_WORKLOAD, pgm_iters=500), 'Independent': independent.IndependentConfig(), diff --git a/tests/serialize_test.py b/tests/serialize_test.py index f23ec08..b6c0ae2 100644 --- a/tests/serialize_test.py +++ b/tests/serialize_test.py @@ -24,7 +24,6 @@ from dpsynth import relational from dpsynth import serialize from dpsynth.discrete_mechanisms import aim -from dpsynth.discrete_mechanisms import aim_gdp from dpsynth.discrete_mechanisms import direct from dpsynth.discrete_mechanisms import discrete from dpsynth.discrete_mechanisms import independent @@ -92,15 +91,6 @@ def test_direct_config_roundtrip(self): loaded = serialize.from_yaml(yaml_str) self.assertEqual(loaded, config) - def test_aim_gdp_config_roundtrip(self): - config = aim_gdp.AIMGDPConfig( - pgm_iters=500, - max_rounds=10, - ) - yaml_str = serialize.to_yaml(config) - loaded = serialize.from_yaml(yaml_str) - self.assertEqual(loaded, config) - def test_discrete_config_roundtrip(self): config = discrete.DiscreteConfig( mechanism=aim.AIMConfig(pgm_iters=400),