mixle.inference.leaf_hotswap module

Leaf hot-swap / analytic roll-up – swap a plateaued gradient leaf for a closed-form surrogate, with a retained path back on misfit (workstream D4).

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 and an update_kind classification – in particular "gradient" for a GradLeaf node, whose M-step is m_steps iterations of SGD/Adam rather than a closed-form update. D2 (mixle.inference.freeze_rollup) and D3 (mixle.inference.block_em) both spend a converged/near-zero D1 Q-gain on SCHEDULING decisions (skip an E-step recompute, skip a turn in this round’s M-step) while leaving the node’s own model object untouched. D4 goes one step further for gradient leaves specifically: once a gradient leaf’s own Q-gain has plateaued (gradient descent has stopped making meaningful progress, so the remaining M-step compute it would otherwise spend is close to wasted), swap it for a closed-form surrogate – a moment-matched MultivariateGaussianDistribution fit against the SAME (responsibility-weighted) data the gradient leaf was trained on – so every later round’s cost for that node collapses from O(param_count * _GRADIENT_STEPS) (D1’s own gradient M-step cost proxy) to O(param_count) (a closed-form MLE).

Correctness backbone (unchanged from the rest of the D-track): this is a SCHEDULING/specialization optimization only. Swapping a node’s model object for an approximation is more aggressive than D2/D3’s “leave the object alone, just skip recomputing/re-fitting it” story, so D4 earns back the Neal-Hinton guarantee two ways instead of one: (1) the swap itself is gated on a genuine, locally computed misfit RECEIPT (misfit_receipt()) – not merely assumed to be a good approximation – and (2) the ORIGINAL gradient leaf is never discarded (SwapRecord.original), so if the receipt later shows the surrogate drifting away from the real held-out data (e.g. the underlying regime shifts after the swap), swap_back() restores the exact retained object and gradient fitting resumes exactly where it left off – “never truly forget”, the same policy D2’s freeze/ roll-up commits to for frozen mixture components. F itself (the real Neal-Hinton free energy) is still the audit receipt: run_em_with_hotswap() gates every round’s proposal – the swap-in itself AND the following M-step (gradient OR closed-form), as one atomic unit – behind the same accept/reject monotone-F test D2/D3 already use, so a bad swap can cost speed (a rejected-and- reverted round, or a later swap-back once already committed) but never correctness. This is the gate that actually makes the mechanism safe: moment_matched_surrogate() on its own is a plain Gaussian MLE fit with no guarantee of matching an arbitrary (e.g. multi-modal) gradient leaf’s held-out density – see that function’s own docstring for a worked adversarial example and mixle.tests.leaf_hotswap_test.MonotoneObjectiveGateCatchesBadSwapTestCase for the regression test proving the gate rejects and fully reverts (not merely skips the M-step of) exactly that case. Earlier revisions of this module applied a plateau-triggered swap to the working tree unconditionally, before the round’s accept/reject check, and only skipped model = candidate on rejection – so a rejected round still silently returned the corrupted surrogate; the gate now rolls the swap itself back too when a round is rejected.

Scope: like D2/D3, this targets one gradient leaf embedded as a component of a MixtureDistribution (mixed freely with classical families, per mixle.models.grad_leaf’s whole point) – the “tree” in swap_leaf(tree, leaf_path, surrogate) is that mixture, and leaf_path is the component index. A single gradient leaf not embedded in any combinator (tree is the leaf itself) is also supported directly (leaf_path is ignored) since it is strictly the num_components=1 degenerate case of the same operation – useful for isolating the swap/misfit/swap-back mechanics from mixture E-step bookkeeping in tests. Generalizing past MixtureDistribution to arbitrary composite/sequence trees is the same “later items are expected to widen it” carve-out D2 documents for its own combinator scope.

