mixle.inference.block_em module

Block-coordinate-ascent EM scheduler – greedy gain-per-cost block selection (workstream D3).

Frame (see the ConditionalJIT track, D1-D6): the estimator tree is an IR. D1 (mixle.inference.node_report) instruments every node with a per-round residual/Q-gain report, an update_kind classification, and an E/M cost proxy. D2 (mixle.inference.freeze_rollup) spent that report on one particular schedule: freeze a subtree once it looks converged and skip it forever after. D3 generalizes the scheduling question one level up: on EVERY round, rank the blocks (mixture components) that are NOT already D2-frozen by D1’s gain-per-cost (q_gain / (e_step_cost + m_step_cost)) and update only the highest-value ones within a per-round cost budget, leaving the rest untouched for that round – a genuine block-coordinate-ascent / ECM schedule, not just “freeze forever”.

Correctness backbone (unchanged from the rest of the D-track): this module is a SCHEDULING optimization only. Any interleaving of a partial E-step (fresh log-density for the active blocks, cached log-density for the inactive ones) and a per-block conditional M-step (only the active blocks are re-estimated; every other block’s model object is carried forward byte-for- byte unchanged) is still coordinate ascent on the SAME Neal-Hinton free energy F vanilla EM climbs – so an accept/reject gate on the round’s candidate objective (identical in spirit to mixle.inference.em.MonotonicEM and D2’s own run_em_freeze_rollup) is what turns “should be monotone” into “IS monotone, mechanically, every round” here too.

Composition with D2: a component D2’s FreezeRollupCache already reports frozen (mixle.inference.freeze_rollup.detect_frozen()) is, from this scheduler’s point of view, exactly a “zero e-step cost, zero m-step cost” block – it is excluded from the gain-per-cost ranking entirely (nothing to rank: it costs nothing and is never scheduled to move) and is served for free from the SAME cache this module reuses for its own this-round-only inactive blocks. A block the scheduler chooses not to update THIS round is, from FreezeRollupCache’s point of view, indistinguishable from a D2-frozen block for exactly one round: its parameter signature has not moved, so the cached per-datum log-density is still a byte-identical cache hit – no separate caching mechanism is needed for D3, only a wider set of “don’t touch this round” indices fed into the same frozen= parameter D2 already threads through _component_log_density_matrix() and _m_step().

Gain estimate, documented per the D1 module’s own note that later track items may re-estimate a cheap proxy fresh every round rather than reuse a stale report: this module calls mixle.inference.node_report.node_report() on every eligible component EVERY round (a small, fixed-size Monte-Carlo self-residual – _DEFAULT_MC_SAMPLES samples, not a real-data pass), always with the SAME seed across components (unlike mixle.inference.node_report. flat_report_table(), which offsets the seed per row for a deduplicated tree walk) so that structurally-identical components draw directly-comparable residuals – this is what makes the “no useful discrimination to make” acceptance criterion (identical components => identical scores => no real ranking) hold deterministically rather than by RNG luck.

Degeneration to vanilla full-tree EM: when every eligible block’s gain-per-cost score is the same (within tie_tol) – either because the blocks truly are indistinguishable, or because the caller passes full_tree_every_round=True as an explicit escape hatch – the scheduler has no real choice to make, so it updates every eligible block, exactly like a plain full-tree EM round would. See mixle.tests.block_em_test for a literal test of this property.

class BlockEMStats(round_index, n_components, n_active, n_frozen, n_zero_weight, n_scheduled_inactive, n_log_density_evals, objective, accepted=True, degenerate_round=False)[source]

Bases: object

One round’s accounting for the block-EM scheduler – the acceptance-criteria receipt.

Mirrors mixle.inference.freeze_rollup.FreezeRollupStats (same n_log_density_evals wall-clock proxy and real Neal-Hinton objective), plus the extra fields specific to gain-per-cost SCHEDULING within a round rather than permanent freezing.

Parameters:

  • round_index (int)

  • n_components (int)

  • n_active (int)

  • n_frozen (int)

  • n_zero_weight (int)

  • n_scheduled_inactive (int)

  • n_log_density_evals (int)

  • objective (float)

  • accepted (bool)

  • degenerate_round (bool)

property active_fraction: float

Fraction of components genuinely (re-)evaluated this round, out of all components.

is_block_em_eligible(model, estimator)[source]

Whether model/estimator are a plain MixtureDistribution/MixtureEstimator pair the D3 scheduler can drive. Lives here (not in estimation.py) so the high-level estimation driver never imports a concrete distribution type – see compute_metadata_test.py’s layering check.

