For now, I ran a quick experiment in Colab:
I think there are two useful things to separate in the current results:
- The final router statistics do not look like an obvious global expert-collapse case, but final load/entropy is not the same thing as routing stability over training.
- The current Top-1 vs Top-4 comparison changes not only
K, but also the router weighting / aggregation mode (straight_through vs probability). In the current implementation, that is a real behavioral difference, not just a logging detail.
I did a small matched check on the current SplitMoE repo, first with a deterministic semantics test and then with the exact Top-4 model shape on a small real-data mixture. At fixed K=4, changing only the router mode was enough to change short-horizon convergence. So I would probably treat router mode as a control variable before interpreting the Top-1/Top-4 change as a pure K effect.
At the same time, the experiment did not show that router mode explains the architecture-ranking change. In my short real-data run, switching Top-4 to straight-through did not restore a Top-1-like advantage for Split; if anything, the one-seed interaction went weakly the other way. So I think the useful conclusion is narrower: there is a confound worth separating, but it is not currently an explanation of the result.
For the routing-stability question specifically, my default path would be:
final load balance / collapse → temporal route churn → boundary sensitivity/margin
as three separate measurements. If you run the model again, logging the selected expert IDs for a small fixed validation set at evaluation steps would probably give more information than adding another large benchmark.
What I tested in Colab
I pinned the repo to commit cfbb4a08d912dc938cd65b0a530d7cffc5f1c61f.
1. Router-mode semantics
The current implementation makes the three router modes meaningfully different in model.py.
For selected experts:
probability uses the normalized selected router probabilities as the forward aggregation weights;straight_through makes the forward selected weights uniform (1/K), while retaining a gradient path through the router probabilities;none also uses uniform 1/K forward weights, but removes that LM-loss gradient path.
A deterministic check gave exactly the expected result:
straight_through vs none: identical forward logits for fixed selected routes;straight_through vs probability: different forward logits;- router LM-gradient under
none: zero;
- router LM-gradient under
straight_through: non-zero.
So router_weight_mode is not just an optimizer implementation detail: probability changes the actual forward mixture.
This matters because the published Top-1 setup uses straight-through routing, e.g. paper_allmoe_standard_1024.json, while the Top-4 configurations use probability weighting, e.g. kaggle_allmoe_16e_top4_standard_seed1337.json.
2. Small synthetic fixed-Top-4 gate
Before spending GPU time on the full shape, I tried:
- Standard / probability
- Standard / straight-through
- Split / probability
- Split / straight-through
with identical initialization and batch ordering inside the matched pairs.
Probability weighting had a repeatable short-horizon loss advantage in that toy setup, but the architecture × router-mode interaction was weak/inconsistent. Straight-through also increased the Top-K boundary margin much more than it changed actual route-set persistence.
That already suggested two things:
- mode can affect finite-horizon training even at fixed
K;
- a large routing-logit margin should not automatically be called temporal stability.
3. Paper-like real-data gate
I then used the exact current Top-4 model shape, not a reduced toy shape:
|
Standard |
Equal-active Split |
| layers |
8 |
8 |
d_model |
512 |
512 |
| experts |
16 |
16 |
K |
4 |
4 |
| Standard expert width |
1024 |
— |
| shared width |
— |
256 |
| private width |
— |
960 |
| active parameters/token |
84,091,392 |
84,091,392 |
| total parameters |
235,086,336 |
225,649,152 |
The exact shape ran on an L4 without fallback.
For data, I used small balanced block pools from the same four source families listed in the repo’s data_sources.example.json:
This was deliberately small: one seed, 140 optimizer steps, and a much smaller token pool than the paper runs. So I would call it a mechanism/control gate, not a reproduction.
Final fixed-probe LM loss:
| architecture |
probability |
straight-through |
ST − probability |
| Standard |
7.561074 |
7.587149 |
+0.026075 |
| Equal-active Split |
7.570499 |
7.605377 |
+0.034878 |
The difference appeared during training rather than being an initialization offset.
For Standard, ST - probability loss moved roughly:
step 0: -0.0008
step 40: +0.0007
step 80: +0.0219
step 100: +0.0326
step 140: +0.0261
For Split:
step 0: -0.0003
step 40: +0.0024
step 80: +0.0205
step 100: +0.0342
step 140: +0.0349
So on this short real-data run, router weighting materially changed finite-horizon optimization even with K=4, architecture, initialization and batch schedule controlled.
But the architecture × mode loss interaction was only:
(Split mode effect) - (Standard mode effect) = +0.008803
and this is one short seed. I would not use it to claim a robust interaction.
More importantly, it does not support the simple hypothesis that “Top-4 Split only lost because Top-4 used probability weighting.” Switching to straight-through did not rescue Split in this gate.
4. Routing dynamics
I also tracked the selected Top-4 sets on a fixed probe.
Final mean K vs K+1 router-logit margin:
|
probability |
straight-through |
| Standard |
0.05879 |
0.07012 |
| Split |
0.04911 |
0.04867 |
So straight-through increased the margin for Standard, but essentially not for Split.
Mean step-to-step Top-4-set overlap:
|
probability |
straight-through |
| Standard |
0.55563 |
0.55872 |
| Split |
0.53455 |
0.52196 |
And only for the late 120 → 140 interval:
|
probability |
straight-through |
| Standard |
0.8244 |
0.8387 |
| Split |
0.8062 |
0.7538 |
So the margin and actual temporal persistence can clearly move differently.
Another useful detail: by step 140, the exact same Top-4 set as initialization was almost nonexistent (well below 1% for every cell), while late step-to-step overlap had already risen to roughly 0.75–0.84.
That suggests a useful distinction between:
- global reorganization from initialization, and
- local late-training stability.
Again, this is only a short run, but I think it is enough to say that final load balance, router margin and temporal assignment stability should be measured separately.
What I would call routing stability here
For this project I would probably split “routing stability” into three observables.
A. Load balance / collapse
Things like:
- per-expert token fraction;
- router entropy;
- dead experts;
- dominant experts.
The committed router_summary.csv is useful for this.
From those final summaries, I do not see an obvious story where the router simply collapses globally onto one expert. There are layer/seed-specific imbalances, but that is a different claim from global collapse.
What this does not tell us is whether the same input kept changing experts during training.
B. Temporal stability / churn
This is much closer to the “routing fluctuation” definition used in StableMoE: the same input may change its target expert while the model is training.
That can be measured directly with a fixed evaluation set:
- exact route-set persistence;
- Jaccard/overlap of selected expert sets;
- number of route changes;
- last step at which the route changed;
- layerwise churn.
I think this would be a particularly informative metric for SplitMoE because it separates:
“the final router looks reasonably balanced”
from
“the routed private branches actually acquired stable token assignments.”
C. Boundary sensitivity
A cheap proxy is the router-logit margin:
margin = logit(Kth selected expert) - logit((K+1)th expert)
A tiny margin means the Top-K boundary is easy to flip under small router changes.
But I would call this sensitivity / fragility, not stability itself.
The Colab checks are one concrete reason: changing router mode could enlarge the margin without comparably increasing actual step-to-step route persistence.
Minimal instrumentation
If you train again, I do not think you need to save every full checkpoint just to study this.
At each normal validation point, for a small fixed sample, something as compact as:
step
layer
example/token ID
selected expert IDs
Kth router logit
(K+1)th router logit
would be enough to reconstruct most of the useful diagnostics later:
- load;
- churn;
- exact persistence;
- overlap;
- last-change time;
- boundary margin.
That seems like a relatively cheap way to make the “routing stability” claim much more concrete.
One extra caution: stable routing is not automatically good routing. OpenMoE, for example, found routing patterns that become established quite early while also identifying limitations in the semantic/contextual quality of that specialization. So I would keep “stable”, “balanced” and “useful” as three separate questions.
Why I would separate K from router aggregation
The Top-1/Top-4 change is interesting, but I would currently describe it as a regime change, rather than attributing it to K alone.
At least these axes move:
Top-K / selected expert count
router weighting / aggregation semantics
output-scaling choices
shared/private capacity allocation
The most important one to isolate first seems to be aggregation, because the current implementation makes probability and straight_through genuinely different forward functions.
There is also some recent methodological support for making exactly this separation. Beyond Routing: Decoupling Expert Dispatch and Aggregation in Sparse Mixture-of-Experts keeps the selected expert IDs fixed and changes the within-set aggregation, finding meaningful effects from aggregation alone. It is a recent preprint, so I would treat it as supporting evidence rather than settled consensus, but the distinction maps unusually well onto this case.
A simple interpretation tree would be:
Fixed Top-4 probability vs straight-through:
If Standard and Split keep roughly the same ordering
-> the K/shared-capacity interpretation becomes cleaner.
If the architecture ordering changes materially
-> there is a meaningful architecture × aggregation interaction.
If mainly convergence speed changes
-> router mode is more likely a finite-horizon/training-dynamics variable
in this regime than the explanation of the architecture ordering.
The short Colab checks make the third branch look plausible, but they are nowhere near sufficient to settle the 6,500-step multi-seed result.
If you already wanted one extra confirmatory run, a longer multi-seed fixed-Top-4 probability-vs-ST control seems much higher-information than immediately launching a broad new scaling sweep.
But I would not call that experiment necessary for the existing results to be useful—the current Top-1 and Top-4 results are already useful as observations of two different regimes.
Scaling: I would split this into several axes
I would also avoid reducing “does SplitMoE scale?” to one model-size axis.
For routed models, useful axes include at least:
| axis |
question |
| total parameters |
how much stored capacity is available? |
| activated parameters/token |
how much parameter capacity is used per token? |
N experts |
how many routed choices exist? |
active K |
how many are composed per token? |
| shared fraction |
how much computation/capacity is always reusable? |
| width / depth |
does the preferred allocation move with model geometry? |
| training tokens / horizon |
does a short-run optimum survive longer training? |
| FLOPs |
idealized arithmetic cost |
| throughput / latency |
actual systems behavior |
Unified Scaling Laws for Routed Language Models is useful background here because it explicitly treats parameter count and computation as independent scaling axes for routing networks rather than collapsing them into one notion of “size”.
For SplitMoE specifically, one way I would frame the design question is:
How much reusable shared capacity can be introduced before it starts costing too much conditional-combination flexibility?
That trade-off does not have to be monotonic.
There is useful evidence in both directions.
DeepSeekMoE deliberately combines fine-grained expert segmentation with shared expert isolation, motivated by putting common information into shared experts and reducing redundancy among routed experts.
But the OLMoE study gives an informative counterexample rather than a contradiction. In the OLMoE paper, the shared-expert ablation matches total parameters, active parameters and FLOPs. The shared-expert version was slightly worse in that tested setting, and the authors point out that replacing a routed expert with an always-active one sharply reduces the number of possible routed-expert combinations.
That makes me think the interesting variable may be less:
“shared experts: yes/no”
and more:
“where is the best point between reusable capacity and conditional combinatorics for this N, K, depth and training budget?”
That also seems consistent with your own result that the preferred Split/shared configuration changes between Top-1 and Top-4.
One additional branch that might become interesting later is layerwise allocation. The current shared/private activation-norm ratios are strongly layer-dependent rather than uniform across depth. I would only treat this as a future hypothesis, but it suggests that a single global shared fraction may eventually be more restrictive than necessary.
Some related work that seems useful for positioning this
I would not use these as a novelty verdict; they are mostly useful for locating the design space.
DeepSeekMoE
DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models
Probably the closest obvious reference point because it explicitly separates:
- fine-grained routed expert segmentation;
- always-active shared experts.
The similarity is useful, but I would still keep the implementations conceptually separate: SplitMoE is changing the shared/private FFN-width allocation inside this experimental architecture rather than simply reproducing DeepSeek’s expert-bank construction.
OLMoE
OLMoE: Open Mixture-of-Experts Language Models
The shared-expert appendix is particularly useful because it shows why there may not be a universally optimal answer. Their matched shared-expert test was slightly worse, while the authors argue that shared capacity also reduces routed-expert combination flexibility.
That is a useful counterweight to any assumption that more shared capacity should monotonically improve redundancy.
MoEBERT / MPOE
Two older neighboring examples:
They are structurally different from SplitMoE, but they are useful prior art around decomposing/transforming expert FFNs and sharing parts of expert parameterization rather than assuming every expert must be fully independent.
MPOE is especially relevant conceptually: it shares a central tensor between experts while leaving auxiliary expert-specific tensors, explicitly targeting parameter redundancy.
StableMoE
StableMoE
The most directly useful reference for the routing-stability question because it treats same-input route changes during training as the quantity of interest.
Dispatch vs aggregation
Beyond Routing
Recent preprint, but directly relevant to the Top-1/Top-4 interpretation because it shows that expert selection and within-selected-set weighting need not be the same optimization problem.
If you want one or two extra checks
I would branch it depending on what artifacts or training budget you already have rather than treating all of this as a checklist.
If you retained periodic checkpoints:
-> replay a small fixed validation sample
-> measure route churn / last-change directly
If you only retained final checkpoints:
-> measure Top-K boundary margins and alternative routes
-> call this sensitivity, not temporal stability
If another training run is already planned:
-> add compact fixed-input route fingerprints at eval steps
-> no need to retain every full checkpoint
If one causal control is affordable:
-> fixed Top-4
-> probability vs straight-through
-> Standard + equal-active Split
-> preferably multiple seeds / longer horizon
If scaling is the main goal:
-> first decide whether the claim is about N/K,
shared fraction, model size, training horizon,
stored/active parameters, or actual system throughput
If none of those additional runs are attractive, I think there is still a clean result to report:
- the final router statistics do not suggest a trivial global-collapse explanation;
- temporal routing stability remains a separate, currently under-measured quantity;
- Top-1 vs Top-4 is currently a multi-axis comparison, not a pure
K intervention;
- the existing results can therefore be described as regime-specific evidence without requiring a stronger causal explanation.
One systems caveat about 'scaling'
I would keep systems scaling separate from model-quality scaling.
A reduction in stored parameters, active parameters or nominal FLOPs does not automatically imply the same reduction in wall-clock latency or training time for an MoE implementation.
MegaBlocks is a good example of why: dynamic expert workloads interact with padding, token dropping, sparse kernel structure and GPU utilization, and specialized kernels can substantially change the realized performance.
So I would read the current reference PyTorch throughput as an implementation measurement, while keeping any broader systems-efficiency claim separate until the comparison is done with an optimized MoE backend / grouped or fused expert execution.
That does not weaken the parameter-efficiency result; it just keeps parameter efficiency, compute efficiency, and systems efficiency from becoming the same claim by accident.
Overall, I think the Top-1 result is still interesting as a local empirical result, and the Top-4 result is useful rather than a failure: it exposes a boundary where the preferred shared/private allocation changes.
The cleanest conceptual decomposition I can see now is:
shared/private capacity allocation
×
dispatch (which experts / K)
×
aggregation (how selected experts are weighted)
×
temporal routing stability
Separating those four pieces would make it much easier to say where SplitMoE helps, without requiring a change to the core idea.
GitHub Repo: https://github.com/Priyanshu-5257/SplitMoE