#5 LSTM + Mixture of Experts in JAX / Flax / Optax

A hand-written LSTM with a sparse Mixture-of-Experts FFN between two recurrent layers, trained on Tiny Shakespeare. Anchor: Shazeer et al., Outrageously Large Neural Networks (2017) — the original LSTM-with-MoE paper.

Full runnable notebook: createcentury/jax-flax-optax-lab — 02-lstm-moe.

Why this stack

Two angles. First, MoE is back: Mixtral, DeepSeek-V3, and most frontier LLMs use sparse experts because the compute-per-token decouples from the parameter count. The mechanism is small but subtle — a router, a top-$k$ selection, and an auxiliary loss to keep things balanced. Building it once from scratch is the cleanest way to see the moving parts.

Second, JAX/Flax/Optax fits this kind of mechanism unusually well. The router state, the gating weights, and the load-balancing loss all live as explicit tensors and intermediates. Nothing is hidden inside a forward() that mutates a self somewhere — every byte of state is in the PyTree.

Task

Character-level next-token prediction on Tiny Shakespeare. ~1MB of text, 65 unique characters, batches of (B=64, T=128). Loss is the average per-position cross-entropy:

$$\mathcal{L}_\text{CE} = -\frac{1}{B T}\sum_{b,t} \log p_\theta(y_{b,t} \mid x_{b, \le t})$$

Chance level is $\log 65 \approx 4.17$ nats. A trained char-LSTM in this regime sits at around 1.0–1.2 with top-1 accuracy 55–62%.

Architecture

tokens [B, T]
  → Embed
  → LSTM₁              (hidden = 256)
  → MoE FFN block      (4 experts, top-2, residual)
  → LSTM₂              (hidden = 256)
  → Linear head → next-token logits

The MoE substitutes for what would otherwise be a single shared FFN between LSTM layers. Each token picks two experts out of four; the other two contribute zero.

1. LSTM cell

I wrote the cell explicitly rather than calling flax.linen.LSTMCell — the recurrence is the point.

$$\begin{aligned} i_t &= \sigma(W_{xi} x_t + W_{hi} h_{t-1} + b_i) \\ f_t &= \sigma(W_{xf} x_t + W_{hf} h_{t-1} + b_f) \\ g_t &= \tanh(W_{xg} x_t + W_{hg} h_{t-1} + b_g) \\ o_t &= \sigma(W_{xo} x_t + W_{ho} h_{t-1} + b_o) \\ c_t &= f_t \odot c_{t-1} + i_t \odot g_t \\ h_t &= o_t \odot \tanh(c_t) \end{aligned}$$

The four input projections share the same matrix shape, so they get fused into a single Dense(4D) followed by a split. Same for the hidden projections. In Flax:

class LSTMCell(nn.Module):
    hidden_dim: int

    @nn.compact
    def __call__(self, carry, x):
        h, c = carry
        D = self.hidden_dim
        gates = (
            nn.Dense(4 * D, use_bias=True,  name='x2g')(x) +
            nn.Dense(4 * D, use_bias=False, name='h2g')(h)
        )
        i, f, g, o = jnp.split(gates, 4, axis=-1)
        i, f, o = jax.nn.sigmoid(i), jax.nn.sigmoid(f), jax.nn.sigmoid(o)
        g = jnp.tanh(g)
        new_c = f * c + i * g
        new_h = o * jnp.tanh(new_c)
        return (new_h, new_c), new_h

nn.scan lifts this cell to a sequence module — params are broadcast across time, RNG keys are not split. That's the canonical Flax pattern for recurrence.

2. The MoE block

Router

A single dense projection from the residual stream to per-expert logits:

$$\ell = W_g\, x \in \mathbb{R}^E$$

with $W_g \in \mathbb{R}^{D \times E}$, $D$ = hidden dim, $E$ = number of experts.

Top-$k$ gating

For each token, keep the $k$ largest logits, zero the rest, and re-softmax within the surviving set:

$$T = \operatorname{top-}k(\ell), \quad g_i = \begin{cases} \dfrac{\exp(\ell_i)}{\sum_{j \in T} \exp(\ell_j)} & i \in T \\[6pt] 0 & i \notin T \end{cases}$$

