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.

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

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(

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

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)