Moment-matching machinery: G1’s moment_propagation.py (origin/moment-propagation, unmerged into this branch’s D1/D2/D3 chain – see this module’s PR description for how it was read via git show rather than a cross-branch merge) propagates a GAUSSIAN LAW through the specific layer types of mixle.models.transformer.CausalLM (Linear/LayerNorm/GELU/Attention), which is a different, narrower object than “moment-match an arbitrary gradient leaf’s behavior against arbitrary data” – it has no entry point that takes a generic GradLeaf and a data sample. Reusing it here would mean either constraining D4 to transformer-shaped leaves only (out of scope – GradLeaf wraps ANY torch density module) or re-deriving a generic version of its per-layer-law machinery, which is its own multi-week item. moment_matched_surrogate() therefore uses straightforward closed-form moment matching instead: draw/collect the SAME (possibly responsibility-weighted) data the gradient leaf was trained on and fit a MultivariateGaussianDistribution to it via the existing, real closed-form MLE machinery (MultivariateGaussianAccumulator / MultivariateGaussianEstimator) – an honest, fully-real “G1 machinery” scope per the roadmap item’s own explicit fallback clause.

class PlateauMonitor(*, q_gain_tol=_DEFAULT_PLATEAU_Q_GAIN_TOL, patience=_DEFAULT_PLATEAU_PATIENCE, n_mc=_DEFAULT_PLATEAU_N_MC)

Bases: object

Tracks, per leaf path, how many consecutive rounds a D1 NodeReport has reported a near-zero Q-gain for a "gradient"-update-kind node – the “gradient descent has stopped making meaningful progress” signal run_em_with_hotswap() swaps on.

Deliberately keyed and reset exactly like mixle.inference.freeze_rollup. FreezeRollupCache’s own _frozen_streak: a report that stops looking plateaued (the node moved again, or was swapped back to the original – see reset()) resets the streak to 0 rather than latching a stale verdict.

Parameters:

• q_gain_tol (float)

• patience (int)

• n_mc (int)

reset(path)

Clear the tracked history for path (e.g. after a swap or a swap-back).

Parameters:

path (Any)

Return type:

None

is_plateaued(path, leaf, *, nobs=None, seed=_SCORE_SEED)

Return whether leaf (identified by path) has plateaued this round.

Only a D1 update_kind == "gradient" node can plateau in this module’s sense (a closed-form/conjugate/frozen/em node has no “wasted SGD compute” to reclaim by swapping); anything else always returns False and resets the streak, so a leaf that has already been swapped for its (closed-form) surrogate is correctly reported as “not plateaued” going forward – there is nothing left to swap.

D1’s own residual is a MONTE-CARLO estimate (-mean(log_density(x)) over the node’s own self-samples, see mixle.inference.node_report’s module docstring) – for a genuinely converged gradient leaf, round-to-round Q-gain is dominated by MC sampling noise rather than any real drift in the fit, so this class uses a much larger n_mc than D1’s own default (_DEFAULT_PLATEAU_N_MC vs D1’s _DEFAULT_MC_SAMPLES=64) to push that noise floor down, and a correspondingly looser q_gain_tol than D2/D3’s exact-residual default (their residual is deterministic given unchanged parameters; this one never is).

Parameters:

• path (Any)

• leaf (Any)

• nobs (float | None)

• seed (int)

Return type:

bool

class SwapRecord(leaf_path, original, surrogate, swap_round, baseline_misfit, misfit_history=<factory>, swapped_back=False, swap_back_round=None)

Bases: object

The retrievable receipt of one leaf hot-swap – “never truly forget” (D2’s own phrase for its freeze/roll-up cache) applied to a swapped-out gradient leaf: original is the exact GradLeaf object that was in the tree before the swap, always retrievable, never discarded, so swap_back() can restore it byte-for-byte.

Parameters:

• leaf_path (Any)

• original (GradLeaf)

• surrogate (MultivariateGaussianDistribution)

• swap_round (int)

• baseline_misfit (float | None)

• misfit_history (list[float])

• swapped_back (bool)

• swap_back_round (int | None)

class LeafHotswapStats(round_index, n_components, n_frozen, n_swapped, n_log_density_evals, n_gradient_m_steps, n_closed_form_m_steps, objective, accepted=True, swapped_this_round=(), swapped_back_this_round=())