The re-softmax matters: the surviving gate weights sum to 1, which keeps the output norm sane regardless of $k$.

Output

Each expert $E_i: \mathbb{R}^D \to \mathbb{R}^D$ is a small FFN (Linear → GELU → Linear). The block output is a residual:

$$y = x + \sum_{i=1}^E g_i \cdot E_i(x)$$

Because $E - k$ of the gates are zero, only $k$ experts actually contribute per token. In principle. In this notebook I compute all experts on all tokens and zero out the unwanted contributions afterwards — wasteful, but conceptually transparent. The sparse-dispatch version (token-to-expert scatter, expert-to-token gather) belongs in a follow-up.

Aux loss — the part that actually matters

Top-$k$ routing has a degenerate solution: route every token to the same expert and let the other $E-1$ go unused. The loss decreases just fine because that one expert can specialise on the whole distribution. To prevent collapse, Shazeer 2017 / Switch Transformer (Fedus 2022) add an auxiliary load-balancing loss.

Let $N$ be the number of tokens in the batch.

• Hard usage — fraction of tokens for which expert $i$ is in the top-$k$:

$$f_i = \frac{1}{N} \sum_{n=1}^{N} \mathbf{1}[i \in T_n]$$

• Soft probability — mean router softmax for expert $i$:

$$P_i = \frac{1}{N} \sum_{n=1}^{N} \mathrm{softmax}(\ell_n)_i$$

• Auxiliary loss:

$$\mathcal{L}_\text{aux} = E \sum_{i=1}^E f_i \cdot P_i$$

This term enters as $\mathcal{L} = \mathcal{L}_\text{CE} + \alpha \mathcal{L}_\text{aux}$ with $\alpha \sim 0.01$.

The shape of this loss has a useful property. $f_i$ is non-differentiable (a hard indicator), so gradients flow only through $P_i$. The product $f_i P_i$ reads as: for over-used experts, push the soft probability down. The $f_i$ acts as a fixed weight that scales the gradient on $P_i$, harder for the more-used experts.

The minimum value of the aux loss

Worth deriving once because it's a useful sanity check during training.

Under top-$k$ routing, each token contributes exactly $k$ ones across the experts in its top-$k$, so:

$$\sum_{i=1}^E f_i = k, \qquad \sum_{i=1}^E P_i = 1$$

At perfectly uniform routing $f_i = k/E$ for all $i$, $P_i = 1/E$ for all $i$, and:

$$\mathcal{L}_\text{aux}^{\min} = E \cdot \sum_{i=1}^{E} \frac{k}{E} \cdot \frac{1}{E} = E \cdot E \cdot \frac{k}{E^2} = k$$

So with $k=2$, $E=4$ the floor is exactly 2.0, not 1.0. Switch Transformer's quoted "minimum is 1" assumes top-1. During training, aux settling near $k$ means routing is well-balanced — not that something is broken.

