Stepped MoE: Segment-Level Routing with Configurable Inference Complexity

• Flexible Stepped-MoE targets edge-serving constraints where token-level expert swapping can turn inference into an I/O-bound workload.Source: §1 · pp. 1–2
• It reduces IFP’s pruning search space by selecting MoE experts instead of individual feed-forward weight rows and columns.Source: §1 · p. 2
• A unified routing mechanism uses part of the model as the pruner, avoiding a separate pruner model.Source: §1 · p. 2
• Experts are reloaded every S inference steps, trading off representational capacity against the latency cost of switching too frequently.Source: §1 · p. 2
• Users can select configurable active-parameter budgets from one trained model to match task and hardware constraints.Source: §1 · p. 2
• A sparsity-aware embedding signals active capacity, while the mask generator ranks experts across layers and selects a variable top k.Source: §1 · p. 2
• Stepped-MoE is much faster at inference than Vanilla MoE and is reported more accurate than dense counterparts.Source: §1 · p. 3

2 Methodology

• The methodology targets local LLMs that flexibly allocate compute and load parameters efficiently when device DRAM is limited.Source: §2 · p. 3

2.1 MoE background

• Per-token MoE routing changes active experts at every decoding step, prompting weight swaps that can limit throughput on DRAM-constrained devices.Source: §2.1 · p. 3

2.2 IFP background

• IFP keeps one query’s selected weights resident, but cannot adapt experts as context evolves and searches over individual hidden dimensions.Source: §2.2 · p. 3

2.3 Model architecture

• The dense decoder backbone selects experts for successive layers at once, allowing later-layer weights to load while earlier layers execute.Source: §2.3 · p. 4

2.4.1 Pre-training

• Pre-training selects a segment’s experts from its first embedding and replicates batch elements so each sample trains across supported active sizes.Source: §2.4.1 · p. 4

2.4.2 Loss formulation

• The loss combines examples across sparsity levels; scaling penalizes denser configurations more to compensate for their less-frequent expert updates.Source: §2.4.2 · pp. 4–5

2.4.3 Supervised Fine-tuning

• Supervised fine-tuning concatenates short examples, computes loss on target tokens, and selects experts using the last prompt token and segment starts.Source: §2.4.3 · p. 5

2.5 Inference

• At inference, the dense part stays in DRAM while experts selected for the user’s active size remain valid for S generation steps before refresh.Source: §2.5 · p. 5

3.1 Pre-training

• Pre-training uses 1 trillion web-scraped tokens, an 8192-token context, and batch size 2048; comparisons include dense, static S-MoE, flexible S-MoE, and a 3B-active Vanilla MoE upper bound.The architecture has 56 layers, including 12 dense layers that generate routing masks for the remaining 44 MoE layers.Source: §3.1 · p. 5
• Pre-training evaluations cover core language understanding and reasoning, MMLU, GSM8K mathematics, and HumanEval-python coding, using zero- or few-shot metrics.Core Text includes seven zero-shot tasks and two one-shot question-answering tasks.Source: §3.1 · pp. 5–6
• Across most sizes, flexible S-MoE outperforms dense counterparts on almost all benchmarks but remains below Vanilla MoE; gains are more pronounced on knowledge-based tasks such as MMLU.The authors hypothesize that 12 dense overhead layers, or 21% of layers at 1B scale, leave insufficient capacity for MoE-specific learning.Source: §3.1 · p. 6, Table 1 · p. 6

3.2 Pre-training Ablations

• Batch replication addresses the undertraining risk of flexible S-MoE; mask-weighted cross-entropy improves larger models but reduces small-active-size performance, while active-size embeddings help bridge the gap.The authors use batch replication so ablation models see the same data, estimating that training all extractable flexible models costs roughly as much as training four dense models.Source: §3.2 · p. 6, Table 2 · p. 7
• A segment length of 16 outperforms 32 across benchmarks, while shorter segments increase routing overhead and longer ones may limit adaptation to evolving context.The authors treat segment length as configurable because its optimum may depend on sequence length, model scale, and deployment constraints.Source: §3.2 · pp. 6–7, Table 3 · p. 7

3.3 Supervised fine-tuning

• Supervised fine-tuning uses batch replication for flexible S-MoE and evaluates instruction following, coding, mathematics, and tool use; Table 4 compares dense, static S-MoE, and flexible S-MoE models.All S-MoE models in the comparison have 12B total parameters.Source: §3.3 · pp. 7–8, Table 4 · p. 8

3.4 Inference profiling

• Stepped-MoE closely tracks 3B dense-model latency across cache hit rates and sequence lengths, stays faster than 6B dense models, and matches 4B dense latency while being more accurate.Vanilla MoE latency decreases as cache hit rate rises from 0 to 0.6 but remains exponentially higher than dense baselines.Source: §3.4 · p. 8, Figure 3 · p. 8

Sparse activation

• MoE routes tokens to specialist experts to increase model capacity without proportional compute, but per-token routing requires expert swapping and fixed sparsity limits platform flexibility.The paper adopts always-active shared experts alongside routed experts, following DeepSeekMoE [Dai et al., 2024].Source: §4 · pp. 8–9

Instruction-based pruning