Bases: object

One round’s accounting for the hot-swap EM driver – mirrors mixle.inference.freeze_rollup.FreezeRollupStats / mixle.inference.block_em. BlockEMStats (same n_log_density_evals wall-clock proxy and real Neal-Hinton objective), plus the D4-specific n_gradient_m_steps / n_closed_form_m_steps split that is literally what “faster to same F” is measured against: a swapped component’s per-round M-step cost drops from D1’s param_count * _GRADIENT_STEPS proxy to param_count.

swapped_this_round is which components a plateau proposed a swap for this round – it is recorded even when accepted is False (a rejected round rolls the swap itself back out of the returned model/swap_records, but the attempt still happened and is worth a receipt); use n_swapped (or check idx in swap_records) for which swaps are actually COMMITTED as of this round.

Parameters:

• round_index (int)

• n_components (int)

• n_frozen (int)

• n_swapped (int)

• n_log_density_evals (int)

• n_gradient_m_steps (int)

• n_closed_form_m_steps (int)

• objective (float)

• accepted (bool)

• swapped_this_round (tuple[Any, ...])

• swapped_back_this_round (tuple[Any, ...])

moment_matched_surrogate(gradient_leaf, data, weights=None)

Fit a closed-form MultivariateGaussianDistribution that moment-matches data – the SAME (optionally responsibility-weights-ed) data gradient_leaf was trained on – via the real closed-form MLE machinery (weighted mean/covariance), not a re-derived formula.

This is a genuine fit, not a placeholder: it reads no attribute off gradient_leaf at all (a torch module has no portable closed-form “current mean/covariance” to read off directly – scoring/sampling is the only generic contract), so “moment-matched against the gradient leaf’s current behavior” means “against the data its current fit was scoring/trained on”, exactly the module-docstring’s documented (and roadmap-sanctioned) fallback to plain closed-form moment matching in lieu of G1’s transformer-specific law-propagation machinery.

gradient_leaf is accepted (rather than a bare module) purely as a type/documentation signal of intent – see the module docstring’s “why not G1” note – it is not otherwise used.

IMPORTANT, honest limitation – read before trusting this function’s output alone: a single Gaussian can only ever be as good an approximation as the true fitted density IS Gaussian. Against a unimodal, roughly-Gaussian-shaped leaf (the common case for a single mixture component pulling its own well-separated slice of responsibility-weighted data) this is an excellent approximation. Against a leaf whose OWN fitted density is multi-modal, heavy-tailed, or otherwise non-Gaussian, this function will silently produce a POOR approximation with no warning – e.g. on a genuinely bimodal GradLeaf (two well-separated modes, mixle.tests.leaf_hotswap_test.BimodalGauss) the moment-matched surrogate collapses both modes into one wide Gaussian sitting between them, degrading held-out NLL by roughly 80% relative to the original leaf (see mixle.tests.leaf_hotswap_test.MonotoneObjectiveGateCatchesBadSwapTestCase). This function provides NO guarantee, on its own, that the surrogate’s held-out density tracks the original leaf’s – that guarantee, to the extent one exists, is earned entirely by run_em_with_hotswap()’s per-round monotone-F accept/reject gate (see that function’s own docstring): a swap-plus-refit round that does not improve the real Neal-Hinton objective is rejected and reverted, INCLUDING the swap itself, not merely the following M-step. Do not call moment_matched_surrogate() outside that gated driver (or an equivalent one) and assume the result is a safe stand-in for gradient_leaf – verify with a real misfit receipt (misfit_receipt()) against genuinely held-out data first.

Parameters:

• gradient_leaf (GradLeaf)

• data (Any)

• weights (ndarray | None)

Return type:

MultivariateGaussianDistribution

misfit_receipt(surrogate, holdout_data)

