mixle.inference.block_em module

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

The estimator tree is an update IR. After an initial full sweep, this scheduler ranks mixture components by the complete-data Q improvement measured the last time that component was updated, divided by a structural parameter-count cost. It updates the highest-value blocks within a per-round budget and bounds starvation with a forced periodic update.

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 a generalized-EM coordinate update. A nonnegative complete-data Q gain certifies that the observed objective cannot decrease. Periodic and final exact objective audits verify that certificate; a candidate without a finite certificate falls back to an exact gate. With parameter priors, every round is gated on the MAP objective.

The scheduler reuses FreezeRollupCache only for blocks left inactive in the current round. It does not permanently freeze a component from a heuristic residual; max_skip_rounds guarantees every nonzero-weight block is revisited. Callers wanting the explicit near-zero-weight permanent-freeze policy can use run_em_freeze_rollup directly.

The measured gain is tied to the real E-step responsibilities: for active component k it is sum_i gamma_ik [log p_new(x_i,k) - log p_old(x_i,k)], including the weight-coordinate term. It is therefore a stale scheduling estimate, not a convergence proof. Correctness comes from the Q certificate and observed-objective audits; ranking only changes work order.

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 BlockEMAssumptionReceipt(declared_budget_fraction, eligible_cost, selected_cost, forced_cost, expected_density_evaluations, observed_density_evaluations, expected_score_reuses, observed_score_reuses, schedule_reused=False, full_refresh=False, cost_basis='structural_parameter_count', expected_responsibility_columns=0, observed_responsibility_columns=0, incremental_normalizer=False, normalizer_audit_max_abs_error=0.0, normalizer_logspace_rows=0, failures=())[source]

Bases: object

Observed checks for assumptions connecting selected work to saved work.

Parameters:

  • declared_budget_fraction (float)

  • eligible_cost (float)

  • selected_cost (float)

  • forced_cost (float)

  • expected_density_evaluations (int)

  • observed_density_evaluations (int)

  • expected_score_reuses (int)

  • observed_score_reuses (int)

  • schedule_reused (bool)

  • full_refresh (bool)

  • cost_basis (str)

  • expected_responsibility_columns (int)

  • observed_responsibility_columns (int)

  • incremental_normalizer (bool)

  • normalizer_audit_max_abs_error (float)

  • normalizer_logspace_rows (int)

  • failures (tuple[str, ...])

class BlockEMTimingReceipt(scheduling_seconds, estep_density_seconds, responsibility_seconds, snapshot_seconds, mstep_seconds, candidate_density_seconds, validation_seconds, accounting_seconds, total_seconds, component_density_seconds=(), component_update_seconds=(), component_block_seconds=())[source]

Bases: object

Wall-time decomposition for one scheduler round.

Parameters:
class BlockEMStats(round_index, n_components, n_active, n_frozen, n_zero_weight, n_dormant, n_scheduled_inactive, n_log_density_evals, n_responsibility_columns, objective, accepted=True, acceptance_basis='observed', objective_exact=True, degenerate_round=False, measured_q_gain=0.0, certified_q_gain=0.0, wall_time_seconds=0.0, final_audit_seconds=0.0, assumptions=None, timing=None, stop_reason=None, escape_attempted=False)[source]

Bases: object

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

n_log_density_evals records removed model work, while wall_time_seconds records actual elapsed time. objective is either the exact observed-data MLE/MAP objective or its certified EM lower bound, as disclosed by objective_exact; measured_q_gain is the complete-data gain used for the next scheduling decision.

escape_attempted marks the one full-tree proposal allowed after repeated rejections. That proposal remains objective-gated; acceptance_basis == "vanilla_escape" means it passed the gate. stop_reason is set on the final entry when no gated operator progresses.