(This caught me off-guard the first time: I'd drawn the "balanced" line at 1.0 in the notebook plot.)

3. Preventing router collapse — noise vs aux loss

Top-$k$ routing has a classic failure mode worth naming. With zero intervention the router exhibits a winner-take-all dynamic:

1. One expert is initialised slightly better by luck.
2. The router favours it on a few tokens.
3. Only that expert receives gradient → only it improves.
4. The router favours it more strongly.
5. The losing experts stay stuck at random init, useless.

Breaking this loop needs either:

• (A) Force tokens to losing experts, so they get gradient anyway, or
• (B) Penalise the router directly for asymmetric distributions.

Shazeer 2017's solution: noisy top-$k$ (mechanism A)

Inject Gaussian noise into the router logits during training, with a learnable per-expert variance:

$$H(x)_i = (W_g\, x)_i + \mathcal{N}(0, 1) \cdot \mathrm{softplus}\bigl((W_\text{noise}\, x)_i\bigr)$$

$W_\text{noise}$ is its own dense projection; softplus keeps the noise std positive. At inference time the noise is zeroed out — routing is deterministic. The mechanism is random exploration: losing experts occasionally bubble into the top-$k$ by chance, receive tokens, get gradient, and improve. Untargeted but reliable — the noise doesn't know which expert needs help; it just shakes the system.

Aux loss (mechanism B) — what we have

The load-balancing loss penalises asymmetric routing directly. Consider a router collapsed onto experts 0 and 1 only:

• $f_0 = f_1 = 1$, $f_2 = f_3 = 0$ (every token's top-2 is $\{0, 1\}$)
• $P_0, P_1 \approx 0.5$, $P_2, P_3 \approx 0$ (router probability mass on the winners)
• $\mathcal{L}_\text{aux} = 4 \cdot (1 \cdot 0.5 + 1 \cdot 0.5 + 0 + 0) = 4.0$

versus the uniform minimum of $\mathcal{L}_\text{aux}^{\min} = k = 2.0$. Collapse is exactly $2\times$ the floor. The gradient

$$\frac{\partial \mathcal{L}_\text{aux}}{\partial \ell_i} \propto f_i \cdot \frac{\partial P_i}{\partial \ell_i}$$

means over-used experts get their logits pushed down with weight $f_i$ — the more over-used, the harder the push. Dead experts get their logits relatively boosted (softmax redistribution). It's a closed-loop controller: the gradient knows which expert to nudge in which direction.

Why aux loss alone is enough at modern scale

	noise injection	aux loss
exploration	random (untargeted)	targeted by $f_i$
feedback	none — fixed perturbation	proportional to imbalance
at convergence	keeps stirring forever	self-quiets at the floor
failure mode	over-perturbs late training	needs stable $f_i$ estimate

Three reasons modern MoE drops noisy gating in favour of just the aux loss:

1. $f_i$ is estimated precisely at large batch sizes. With $N = B \cdot T$ tokens per batch ($\sim 10^4$ here, $\sim 10^6$+ in production), the per-expert usage estimate has standard error well under a percent. The aux gradient is a sharp signal, not a noisy one.
2. Targeted beats random. Noise perturbs all experts equally — including the ones already routed to correctly. Aux loss only pushes against the actually-misbehaving ones. No wasted gradient on already-good routing.
3. Self-stopping. Once routing balances, $\nabla \mathcal{L}_\text{aux} \to 0$ and the loss stops perturbing the router. Noise keeps stirring at convergence, which becomes a regularisation cost rather than a benefit.

This is the trajectory of MoE research: Shazeer 2017 used noise + aux, Switch Transformer 2021 dropped to aux only, DeepSeek-V3 2024 simplified further to expert biases only (a tiny learnable per-expert offset added to the router logit, no aux loss at all). Each step trades an explicit mechanism for stronger inductive bias toward what works at scale.

When noisy gating still helps

At small batch and small $E$ the $f_i$ estimate is too noisy for the aux loss alone to give a sharp gradient. Initialisation luck can stall recovery for many steps before the aux loss can pull the routing back. Mechanical exploration (noise) gives a deterministic exploration floor — losing experts get token assignments even when the gradient signal is weak.

In this notebook ($B \cdot T = 8192$, $E = 4$) the aux loss is comfortably sufficient — the aux trace pins to 2.0 within ~50 steps and stays. But that's a function of scale, not a universal truth.

4. Composing it in Flax

The full model is short:

class LSTMMoE(nn.Module):
    vocab_size: int
    embed_dim:  int = 128
    hidden_dim: int = 256
    num_experts: int = 4
    top_k: int = 2
    expert_hidden: int = 512

    @nn.compact
    def __call__(self, x):
        x = nn.Embed(self.vocab_size, self.embed_dim)(x)
        x = LSTMLayer(hidden_dim=self.hidden_dim, name='lstm1')(x)
        x = MoEBlock(
            num_experts=self.num_experts, top_k=self.top_k,
            expert_hidden=self.expert_hidden, out_dim=self.hidden_dim,
            name='moe',
        )(x)
        x = LSTMLayer(hidden_dim=self.hidden_dim, name='lstm2')(x)
        return nn.Dense(self.vocab_size, name='head')(x)

About 2M parameters: the LSTMs hold a bit under half, the four experts the other half, embed and head are negligible.

The sow / mutable trick

The aux loss is computed inside the MoE module but needs to be added to the main loss in the training step. Flax handles this with sow — a side channel that lives in a separate variable collection:

# inside MoEBlock.__call__
self.sow('intermediates', 'aux_loss', aux_loss)

# inside the training step
logits, mut = state.apply_fn(
    {'params': params}, x,
    mutable=['intermediates'],
)
aux = mut['intermediates']['moe']['aux_loss'][0]
loss = ce + AUX_COEF * aux

sow is the canonical way to surface things from deep in a Module tree (intermediate activations for visualisation, statistics, here an extra loss) without breaking the pure-function contract.

5. Training loop

jit-ed train_step, single-batch sampling from the raw byte stream, warmup-cosine LR, AdamW, gradient clipping. The full step:

@jax.jit
def train_step(state, x, y):
    def loss_fn(params):
        logits, mut = state.apply_fn(
            {'params': params}, x,
            mutable=['intermediates'],
        )
        ce  = optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()
        aux = mut['intermediates']['moe']['aux_loss'][0]
        loss = ce + AUX_COEF * aux
        return loss, (ce, aux, logits)
    (loss, (ce, aux, logits)), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params)
    return state.apply_gradients(grads=grads), {'loss': loss, 'ce': ce, 'aux': aux}

1500 steps. Per-step about 0.2s on a Colab T4 after JIT warms up.

6. Results

Tracking three numbers across training:

• train_ce / val_ce — main next-token loss, on training and held-out splits
• aux — should hover near the theoretical floor of $k = 2$ once routing settles
• val_acc — top-1 next-character accuracy

1500 steps in 122.3s on a free Colab T4 (about 80ms / step after JIT warm-up):

step	train_ce	aux	val_ce	val_acc
1	4.18	2.01	4.18	1.0%
100	2.30	2.00	2.22	36.5%
300	1.64	1.98	1.76	47.6%
600	1.39	2.00	1.59	52.8%
1000	1.27	2.00	1.54	53.8%
1500	1.16	2.00	1.49	55.2%

Three things to read:

1. aux snaps to ~2.00 within the first ~50 steps and stays there. Routing is balanced essentially from initialization onward — the gradient through $P_i$ is strong enough to keep $\mathcal{L}_\text{aux}$ pinned to the floor. The aux loss is doing its job, quietly.
2. ce drops cleanly without over-training. Final train-val gap is 0.32 nats — mild overfitting beginning, not yet hurting.
3. Final val_acc 55.2% — in the same neighbourhood as a vanilla LSTM at this compute budget. The MoE is not the win here; the point of this run is the mechanism, not the headline number. Real returns come from sparse dispatch + many more experts.

Generation sample

After training, sampling from the model with prompt='ROMEO:\n', temperature 0.8:

ROMEO:
Thus he was lie in hazarding, their loves
Than they the house of Bolingpoodest that dost would ask those with wind hours
To fright the thoroun dear death,
And when thy wife first as his oratory,
Which best your honour strike it be my wounds,
That show'd for that hand in the noble strange of my love
That now be rought the which in the hands:
Your facting stolen abroad, in sudden and this heavy abse

Iambic-ish line lengths, character names, line breaks, even a fake Shakespearean word ("Bolingpoodest"). The model has learned the shape of the corpus — character distribution, line structure, capitalisation — without any semantic coherence. Exactly what you'd expect from 1500 steps of char-LM at this scale.

7. What MoE specialisation looks like

The harder question — is the MoE doing anything the FFN couldn't? Numerically, aux near $k$ only tells you usage is balanced; it doesn't tell you experts specialise.

Measured at the end of training, the fraction of tokens for which each expert appears in the top-2:

expert	usage
0	50.8%
1	49.6%
2	49.8%
3	49.8%
ideal (top-2/4)	50.0%

All four within 0.5pp of perfect. That answers the load-balancing question but leaves the specialisation one open. The notebook plots a per-sequence routing heatmap ([T, E] grid, colour = gating weight) to look at the shape of the assignment over time. Patterns to look for:

• Whitespace / punctuation experts: in a char-LM, certain experts grab spaces, newlines, commas — easy "boundary detection" specialisation.
• Vowel / consonant experts: less common but does show up if the model commits hard enough.
• Position-dependent splits: experts that activate on line starts vs line interiors.

The signal is strongest at the end of training when the routing distribution stops moving. At only 1500 steps and four experts the patterns are subtle; pushing $E$ up and training longer is where specialisation gets visually obvious.

8. Forcing specialisation: a 2-domain experiment

The Shakespeare-only setup leaves an awkward gap: aux near $k$ proves balance but doesn't prove specialisation. At 4 experts on a single-domain corpus the model has no strong reason to assign experts to a clear subset of inputs — it can balance load by routing each token to whichever pair happens to be slightly preferred.

A sharper test: train on two structurally different corpora with a shared character vocabulary, then ask whether expert usage diverges by domain.

Corpus:

• Shakespeare — Tiny Shakespeare (~1MB).
• Python source — a handful of CPython stdlib files (inspect, typing, argparse, dataclasses, json/{decoder,encoder}, asyncio/base_events) concatenated.

Both encoded with the same char vocabulary (the union of characters across both), so the router can't trivially partition by character range — only by structural patterns (indentation, punctuation, casing, line shape).

Batching is per-item random: each of the 64 items in a training batch comes from Shakespeare or Python with 50/50 probability. The aux loss looks at the full batch, so global balance still pins $\mathcal{L}_\text{aux}$ near $k = 2$. But that's globally balanced — it's silent about per-domain balance.

Evaluation: run the trained model on pure Shakespeare batches and pure Python batches separately, gather routing weights, and compute per-expert usage on each domain. Asymmetry across the two columns is the specialisation signal.

Training summary

2000 steps in 147.3s on a free T4. Validation tracked separately for the two domains:

step	train_ce	aux	shake val_ce	shake val_acc	py val_ce	py val_acc
1	4.57	2.01	4.57	0.9%	4.57	0.6%
400	1.50	1.99	1.78	47.0%	1.47	61.1%
1000	1.12	1.99	1.56	53.3%	1.25	67.8%
1500	0.97	2.01	1.54	54.8%	1.10	71.2%
2000	0.94	2.00	1.48	55.6%	1.10	71.8%

Python is consistently easier than Shakespeare — about 0.4 nats lower CE and 16pp higher accuracy throughout. Char-LMs on code corpora benefit hugely from the mechanical repetition ( indentation, def/return/self tokens, ): line endings) that Shakespeare doesn't have.

