From 9bd8f5b819b6bd4cc2b1b56f1cf73b6199a8bc83 Mon Sep 17 00:00:00 2001 From: Ryan McKenna Date: Fri, 9 Oct 2026 09:23:59 -0700 Subject: [PATCH] Add `estimator` configuration field to `DirectConfig`. The default MirrorDescent estimator is pretty robust but only has a O(1/t) convergence rate. mbi supports other estimators that enjoy improved O(1/t^2) convergence rates, which can be particularly helpful in regimes with high signal to noise ratio where getting close to the optimum is important. PiperOrigin-RevId: 996576770 --- dpsynth/discrete_mechanisms/direct.py | 5 +++- dpsynth/serialize.py | 13 +++++++++ tests/discrete_mechanisms/direct_test.py | 34 ++++++++++++++++++++++++ tests/serialize_test.py | 18 +++++++++++++ 4 files changed, 69 insertions(+), 1 deletion(-) diff --git a/dpsynth/discrete_mechanisms/direct.py b/dpsynth/discrete_mechanisms/direct.py index 78846dc1..407a650f 100644 --- a/dpsynth/discrete_mechanisms/direct.py +++ b/dpsynth/discrete_mechanisms/direct.py @@ -35,6 +35,7 @@ def configure(self, _=None, *, budget, delta=0): gdp_budget=accounting.zcdp_to_gdp(budget), ) + estimator: mbi.Estimator = mbi.estimation.MirrorDescent() marginal_oracle: mbi.MarginalOracle | None = None pgm_iters: int = 5000 prespecified_marginal_queries: list[tuple[str, ...]] = dataclasses.field( @@ -79,7 +80,9 @@ def __call__( logging.info('[%s]:\n%s', type(self).__name__, summary) # Kick off async AOT compilation of the estimator while we measure. - estimator = mbi.estimation.MirrorDescent(self.config.marginal_oracle) + estimator = self.config.estimator + oracle = estimator.marginal_oracle or self.config.marginal_oracle + estimator = dataclasses.replace(estimator, marginal_oracle=oracle) pgm_future = estimator.precompile( data.domain, list(initial_measurements), diff --git a/dpsynth/serialize.py b/dpsynth/serialize.py index 47196b2f..8e3f59ee 100644 --- a/dpsynth/serialize.py +++ b/dpsynth/serialize.py @@ -27,6 +27,7 @@ from dpsynth import reporting from dpsynth.relational import domain as relational_domain from etils import epath +import mbi import yaml PathType = epath.PathLike @@ -49,6 +50,10 @@ def _resolve_type(type_name: str) -> type[Any] | None: candidate, dp_accounting.DpEvent ): return candidate + if hasattr(mbi.estimation, type_name): + candidate = getattr(mbi.estimation, type_name) + if isinstance(candidate, type) and issubclass(candidate, mbi.Estimator): + return candidate return None @@ -90,6 +95,10 @@ def _make_converter() -> cattrs.Converter: lambda cl: isinstance(cl, type) and issubclass(cl, dp_accounting.DpEvent), lambda cl: _unstructure_dataclass(cl, conv), ) + conv.register_unstructure_hook( + mbi.Estimator, + lambda obj: _unstructure_dataclass(obj.__class__, conv)(obj), + ) conv.register_unstructure_hook( tuple, lambda val: [conv.unstructure(x) for x in val], @@ -107,6 +116,10 @@ def _make_converter() -> cattrs.Converter: dp_accounting.DpEvent, lambda data, _: _structure_polymorphic(data, _, conv), ) + conv.register_structure_hook( + mbi.Estimator, + lambda data, _: _structure_polymorphic(data, _, conv), + ) # 3. Structure Sequence[T] consistently as list conv.register_structure_hook_func( diff --git a/tests/discrete_mechanisms/direct_test.py b/tests/discrete_mechanisms/direct_test.py index 0d1e05f2..9eace302 100644 --- a/tests/discrete_mechanisms/direct_test.py +++ b/tests/discrete_mechanisms/direct_test.py @@ -42,6 +42,40 @@ def test_calibrate_gives_expected_gdp_budget(self): calibrated = direct.DirectConfig().calibrate(epsilon=1.0, delta=1e-5) self.assertAlmostEqual(calibrated.gdp_budget, 0.07185134, places=6) + def test_custom_estimator(self): + data = mbi.Dataset.synthetic(mbi.Domain(['a', 'b', 'c'], [3, 4, 5]), N=1000) + prespecified_queries = [('a', 'b'), ('a', 'c'), ('b', 'c')] + config = direct.DirectConfig( + prespecified_marginal_queries=prespecified_queries, + estimator=mbi.estimation.InteriorGradient(), + pgm_iters=500, + ) + result = config.configure(budget=10000)(np.random.default_rng(0), data) + + 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_marginal_oracle_propagated_to_estimator(self): + data = mbi.Dataset.synthetic(mbi.Domain(['a', 'b'], [2, 3]), N=100) + calls = [] + + def tracking_oracle(potentials, total=1, constraints=()): + calls.append(total) + return mbi.marginal_oracles.brute_force_marginals( + potentials, total, constraints=constraints + ) + + config = direct.DirectConfig( + prespecified_marginal_queries=[('a', 'b')], + estimator=mbi.estimation.InteriorGradient(), + marginal_oracle=tracking_oracle, + pgm_iters=10, + ) + config.configure(budget=100)(np.random.default_rng(0), data) + self.assertNotEmpty(calls) + if __name__ == '__main__': absltest.main() diff --git a/tests/serialize_test.py b/tests/serialize_test.py index b6c0ae2a..53246a09 100644 --- a/tests/serialize_test.py +++ b/tests/serialize_test.py @@ -30,6 +30,7 @@ from dpsynth.discrete_mechanisms import mst from dpsynth.discrete_mechanisms import swift from dpsynth.local_mode import initialization +import mbi import yaml @@ -88,6 +89,23 @@ def test_direct_config_roundtrip(self): prespecified_marginal_queries=[('a', 'b'), ('c',)], ) yaml_str = serialize.to_yaml(config) + raw_dict = yaml.safe_load(yaml_str) + self.assertNotIn('estimator', raw_dict) + loaded = serialize.from_yaml(yaml_str) + self.assertEqual(loaded, config) + + def test_direct_config_with_estimator_roundtrip(self): + config = direct.DirectConfig( + pgm_iters=1500, + prespecified_marginal_queries=[('a', 'b')], + estimator=mbi.estimation.UniversalAcceleratedMethod(linesearch=True), + ) + yaml_str = serialize.to_yaml(config) + raw_dict = yaml.safe_load(yaml_str) + self.assertEqual( + raw_dict['estimator'], + {'type': 'UniversalAcceleratedMethod', 'linesearch': True}, + ) loaded = serialize.from_yaml(yaml_str) self.assertEqual(loaded, config)