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)

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:

• routing_history (Sequence[Any])

• merged_effective_frac (float)

• shattered_instability (float)

• shattered_edge_threshold (float)

• shattered_edge_frac (float)

Return type:

dict

factorize_dense_mlp_to_moe(dense_mlp, n_experts, *, common_fraction=0.5, top_k=1, aux_loss_weight=0.01, disjoint_loss_weight=0.01, probe_tokens=64, seed=0)

Split a dense transformer FFN into a common trunk and identical sparse residual experts.

Hidden neurons are ranked by the product of their input and output weight norms. The highest scoring common_fraction are evaluated once in the shared path; the rest seed every residual expert. A zero gate gives uniform probabilities and the forward-normalized top-k combination makes the selected identical experts sum to the original residual exactly. The conversion is therefore function-preserving up to floating-point summation before experts begin specializing.

Parameters:

• dense_mlp (Any)

• n_experts (int)

• common_fraction (float)

• top_k (int)

• aux_loss_weight (float)

• disjoint_loss_weight (float)

• probe_tokens (int)

• seed (int)

Return type:

tuple[Any, dict]

upcycle_dense_to_moe(dense_block, n_experts, *, top_k=1, seed=0, noise_std=0.01, probe_tokens=64, preservation_tolerance=0.2)

“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). Selected gate weights are forward-normalized, so noise_std=0 preserves the dense function up to the declared floating-point tolerance even with top-1 routing. With nonzero noise, a fixed probe batch measures the approximation and construction fails if the relative L2 gap exceeds preservation_tolerance.

Returns (moe_block, receipt) where the receipt records the measured gap, declared tolerance, preservation result, source substrate, and construction controls.

Parameters:

• dense_block (Any)

• n_experts (int)

• top_k (int)

• seed (int)

• noise_std (float)

• probe_tokens (int)

• preservation_tolerance (float)

Return type:

tuple[Any, dict]

upcycle_dense_to_shared_residual_moe(dense_block, n_experts, *, common_fraction=0.5, top_k=1, aux_loss_weight=0.01, disjoint_loss_weight=0.01)

Copy a transformer block and exactly factor its MLP into shared plus routed residual work.

Parameters:

• dense_block (Any)

• n_experts (int)

• common_fraction (float)

• top_k (int)

• aux_loss_weight (float)

• disjoint_loss_weight (float)

Return type:

tuple[Any, dict]