Expert usage by domain

expert	Shakespeare	Python	asymmetry
0	50.9%	46.9%	+4.0pp
1	42.5%	48.1%	−5.6pp
2	44.4%	60.1%	−15.7pp
3	62.2%	44.9%	+17.3pp

Partial specialisation emerged. Two experts (0 and 1) stayed roughly balanced across domains — they act like shared general-purpose FFNs. The other two (2 and 3) clearly committed:

• Expert 3 → Shakespeare (+17.3pp asymmetry, fires on 62% of Shakespeare tokens vs 45% of Python)
• Expert 2 → Python (−15.7pp, fires on 60% of Python tokens vs 44% of Shakespeare)

So with $k=2$ and $E=4$, the model spontaneously settled into a "one specialist per domain plus two generalists" routing strategy. The aux loss kept each expert near 50% globally; the per-domain split is what specialised within that constraint.

This is genuinely a real signal — not the clean 80/20 partition you'd see at large $E$, but unambiguously not symmetric either.

Generations switch register cleanly

With prompt='ROMEO:\n', temperature 0.8:

ROMEO:
Must have been sing me to the own with far
and more than the supplied by vain of post-swift and chose as sisters shall gaid.
What know not out of love, made it may take with whise easy thou didst present
Which you for a place wretch, what hither,
Of your market-perform in the colours,

