mbi.estimation.UniversalAcceleratedMethod

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

Methods

`__init__`
`estimate`	Estimate a Model from noisy marginal measurements.
`precompile`	Warm up the JIT cache for `estimate` asynchronously.

Attributes

`linesearch`
`marginal_oracle`
`max_iter_search`
`norm`
`target_acc`

marginal_oracle: MarginalOracle | None = None

max_iter_search: int = 30

target_acc: float = 0.0

norm: int = 2

linesearch: bool = False