A genuine, locally-computed misfit scalar for surrogate on REAL held-out data: its own negative log-likelihood, -mean(log_density(x)). Not a placeholder – recomputed from real samples every call, exactly the same style of receipt D1’s residual and G1’s closure_error both are (see the respective module docstrings). Lower is better; compared against SwapRecord.baseline_misfit (the receipt measured right after the swap) by should_swap_back() to decide whether the surrogate has since drifted away from the real data it was swapped in to approximate.

Parameters:

• surrogate (MultivariateGaussianDistribution)

• holdout_data (Any)

Return type:

float

swap_leaf(tree, leaf_path, surrogate, *, round_index=0)

Replace the plateaued gradient leaf at leaf_path in tree with surrogate.

Returns (new_tree, swap_record): new_tree has the surrogate in place (every generic D1/D2/D3 mechanism – node_report(), detect_frozen(), seq_log_density – sees an ordinary MultivariateGaussianDistribution node from here on, no special-casing required upstream); swap_record.original retains the exact GradLeaf that was swapped out, per this module’s “never discard capacity” policy (see SwapRecord).

Parameters:

• tree (Any)

• leaf_path (Any)

• surrogate (MultivariateGaussianDistribution)

• round_index (int)

Return type:

tuple[Any, SwapRecord]

swap_back(tree, swap_record, *, round_index=0)

Restore swap_record.original (the retained gradient leaf) into tree at swap_record.leaf_path, undoing swap_leaf(). Marks swap_record as swapped back (in place) so callers/tests can assert the misfit receipt actually fired.

Parameters:

• tree (Any)

• swap_record (SwapRecord)

• round_index (int)

Return type:

Any

run_em_with_hotswap(enc_data, estimator, initial_model, *, holdout_data=None, max_its=10, delta=1.0e-9, cache=None, monitor=None, freeze_q_gain_tol=1.0e-6, plateau_q_gain_tol=_DEFAULT_PLATEAU_Q_GAIN_TOL, plateau_patience=_DEFAULT_PLATEAU_PATIENCE, plateau_n_mc=_DEFAULT_PLATEAU_N_MC, misfit_tol=_DEFAULT_MISFIT_TOL, accept_tolerance=_DEFAULT_ACCEPT_TOLERANCE)

Run EM over a MixtureDistribution with D4 leaf hot-swap: any component D1 reports as a plateaued gradient leaf (see PlateauMonitor) is swapped for a moment-matched closed-form surrogate (moment_matched_surrogate(), fit against that round’s responsibility-weighted data), which is then updated by a closed-form re-fit every round instead of gradient descent – and swapped back to the retained original the instant a real misfit receipt (misfit_receipt() on holdout_data, when supplied) shows it has drifted (see should_swap_back()).

Reuses D2’s FreezeRollupCache/detect_frozen for ordinary converged-component freezing (composes with D4 exactly like D3 does: a frozen component is excluded from both the swap check and the M-step) and the same per-round objective accept/reject gate D2/D3 use, so history is a real monotone-F receipt: swapping in a surrogate, or swapping back out of one, can only ever be accepted if the round’s real Neal-Hinton objective does not decrease.

holdout_data, if given, is treated as belonging to whichever component the CURRENT model would assign it to most responsibly at swap time (a single fixed sample scored against every swapped component’s surrogate) – a simplification documented here rather than a full per-component responsibility-weighted holdout split, mirroring D2/D3’s own explicit single-combinator scope carve-outs (see this module’s docstring).

Returns (final_model, history, swap_records) where swap_records maps component index to every SwapRecord created during the run (including ones later swapped back), so a caller/test can retrieve the retained original gradient leaf and inspect misfit_history.

Parameters:

• enc_data (Any)

• estimator (MixtureEstimator)

• initial_model (MixtureDistribution)

• holdout_data (Any)

• max_its (int)

• delta (float | None)

• cache (FreezeRollupCache | None)

• monitor (PlateauMonitor | None)

• freeze_q_gain_tol (float)

• plateau_q_gain_tol (float)

• plateau_patience (int)

• plateau_n_mc (int)

• misfit_tol (float)

• accept_tolerance (float)

Return type:

tuple[MixtureDistribution, list[LeafHotswapStats], dict[int, SwapRecord]]