Parameters:

  • round_index (int)

  • n_components (int)

  • n_active (int)

  • n_frozen (int)

  • n_zero_weight (int)

  • n_dormant (int)

  • n_scheduled_inactive (int)

  • n_log_density_evals (int)

  • n_responsibility_columns (int)

  • objective (float)

  • accepted (bool)

  • acceptance_basis (str)

  • objective_exact (bool)

  • degenerate_round (bool)

  • measured_q_gain (float)

  • certified_q_gain (float)

  • wall_time_seconds (float)

  • final_audit_seconds (float)

  • assumptions (BlockEMAssumptionReceipt | None)

  • timing (BlockEMTimingReceipt | None)

  • stop_reason (str | None)

  • escape_attempted (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.

Keyed (tied) estimators are ineligible: tying pools statistics ACROSS sites via the merge_accumulator_keys pass, and block-EM’s sparse per-component M-steps neither run that pass nor could honor it without re-estimating every tied sibling of an updated block – exactly the silent-untying bug class #432 fixed in the posterior-transform strategies. schedule=’auto’ therefore keeps keyed fits on the full-tree path (this predicate), and an EXPLICIT run_block_em call refuses loudly instead of silently untying.

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, escape_after_rejections=3, max_skip_rounds=2, boundary_weight_step=_DEFAULT_BOUNDARY_WEIGHT_STEP, full_refresh_interval=10, schedule_wave_rounds=1, objective_audit_interval=10, compute_dtype=None, cost_model='structural', policy='greedy', controller=None)[source]

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

The first round updates every component. Later rounds rank blocks by the observed Q gain from their last update per structural-cost unit, select the highest-value blocks within budget_fraction, reuse cached component scores for untouched blocks, and transactionally gate the resulting partial M-step on the exact MLE or MAP objective.

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-objective 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.

escape_after_rejections bounds a REJECTION livelock. A rejected round restores the exact prior state (score caches included), so the next round regenerates a near-identical proposal: when every proposal the scheduler can produce regresses observed likelihood (seen in practice with stochastic GradLeaf M-steps once the closed-form blocks reach their conditional fixed point), the loop would otherwise pay full proposal cost forever without moving. After that many CONSECUTIVE rejected rounds, one exact vanilla full-tree proposal is attempted through the same observed-objective gate (escape_attempted=True and, only when accepted, acceptance_basis="vanilla_escape"). If it also rejects, no available operator establishes progress, so the run stops early with stop_reason="rejection_livelock" rather than accepting a loss.

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 (the gated objective 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 objective, fewer evals” (as opposed to just “monotone but permanently stuck”) an honest comparison against vanilla EM.

objective_audit_interval controls exact observed-likelihood audits of Q-certified rounds; None disables periodic audits, but the final model is always audited. A non-finite or negative Q certificate always triggers an exact fallback gate. schedule_wave_rounds can amortize one selection over several rounds, but defaults to one because repeated coordinates can delay complementary updates; starvation immediately invalidates a pending wave.

boundary_weight_step relaxes only a sparse weight update that would turn a previously positive component into an exact zero. The relaxed point lies on the segment from the current feasible weights to the conditional maximizer, so concavity preserves Q ascent while avoiding an irreversible floating-point simplex boundary. Its default leaves a machine-precision-derived interior margin. Components with exactly zero responsibility retain their existing parameters; positive components below weight_tol remain dormant between full refreshes rather than consuming ordinary block budget.

full_refresh_interval schedules a globally coupled EM update periodically. This bounds stale block interactions and gives tiny-but-positive components a dense responsibility update; None disables refreshes. Refresh rounds are disclosed separately from score degeneracy and reset any pending schedule wave.

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 observed-objective improvement, round_value - current_value, i.e. exactly what this round’s active blocks achieved) and realized model-work cost (evals_e + evals_c) 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.

q_gain_tol, weight_delta_tol, and freeze_patience remain accepted for API compatibility with earlier scheduler releases, but permanent freeze detection is no longer part of this path. weight_tol controls temporary dormancy only. Use mixle.inference.freeze_rollup.run_em_freeze_rollup() when those controls are desired.

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)

  • escape_after_rejections (int)

  • max_skip_rounds (int)

  • boundary_weight_step (float)

  • full_refresh_interval (int | None)

  • schedule_wave_rounds (int)

  • objective_audit_interval (int | None | str)

  • compute_dtype (Any)

  • cost_model (str)

  • policy (str | LearnedController)

  • controller (LearnedController | None)

Return type:

tuple[MixtureDistribution, list[BlockEMStats]]