Parameters:
Return type:

bool

run_block_em(enc_data, estimator, initial_model, *, max_its=10, delta=1.0e-9, cache=None, budget_fraction=_DEFAULT_BUDGET_FRACTION, full_tree_every_round=False, tie_tol=_DEFAULT_TIE_TOL, accept_tolerance=_DEFAULT_ACCEPT_TOLERANCE, q_gain_tol=1.0e-6, weight_tol=1.0e-4, weight_delta_tol=1.0e-8, freeze_patience=3, stall_patience=5, max_skip_rounds=2, policy='greedy', controller=None)[source]

Run block-coordinate-ascent EM over a MixtureDistribution (workstream D3, D5).

Each round: rank every component D2 does not already report frozen by D1 gain-per-cost, select the highest-value ones within budget_fraction of the round’s total eligible cost (the rest are left untouched for this round), do a fresh partial E-step (cached log-density reused for every untouched/frozen component – same mechanism D2’s FreezeRollupCache already provides), a per-block conditional M-step over only the active components, and an accept/reject gate on the round’s real objective – exactly D2’s own monotone-F machinery, reused rather than reimplemented.

schedule="auto" at the mixle.inference.estimation.optimize() layer dispatches here; budget_fraction=1.0 or full_tree_every_round=True both degenerate this to (numerically indistinguishable from) vanilla full-tree EM – see mixle.tests.block_em_test for a test of that literal property.

delta convergence is gated by stall_patience consecutive small-gain rounds rather than a single one: a partial (budget-throttled) round can legitimately show a tiny total-F gain even while a still-improving block simply wasn’t scheduled that round (its turn is coming), so a single-round delta check – fine for vanilla EM, where every round touches every block – would risk declaring convergence early here. Requiring the plateau to persist for stall_patience rounds is the same style of robustness D2’s own freeze_patience already uses for its (structurally identical) “has this genuinely stopped moving” question.

max_skip_rounds bounds STARVATION: a purely greedy top-score-wins ranking can, on a real fixture, rank the same eligible block last round after round (its own gain-per-cost score stays a little below its rivals’) and never actually get a turn – which is still valid coordinate ascent (F still only goes up), but converges to a WORSE fixed point than vanilla EM would reach in the same number of rounds, since a legitimately-still-moving block is parked indefinitely instead of merely delayed. Any eligible block skipped max_skip_rounds rounds in a row is therefore force-included in active regardless of its score (an aging boost) – guaranteeing every eligible block gets a turn at least once every max_skip_rounds + 1 rounds, which is what makes “same target F, fewer evals” (as opposed to just “monotone but permanently stuck”) an honest comparison against vanilla EM.

policy (workstream D5, mixle.inference.conditional_jit_controller) selects WHO picks the per-round budget_fraction that feeds the exact same _select_active() ranking above: "greedy" (the default) uses the fixed budget_fraction argument every round, unchanged D3 behavior. "learned_bandit"/"learned_design_model" instead ask a LearnedController for this round’s budget (constructing a default BanditController/ DesignModelController if controller is not supplied) and feed it back the round’s REALIZED gain (this round’s own F improvement, round_value - current_value, i.e. exactly what this round’s active blocks achieved) and REALIZED cost (evals_e + evals_c, the same wall-clock proxy BlockEMStats already reports) after every round – so the controller learns from the SAME accept/reject-gated, provably-monotone rounds D3 already runs, never from an invented parallel objective. Passing an already-constructed LearnedController directly as policy is also accepted (equivalent to passing it as controller with policy="learned_bandit"/"learned_design_model") and is how a caller warm-starts a controller across several calls/fits – the SAME controller object keeps learning, since it is only ever mutated in place, never copied.

Returns (final_model, history) where history[i] is round i’s BlockEMStats.

Parameters:

  • enc_data (Any)

  • estimator (MixtureEstimator)

  • initial_model (MixtureDistribution)

  • max_its (int)

  • delta (float | None)

  • cache (FreezeRollupCache | None)

  • budget_fraction (float)

  • full_tree_every_round (bool)

  • tie_tol (float)

  • accept_tolerance (float)

  • q_gain_tol (float)

  • weight_tol (float)

  • weight_delta_tol (float)

  • freeze_patience (int)

  • stall_patience (int)

  • max_skip_rounds (int)

  • policy (str | LearnedController)

  • controller (LearnedController | None)

Return type:

tuple[MixtureDistribution, list[BlockEMStats]]