With prompt='def factorial(n):\n if':

def factorial(n):
    if not hasattr(self, owner=None):
    """Special typing can from optional arguments to arg_strings and the signature are generated all the
            # that the parser states with the stack like a data
        # and extras and static all help with dataclass or Function.

    Some other function disjo

The same model, same parameters, no prompt-engineering — only the first few characters differ. Output 1 is Shakespearean register (line lengths, archaic vocabulary, no symbols). Output 2 is Python register: indentation, self, docstring, # comments, no period sentences. The MoE has internalised that these are two distinct distributions and switches routing accordingly.

Notebook: 03-lstm-moe-domain-split.

What this tells us

At $E = 4$ with aux loss only, specialisation is real but limited: half the experts pick a side, the other half stay shared. This is exactly the trade-off the aux loss imposes — it enforces global balance, which means an expert can specialise only as long as the other experts compensate to keep the global usage near $k/E$. With only 4 experts to play with, two-and-two is what fits.

Pushing $E$ up (say 8 or 16) would loosen this constraint substantially — there'd be more "budget" for specialisation. That's where the Shazeer 2017 results showing strong domain specialisation came from: $E$ in the hundreds, lots of room for fine-grained partition.

9. Pushing to four domains — what specialisation actually looks like

Two domains gave a partial 1+1+2 split. Four domains tighten the geometry: with $E=4, k=2$, the aux loss enforces global $f_i \approx 0.5$, but each domain only accounts for $1/D = 0.25$ of the tokens. Pure 1-to-1 specialisation is now arithmetically impossible — each expert must either pair up with another or stay shared. Which pair (or which axis) the model picks is the question.

Setup

Same architecture as before. Four corpora, shared character vocabulary (~2,160 chars once Japanese joins):

• Shakespeare — Tiny Shakespeare (classical English drama)
• Python — CPython stdlib (inspect, typing, argparse, dataclasses, json/*, asyncio/base_events)
• Kokoro (こころ) — Sōseki's novel, from Aozora-bunko (Japanese literary prose)
• Dario — Dario Amodei's essay The Adolescence of Technology (modern English long-form)

3000 steps, 25% domain weight per batch. Notebook: 04-lstm-moe-4domain-clustering.

The 4 × 4 usage matrix

After training, fraction of tokens for which each expert appears in the top-2:

	expert 0	expert 1	expert 2	expert 3
Shakespeare	54.1%	44.6%	45.7%	55.7%
Python	66.8%	39.0%	26.2%	68.1%
Kokoro	39.4%	55.3%	72.4%	32.9%
Dario	48.8%	48.6%	41.7%	60.9%

The structure:

• Expert 2 → Kokoro specialist: 72% on Kokoro vs 26% on Python — a 46pp swing. The strongest single-axis specialist of the run.
• Expert 3 → ASCII specialist (anti-Japanese): 56–68% on the three English/code domains, only 33% on Kokoro.
• Expert 0 — mild English/code lean (67% Python).
• Expert 1 — generalist with a small Kokoro tilt.

The axis the model picked is the Unicode range — Japanese vs ASCII — not any of the more interesting sub-axes (prose vs code, classical vs modern, drama vs novel). The strongest structural signal dominated; the aux loss didn't reward fine-grained discrimination within the ASCII cluster.

Generations with the right sampling

The first run used short prompts and temperature=0.8, and generations bled across domains within a few characters. That was a sampling artefact, not a routing failure. With long prompts (~150 chars, near-saturating the 128-token context) and temperature=0.3, the same model produces:

Shakespeare (continuing "ROMEO: But, soft! what light through yonder window breaks…"):

A great to the heavens of his head the soul of the state to him,
That thou shalt not be the honour and the state.

First Murderer:
The world the state of the world and the man that say that the duke
of when the manner and the state.

Stays fully Shakespearean — character names, verse rhythm, archaic vocabulary. Loops on phrases ("the world", "the state") but never falls out of register.

Kokoro (continuing a passage from the novel):

死んだら、お嬢さんは自分の室を離れて、その時の私の知った所にもかかっているという言葉を、それによく書き出した。
「そういう事はない」
　私はそれを忘れてしまった。

Pure Japanese; correct particles; dialogue quote marks 「」 preserved; Sōseki's typical phrase rhythms ("私はその時分から…") emerge naturally. Semantic coherence is loose — this is char-LM at 3000 steps — but the register is locked.

Dario (continuing "The pace of progress in AI…"):

…the questions we must answer about their deployment, governance, and societal impact
that the world with powerful AI will be only to the people to avoid a technology…
…a country of geniuses in a datacenter could be considered a country of geniuses…
…or does not a dataclass in the repr(self.__doc__
        except AttributeError:
            return self.__origin__

Stays in English prose for most of the generation — even using domain-appropriate vocabulary ("country of geniuses in a datacenter" is recognisably Amodei phrasing). Then, near the end, the word dataclass appears in the model's own output, the router flips to the ASCII / Python expert pair, and the rest of the generation is Python.

This is striking: a single trigger token is enough to switch experts. The router has learned that dataclass belongs to a different distribution than the surrounding English prose. The flip is consistent with how production MoE behaves at scale — Mixtral's reported per-domain expert specialisation is the same mechanism, just observable at much larger $E$.

Python (continuing a def parse_arguments(): ... snippet):

def parse_arguments():
        parser = argparse.ArgumentParser(...)
        parser.add_argument(...)
                                                                       [many spaces]

A different failure mode — whitespace mode collapse. After the function definition the highest-probability next character is whitespace (indentation), and at temperature=0.3 the argmax keeps choosing spaces. This isn't an MoE problem; it's a generic greedy/low-temperature decoding pathology that top_p (nucleus) sampling or a repetition penalty fixes.

Three lessons

1. Trigger words flip experts. A single token can shift the router's top-$k$ choice and cascade into a different generation mode. This is specialisation actually working — the same phenomenon production MoE displays, scaled down.
2. Training and inference assumptions must match. This model was trained on pure-domain 128-character sequences. It never saw a Shakespeare → Python transition during training. So at inference, once routing flips, there's no learned recovery path. Robust domain locking at inference requires training on cross-domain sequences too.
3. Sampling quality is orthogonal to the MoE. Domain locking can succeed (good routing) and still produce bad output (whitespace loops, repetition) if the decoder is naive. top_p, repetition penalty, and minimum new tokens all apply.

What axis came out, and what didn't

Predicted axis	Emerged?
Japanese vs ASCII (Unicode range)	✓ Strong (Expert 2 vs Expert 3)
Prose vs code (English)	✗ No clear partition
Classical vs modern (English)	✗ No clear partition
Single 1-to-1 mapping	✗ Geometrically impossible at $E/k = 2$

The model picks the strongest structural axis available and leaves the rest as shared FFN capacity. Pushing $E$ to 8 or 16 would loosen the geometric constraint and likely reveal finer axes — but actually rewarding that finer specialisation needs training-time tricks (capacity factors, expert dropout) on top of aux loss alone.

Reflection

A few things stand out from writing this end-to-end:

• Sparse routing is a small mechanism with a sharp gradient story. The aux loss is two lines of code and a clean derivation, and it's what prevents the whole thing from collapsing to a single expert.
• Dense compute is fine for understanding. The actual "outrageous parameter count" win of MoE comes from sparse dispatch — only forwarding tokens through their assigned experts. That's a separate notebook.
• Flax sow + mutable= is the right abstraction for auxiliary losses. It keeps the Module API clean (a single output) while letting the training step pluck out whatever intermediates it needs. The alternative — returning a tuple from every sub-Module — would force every intermediate level to know about the aux loss.
• The right aux minimum is $k$, not 1. Doesn't matter for training but it matters for plot annotations and for reading the numbers.

What's next

• Sparse dispatch: gather tokens per expert, run each expert on a smaller batch, scatter back. Should be substantially faster at larger $E$.
• Bigger experts: with sparse dispatch in place, scaling expert count without scaling compute per token becomes the point of the exercise.
• Replace LSTM with Transformer: Mixtral-style. Same routing mechanism, different sequence backbone.
• Auxiliary loss alternatives: expert dropout, dispatch capacity factor, DeepSeek-V3's bias-only routing — small variations with measurable training stability differences.

---

References

• Noam Shazeer et al. "Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer" arXiv:1701.06538, 2017.
• William Fedus, Barret Zoph, Noam Shazeer. "Switch Transformer: Scaling to Trillion Parameter Models" arXiv:2101.03961, 2021.
• Albert Q. Jiang et al. "Mixtral of Experts" arXiv:2401.04088, 2024.
• Notebooks: 02-lstm-moe (base) · 03-lstm-moe-domain-split (2-domain) · 04-lstm-moe-4domain-clustering (4-domain). Repo: createcentury/jax-flax-optax-lab.

---

Created: 2026-05-18 / Updated: 2026-05-18