mixle.models.moe module

Mixture-of-experts transformer MLP – the neural half of ConditionalJIT’s structural adaptation (roadmap H2).

The user-stated goal “the model adapts its structure as it trains,” applied to Block’s MLP: instead of one dense feed-forward network every token runs through, MoEBlock gives each token a per-token choice of N expert feed-forward networks via a learned linear gate, so the network’s effective structure (which parameters a given token’s forward pass touches) is decided at train/inference time rather than fixed at init.

This is the same “mixture” idea mixle already has a name for. mixle.stats.latent.mixture.MixtureDistribution defines P(Y) = sum_k P(Y|Z=k) P(Z=k): a soft responsibility (posterior P(Z=k|Y), fit by EM) selects which of K homogeneous component distributions explains a data point. A gradient-trained MoE gate is the same combinator with a different fitting mechanism: softmax(W x) plays the role of the responsibility P(Z=k|Y) and each expert MLP plays the role of a mixture component, but the routing distribution is trained end-to-end by gradient descent through a load-balancing auxiliary loss rather than by EM’s alternating E/M steps (MoE’s per-token hard top-k selection is also a discrete argmax over an otherwise continuous responsibility, unlike EM’s fully soft E-step). That structural identity is not just a metaphor: the gate’s (n_tokens, n_experts) softmax output has exactly the shape of the token-by-component responsibility matrix z that mixle.utils.hvis.topology.model_fit_health()/fuzzy_nerve() consume for probabilistic mixtures, so expert_collapse_receipt() below feeds routing weights into fuzzy_nerve directly – the exact same overlap-nerve computation HViS uses to flag a mixture’s merged/shattered component regimes, re-aimed at expert routing statistics instead of clustering posteriors. See that function’s docstring for the re-aimed semantics.

Two entry points:

  • MoEBlock – drop-in replacement for Block (attention unchanged; the dense MLP is replaced by MoEMLP, N expert MLPs plus a top-k linear gate and the standard Switch-Transformer load-balancing auxiliary loss).

  • upcycle_dense_to_moe() – turn an already-trained dense Block into an MoEBlock by copying attention unchanged and initializing every expert as a near-copy of the dense MLP (the standard “sparse upcycling” trick: Komatsuzaki et al., 2023), carrying a function-preservation-style receipt (how close the freshly-upcycled model’s output is to the original dense block’s output, before any MoE-specific training happens).

  • expert_collapse_receipt() – the balance/collapse receipt: “merged” (routing collapsed onto a handful of experts) and “shattered” (routing so unstable round-to-round that no expert receives a consistent, learnable token distribution), reusing mixle.utils.hvis.topology.fuzzy_nerve().

Torch is imported lazily/guarded exactly like transformer.py so this module still imports (as a no-op) when torch is not installed.

expert_collapse_receipt(routing_history, *, merged_effective_frac=0.5, shattered_instability=0.35, shattered_edge_threshold=0.3, shattered_edge_frac=0.5)[source]

Load-balance / expert-collapse receipt, reusing mixle.utils.hvis.topology.fuzzy_nerve() – the SAME overlap-nerve computation model_fit_health uses to flag a mixture’s merged/shattered component regimes – re-aimed at MoE routing statistics.

routing_history is a sequence of per-round gate softmax matrices (each (n_tokens_r, n_experts), e.g. MoEBlock.routing_weights collected once per training step/round). Each matrix has exactly the shape of the token-by-component responsibility matrix z that fuzzy_nerve was built for, so it is fed in directly – no adapter needed.

Two failure modes, re-aimed from the original clustering semantics:

  • merged – fuzzy_nerve’s masses (per-expert claimed-token mass, pooled over every round) are so concentrated that the effective number of experts in use, exp(entropy(utilization)), drops below merged_effective_frac * n_experts. This is the routing analogue of the original merged-regime detector: instead of “one COMPONENT secretly covers two regimes,” it is “routing has secretly collapsed onto fewer experts than exist,” measured with the same entropy-of-mass machinery.

  • shattered – the original shattered detector flagged near-duplicate components via fuzzy_nerve edge weight; here that is generalized across time: (a) per-ROUND utilization is so unstable (large round-to-round total-variation distance in per-expert mass fractions) that no expert sees a consistent token distribution to specialize on, and/or (b) the POOLED nerve has strong overlap edges (fuzzy_nerve’s literal near-duplicate-component signal) across a large fraction of expert pairs, meaning experts are not actually claiming distinguishable token sets.

A well-balanced run (near-uniform utilization, stable round to round) trips neither flag.

Parameters:
Return type:

dict

upcycle_dense_to_moe(dense_block, n_experts, *, top_k=1, seed=0, noise_std=0.01, probe_tokens=64)[source]

“Sparse upcycling” (Komatsuzaki et al., 2023): build a fresh MoEBlock whose attention is copied unchanged from dense_block and whose n_experts expert MLPs are each initialized as a near-copy of dense_block’s trained dense MLP (exact weights + small seeded Gaussian perturbation per expert, so experts start distinguishable rather than identical dead-gradient copies). Unlike H1’s growth operators this is NOT exactly function-preserving – the gate’s top-k hard selection is a nonlinearity the dense path never had, so upcycled output only APPROXIMATES the original dense output. That approximation is measured, not assumed: a fixed probe batch is pushed through both blocks and the relative L2 output gap is returned in the receipt.

Returns (moe_block, receipt) where receipt has relative_output_diff (the measured gap, expected small but nonzero), n_experts, top_k, noise_std, and seed.

Parameters:
Return type:

tuple[Any, dict]