mbi.approximate_oracles

Approximate marginal oracles with convex counting numbers.

See the paper [“Relaxed Marginal Consistency for Differentially Private Query Answering”](https://arxiv.org/pdf/2109.06153) for more details.

This file implements one approximate marginal inference oracle: Convex-GBP with fixed counting numbers of 1.0 for all regions. We experimented with others, but do not officially support them in this library. If interested, please see the following snapshot of this repository:

https://github.com/ryan112358/private-pgm/tree/approx-experiments-snapshot

Pull requests are welcome to add support for other approximate oracles.

Functions

`build_graph`	Builds the region graph for convex generalized belief propagation.
`convex_generalized_belief_propagation`	Convex generalized belief propagation for approximmate marginal inference.
`mirror_descent`	Mirror descent with approximate marginal inference.

Classes

`ApproxMirrorDescent`	Mirror descent estimator using approximate marginal inference.
`ApproxMirrorDescentState`	State for mirror descent with approximate marginal inference.
`StatefulMarginalOracle`	Defines the callable signature for stateful marginal oracle functions.

class mbi.approximate_oracles.StatefulMarginalOracle(*args, **kwargs)[source]

Bases: Protocol

Defines the callable signature for stateful marginal oracle functions.

A stateful marginal oracle computes (approximate) marginals from log-space potentials while also managing an internal state, often for optimization in iterative algorithms (e.g., preserving messages in message passing).

mbi.approximate_oracles.build_graph(domain: Domain, cliques: Sequence[tuple[str | int, ...]]) → tuple[set[tuple[str | int, ...]], list[tuple[str | int, ...]], dict[tuple[tuple[str | int, ...], tuple[str | int, ...]], Factor], list[tuple[tuple[str | int, ...], tuple[str | int, ...]]], dict[tuple[str | int, ...], list[tuple[str | int, ...]]], dict[tuple[str | int, ...], list[tuple[str | int, ...]]]][source]: Builds the region graph for convex generalized belief propagation.

mbi.approximate_oracles.convex_generalized_belief_propagation(potentials: CliqueVector, total: float = 1, state: dict[tuple[tuple[str | int, ...], tuple[str | int, ...]], Factor] | None = None, mesh: Mesh | None = None, iters: int = 1, damping: float = 0.5) → tuple[CliqueVector, dict[tuple[tuple[str | int, ...], tuple[str | int, ...]], Factor]][source]

Convex generalized belief propagation for approximmate marginal inference.

The algorithms implements the Algorithm 2 in our paper [“Relaxed Marginal Consistency for Differentially Private Query Answering”](https://arxiv.org/pdf/2109.06153), which itself is based on the paper titled [“Tightening Fractional Covering Upper Bounds on the Partition

Function for High-Order Region Graphs”](https://arxiv.org/pdf/1210.4881).

Parameters:

potentials – A CliqueVector object containing the potentials of the graphical model.
total – The total number of records in the dataset.
state – The state of the message passing algorithm (i.e., the messages). Useful when calling this within an iterative procedure for warm starting purposes.
mesh – Specifies how the computation will be sharded across machines.
iters – The number of iterations to run the algorithm.
damping – The damping factor for the messages.

Returns:

A CliqueVector of pseudo-marginals for the cliques in the graphical model.

class mbi.approximate_oracles.ApproxMirrorDescentState(potentials: CliqueVector, mu: CliqueVector, messages: Any)[source]

Bases: NamedTuple

State for mirror descent with approximate marginal inference.

Create new instance of ApproxMirrorDescentState(potentials, mu, messages)

potentials: CliqueVector: Alias for field number 0

mu: CliqueVector: Alias for field number 1

messages: Any: Alias for field number 2

class mbi.approximate_oracles.ApproxMirrorDescent(stepsize: float, oracle_iters: int = 1, damping: float = 0.5, mesh: Mesh | None = None)[source]

Bases: object

Mirror descent estimator using approximate marginal inference.

Uses convex_generalized_belief_propagation as the marginal oracle and warm-starts messages between optimization iterations. Unlike the exact mirror descent in estimation.py, this does not produce a MarkovRandomField because approximate region graphs cannot generate synthetic data.

This class holds optimizer configuration only. All data (loss function, potentials, totals) is passed to methods as arguments.

Example:

estimator = ApproxMirrorDescent(stepsize=1.0)
mu = estimator.estimate(domain, measurements)

stepsize

Fixed step size (required; no line search).

Type:: float

oracle_iters

Belief propagation iterations per optimization step.

Type:: int

damping

Damping factor for belief propagation messages.

Type:: float

mesh

JAX sharding mesh.

Type:: jax._src.mesh.Mesh | None

stepsize: float

oracle_iters: int = 1

damping: float = 0.5

mesh: Mesh | None = None

estimate(domain: ~mbi.domain.Domain, loss_fn: ~mbi.marginal_loss.MarginalLossFn | ~collections.abc.Sequence[~mbi.marginal_loss.LinearMeasurement], *, known_total: float | None = None, potentials: ~mbi.clique_vector.CliqueVector | None = None, iters: int = 1000, callback_fn=<function ApproxMirrorDescent.<lambda>>) → CliqueVector[source]

Run mirror descent and return pseudo-marginals.

Parameters:

domain – The domain over which the model is defined.
loss_fn – A MarginalLossFn or a list of LinearMeasurement.
known_total – The known or estimated number of records.
potentials – Initial potentials.
iters – Number of optimization iterations.
callback_fn – Called with pseudo-marginals at each iteration.

Returns:

Pseudo-marginals as a CliqueVector.

precompile(domain: Domain, measurements: Sequence[LinearMeasurement] | None = None, *, extra_cliques: Sequence[tuple[str | int, ...]] | None = None) → Future[source]

Warm up the JIT cache for estimate.

Parameters:

domain – The domain over which the model is defined.
measurements – Optional list of LinearMeasurement objects.
extra_cliques – Optional additional cliques to compile for.

mbi.approximate_oracles.mirror_descent(domain: Domain, loss_fn, *, known_total: float | None = None, potentials: CliqueVector | None = None, stepsize: float, iters: int = 1000, oracle_iters: int = 1, damping: float = 0.5, callback_fn=<function <lambda>>, mesh=None) → CliqueVector[source]

Mirror descent with approximate marginal inference.

Convenience wrapper around ApproxMirrorDescent.

Parameters:

domain – The domain over which the model should be defined.
loss_fn – A MarginalLossFn or a list of LinearMeasurement.
known_total – The known or estimated number of records.
potentials – Initial potentials.
stepsize – Fixed step size (required; no line search).
iters – Number of optimization iterations.
oracle_iters – Belief propagation iterations per optimization step.
damping – Damping factor for belief propagation messages.
callback_fn – Called with pseudo-marginals at each iteration.
mesh – JAX sharding mesh.

Returns:

Pseudo-marginals as a CliqueVector.