mbi.estimation

Algorithms for estimating graphical models from marginal-based loss functions.

This module provides a flexible set of optimization algorithms, each sharing the the same API. The supported algorithms are: 1. Mirror Descent [our recommended algorithm] 2. L-BFGS (using back-belief propagation) 3. Regularized Dual Averaging 4. Interior Gradient 5. Universal accelerated mirror descent

Each algorithm can be given an initial set of potentials, or can automatically intialize the potentials to zero for you. Any CliqueVector of potentials that support the cliques of the marginal-based loss function can be used here.

Functions

minimum_variance_unbiased_total

Estimates the total count from measurements with identity queries.

Classes

`DualAveraging`	Regularized Dual Averaging estimator for graphical models.
`DualAveragingState`	State for Regularized Dual Averaging (https://proceedings.neurips.cc/paper_files/paper/2009/file/7cce53cf90577442771720a370c3c723-Paper.pdf).
`Estimator`	An object that estimates a Model from a marginal-based loss function.
`InteriorGradient`	Interior Gradient estimator for graphical models.
`InteriorGradientState`	State for Interior Gradient (https://doi.org/10.1137/S1052623403427823).
`LBFGS`	L-BFGS estimator for graphical models.
`LBFGSState`	State for L-BFGS optimization on potentials.
`MirrorDescent`	Mirror descent estimator for graphical models.
`MirrorDescentState`	State for Algorithm 1 of https://arxiv.org/pdf/1901.09136.
`UniversalAcceleratedMethod`	Universal Accelerated Mirror Descent estimator.

class mbi.estimation.Estimator(*, marginal_oracle: MarginalOracle | None = None)[source]

Bases: ABC

An object that estimates a Model from a marginal-based loss function.

Subclasses implement _init, _step, and _finalize. The ABC provides a default estimate loop and a shared asynchronous precompile that works for any estimator with a jitted _step.

State convention: the first element of the state tuple returned by _init / _step is the current solution (passed to callback_fn). Override _callback_value only if the callback needs a transformation.

marginal_oracle

Callable for computing marginals from potentials. If None, auto-selected via default_oracle().

Type:: mbi.marginal_oracles.MarginalOracle | None

Examples of subclasses:

MirrorDescent
DualAveraging
InteriorGradient
extensions.MixtureOfProductsEstimator
extensions.ReweightedDatasetEstimator

marginal_oracle: MarginalOracle | None = None

estimate(domain: Domain, loss_fn: MarginalLossFn | list[LinearMeasurement], known_total: float | None = None, constraints: Sequence[Constraint] = (), iters: int = 1000, callback_fn: Callable | None = None, warm_start: Model | None = None, **kwargs: Any) → Model[source]

Estimate a Model from noisy marginal measurements.

If warm_start is provided, the estimator initializes from a previously estimated model instead of from scratch. For potential-based estimators the model’s potentials are expanded to cover any new cliques; for MixtureOfProducts the model is used directly.

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

Warm up the JIT cache for estimate asynchronously.

Returns a Future that completes when compilation finishes. Callers may ignore the return value (fire-and-forget) or call future.result() to block until compilation is done.

mbi.estimation.minimum_variance_unbiased_total(measurements: list[LinearMeasurement]) → float[source]: Estimates the total count from measurements with identity queries.

class mbi.estimation.MirrorDescentState(mu: CliqueVector, potentials: CliqueVector, alpha: jax.Array | float, loss: ArrayLike)[source]

Bases: NamedTuple

State for Algorithm 1 of https://arxiv.org/pdf/1901.09136.

Create new instance of MirrorDescentState(mu, potentials, alpha, loss)

mu: CliqueVector: Alias for field number 0

potentials: CliqueVector: Alias for field number 1

alpha: Array | float: Alias for field number 2

loss: Array | ndarray | bool | number | bool | int | float | complex: Alias for field number 3

class mbi.estimation.DualAveragingState(w: CliqueVector, v: CliqueVector, gbar: CliqueVector, loss: ArrayLike, lipschitz: jax.Array | float, gamma: jax.Array | float, t: jax.Array | int)[source]

Bases: NamedTuple

State for Regularized Dual Averaging (https://proceedings.neurips.cc/paper_files/paper/2009/file/7cce53cf90577442771720a370c3c723-Paper.pdf).

Create new instance of DualAveragingState(w, v, gbar, loss, lipschitz, gamma, t)

w: CliqueVector: Alias for field number 0

v: CliqueVector: Alias for field number 1

gbar: CliqueVector: Alias for field number 2

loss: Array | ndarray | bool | number | bool | int | float | complex: Alias for field number 3

lipschitz: Array | float: Alias for field number 4

gamma: Array | float: Alias for field number 5

t: Array | int: Alias for field number 6

class mbi.estimation.InteriorGradientState(x: CliqueVector, potentials: CliqueVector, c: jax.Array | float, y: CliqueVector, z: CliqueVector, loss: ArrayLike, inv_lipschitz: jax.Array | float)[source]

Bases: NamedTuple

State for Interior Gradient (https://doi.org/10.1137/S1052623403427823).

Create new instance of InteriorGradientState(x, potentials, c, y, z, loss, inv_lipschitz)

x: CliqueVector: Alias for field number 0

potentials: CliqueVector: Alias for field number 1

c: Array | float: Alias for field number 2

y: CliqueVector: Alias for field number 3

z: CliqueVector: Alias for field number 4

loss: Array | ndarray | bool | number | bool | int | float | complex: Alias for field number 5

inv_lipschitz: Array | float: Alias for field number 6

class mbi.estimation.LBFGSState(potentials: CliqueVector, opt_state: Any)[source]

Bases: NamedTuple

State for L-BFGS optimization on potentials.

Create new instance of LBFGSState(potentials, opt_state)

potentials: CliqueVector: Alias for field number 0

opt_state: Any: Alias for field number 1

class mbi.estimation.MirrorDescent(marginal_oracle: MarginalOracle | None = None, stepsize: float | None = None, mesh: Mesh | None = None)[source]

Bases: Estimator

Mirror descent estimator for graphical models.

This is a first-order proximal optimization algorithm for solving a (possibly nonsmooth) convex optimization problem over the marginal polytope. This is an implementation of Algorithm 1 from the paper “Graphical-model based estimation and inference for differential privacy”.

stepsize

Fixed step size, or None (default) to use Armijo line search.

Type:: float | None

marginal_oracle

The function to compute marginals from potentials. If None (default), uses default_oracle() to auto-select.

Type:: mbi.marginal_oracles.MarginalOracle | None

mesh

JAX sharding mesh.

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

stepsize: float | None = None

marginal_oracle: MarginalOracle | None = None

mesh: Mesh | None = None

class mbi.estimation.DualAveraging(marginal_oracle: MarginalOracle | None = None, mesh: Mesh | None = None)[source]

Bases: Estimator

Regularized Dual Averaging estimator for graphical models.

RDA is an accelerated proximal algorithm for solving a smooth convex optimization problem over the marginal polytope. This algorithm requires knowledge of the Lipschitz constant of the gradient of the loss function.

marginal_oracle

The function to compute marginals from potentials. If None (default), uses default_oracle() to auto-select.

Type:: mbi.marginal_oracles.MarginalOracle | None

mesh

JAX sharding mesh.

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

marginal_oracle: MarginalOracle | None = None

mesh: Mesh | None = None

class mbi.estimation.InteriorGradient(marginal_oracle: MarginalOracle | None = None, mesh: Mesh | None = None)[source]

Bases: Estimator

Interior Gradient estimator for graphical models.

Interior Gradient is an accelerated proximal algorithm for solving a smooth convex optimization problem over the marginal polytope. This algorithm requires knowledge of the Lipschitz constant of the gradient of the loss function. Based on the paper “Interior Gradient and Proximal Methods for Convex and Conic Optimization”.

marginal_oracle

The function to compute marginals from potentials. If None (default), uses default_oracle() to auto-select.

Type:: mbi.marginal_oracles.MarginalOracle | None

mesh

JAX sharding mesh.

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

marginal_oracle: MarginalOracle | None = None

mesh: Mesh | None = None

class mbi.estimation.LBFGS(marginal_oracle: MarginalOracle | None = None)[source]

Bases: Estimator

L-BFGS estimator for graphical models.

Optimizes the potentials (theta) directly via L-BFGS, back-propagating through the marginal inference oracle. The loss is convex w.r.t. marginals but typically non-convex w.r.t. potentials; in practice, L-BFGS still converges well.

See “Learning Graphical Model Parameters with Approximate Marginal Inference”.

marginal_oracle

The function to compute marginals from potentials. If None (default), uses default_oracle() to auto-select.

Type:: mbi.marginal_oracles.MarginalOracle | None

marginal_oracle: MarginalOracle | None = None

class mbi.estimation.UniversalAcceleratedMethod(marginal_oracle: MarginalOracle | None = None, max_iter_search: int = 30, target_acc: float = 0.0, norm: int = 2, linesearch: bool = False)[source]

Bases: Estimator

Universal Accelerated Mirror Descent estimator.

An accelerated first-order method that adapts to any smoothness level. Each optimization step performs an internal line-search via jax.lax.while_loop.

Numerical stability: Internally operates on the unit simplex (total=1) to keep potentials and gradients O(1). The user-facing loss loss_fn is evaluated on the N-scaled marginals via the wrapper f_scaled(p) = loss_fn(N * p). This prevents softmax saturation that otherwise causes the linesearch to degenerate for large N.

See Nesterov (2015) and Roulet & d’Aspremont (2017).

marginal_oracle

The function to compute marginals from potentials. If None (default), uses default_oracle() to auto-select.

Type:: mbi.marginal_oracles.MarginalOracle | None

max_iter_search

Max inner line-search iterations per step.

Type:: int

target_acc

Target accuracy (set > 0 for non-smooth objectives).

Type:: float

norm

Norm measuring smoothness (1 or 2).

Type:: int

linesearch

Whether to use adaptive line-search. Disabled by default because the acceptance condition uses the CliqueVector L2 norm which double-counts overlapping variables, causing the linesearch to accept overly large stepsizes and stall convergence.

Type:: bool

marginal_oracle: MarginalOracle | None = None

max_iter_search: int = 30

target_acc: float = 0.0

norm: int = 2

linesearch: bool = False