• IFP [Hou et al., 2025] selects a query-specific parameter subset that stays resident during generation, avoiding repeated swaps but relying on a separate pruner and limiting pre-training from scratch.The cited work's fixed query-level expert set can also constrain adaptation as context changes.Source: §4 · p. 9

Flexible complexity models

• Elastic models such as Slimmable Networks [Yu et al., 2018], Once-for-All [Cai et al., 2019], DynaBERT [Hou et al., 2020], and Matformer [Kudugunta et al., 2023] share parameters across capacity levels; Stepped-MoE adds input-adaptive routing and sparsity-aware capacity signaling.The paper reports that sparsity-aware embeddings are critical for consistent behavior across sparsity levels.Source: §4 · p. 9

5 Conclusion and Limitations

• Sparse MoE routing adapts per token but creates an I/O bottleneck, while query-level pruning avoids repeated loading but fixes one sub-network for the response.Flexible Stepped-MoE is presented as a compute-adaptable approach for high throughput on low-resource devices.Source: §5 · p. 9
• 3B and 4B active configurations consistently match or beat dense counterparts, and sometimes outperform equivalent Static S-MoE models on instruction-following and math tasks.NPU profiling indicates Stepped-MoE is almost as fast as equally compute-intensive dense models and more accurate than Vanilla-MoE models.Source: §5 · p. 9
• The ablations indicate shorter segments improve accuracy but increase I/O pressure, making the operating point dependent on target-hardware memory bandwidth.The authors say their ablations only scratch the surface of this trade-off.Source: §5 · p. 9
• The study does not demonstrate the effects of longer training, and shared attention parameters may saturate accuracy for larger active sizes.The authors also report pre-training compute almost equal to training equivalent static Stepped-MoE models.Source: §5 · p. 9

A.1 Impact Statement

• The stated goal is one model serving multiple active parameter sizes on resource-constrained devices while improving inference efficiency and LLM performance.The authors say training data was filtered for quality and safety, including removal of personal data from SFT data.Source: §A.1 · p. 12
• The authors state that their algorithm introduces no additional bias or harm risks, and may reduce environmental impact through more efficient inference.They acknowledge potential misuse of machine-learning models and encourage further research on bias and unintended consequences.Source: §A.1 · p. 12

A.2 Scalablity of Stepped-MoE

• Increasing total parameters while keeping active parameters fixed improves accuracy on most benchmarks; Table 5 presents pre-training results for larger-total-parameter S-MoE models.Source: §A.2 · p. 12, Table 5 · p. 12

A.3 Batch replication vs more training data

• Additional unique training tokens improve MMLU performance on larger active sizes, while batch replication versus more data does not significantly improve accuracy on smaller active sizes.Table 6 is identified as pre-training results comparing batch replication with more training data without batch replication.Source: §A.3 · p. 12, Table 6 · p. 12

A.4 Routing mask validation

• Within a single topic, segments use highly similar experts, a pattern that can lead to high CHR.The masks are compared using cosine similarity.Source: §A.4 · p. 13
• As topic A transitions to topic B, expert similarity wears off, indicating that the selected experts adapt to context changes.Source: §A.4 · p. 13
• After a drastic context switch, expert choices change markedly, while a return to an earlier topic after an unrelated diversion brings back similar choices.Source: §A.4 · p. 13

A.5 Fine-Grained Core-En Results

• At S = 16, Core-En-0s rises from 0.640 to 0.666 at 1B and from 0.701 to 0.722 at 3B, exceeding dense scores of 0.662 and 0.704.WinoGrande remains below Dense at both active sizes; the authors suggest short ARC prompts may finish before experts are refreshed with more context.Source: §A.5 · p. 13, Table 7 · p. 14

A.6 Compute-Normalized Training Comparison

• The approximate training-FLOP budget matches the shared checkpoint to the combined two-model baseline portfolio, while each individual baseline receives about half its total budget.The approximation uses C ≈6NT and omits routing, attention, communication, and hardware-utilization overheads.Source: §A.6 · p. 14, Table 8 · p. 14
• Flexible S-MoE beats Dense on MMLU-5s and GSM-8k but trails on both Core-En metrics; Static S-MoE outperforms Flexible S-MoE on all four metrics at both active sizes.The results illustrate a trade-off between supporting multiple active budgets in one checkpoint and optimizing a separate model for each budget.Source: §A.6 · p. 14, Table 8 · p. 14
• The shared Flexible S-MoE run costs 203,435 TPU core-hours, versus Dense costs of 30,720 at 1B and 29,123 at 3B, and Static costs of 92,160 and 82,375, respectively.Despite this practical training-cost overhead, the shared checkpoint fits limited storage while serving multiple active sizes; the FLOP approximation excludes routing, attention, communication, and hardware-utilization overheads.Source: §A.6 · p. 14

A.7 Cross-Segment-Length Evaluation

• Increasing inference segment length from 16 to 32 largely preserves performance and reduces expert-refresh frequency, potentially lowering I/O overhead at minimal accuracy cost.The evaluation uses one checkpoint trained at S_train = 16 and held at 3B active parameters.Source: §A.7 · p. 15, Table 9 · p. 15