Scaling Mixture of Experts: Architecture Search for Billion-Parameter Language Models

Authors: Kshitij Thakkar Date: February 2026 Collection: Large MoE Architecture Search (1B-2B) (22 models) Dataset: kshitijthakkar/large-moe-inference-benchmark

Abstract We present a systematic architecture search for billion-parameter Mixture of Experts (MoE) language models, exploring 22 configurations across model architectures, routing strategies, and training dynamics. Our experiments reveal three key findings: (1) aggressive top-8 routing with 16 experts significantly outperforms conservative top-2 routing, achieving 34% lower loss; (2) smaller batch sizes with longer context windows (bs=1, ctx=2048) dramatically outperform larger batches with shorter contexts, yielding 3.6x lower loss; and (3) shallow-wide architectures (8 layers, 2048 dim) outperform deep-narrow alternatives (16 layers, 1024 dim) at the 1B+ scale. Our best model achieves 0.32 validation loss with 1.08B total parameters (781M active), running at 17.3 tokens/second on A100 GPU. All models and code are publicly available.

Keywords: Mixture of Experts, Architecture Search, Large Language Models, Sparse Models

---

1 Introduction

Scaling language models beyond one billion parameters presents fundamental challenges in memory efficiency, training stability, and computational cost. Mixture of Experts (MoE) architectures address these challenges through conditional computation, activating only a subset of parameters for each input token while maintaining high total capacity.

Despite the success of MoE models like Mixtral [1], Switch Transformer [2], and DeepSeek-MoE [3], the optimal configuration for billion-parameter MoE models remains understudied. Prior work has focused on either very large models (>100B parameters) or mobile-scale models (<500M parameters), leaving a gap in understanding optimal designs for the practical 1-2B parameter range.

This work addresses this gap through systematic architecture search, contributing:

• Comprehensive evaluation of 22 model configurations spanning 1B-1.7B parameters
• Novel finding that context length matters more than batch size for MoE training
• Routing analysis showing top-8 outperforms top-2 at scale
• Open release of all models, training code, and inference benchmarks

2 Related Work

Mixture of Experts. The MoE paradigm, introduced by Jacobs et al. [4] and scaled by Shazeer et al. [5], enables conditional computation where only a subset of "expert" networks process each token. Recent work has demonstrated MoE's effectiveness at scale: Mixtral-8x7B [1] achieves performance comparable to 70B dense models while using only 13B active parameters.

Architecture Search for LLMs. While neural architecture search (NAS) has been applied to vision models extensively, its application to LLMs remains limited. The Smol Training Playbook [6] provides guidelines for small model training but does not address MoE-specific considerations. Our work fills this gap by systematically exploring MoE design choices.

Training Dynamics. Recent work has highlighted the importance of training hyperparameters for LLM quality. The Chinchilla study [7] established compute-optimal scaling laws, while subsequent work [8] has shown these laws may not hold for MoE architectures due to their unique sparse computation patterns.

3 Experimental Setup

3.1 Base Architecture

All models use a Qwen3-style MoE architecture with:

• Normalization: RMSNorm (pre-normalization)
• Activation: SiLU (Swish) in feed-forward layers
• Position Encoding: Rotary Position Embeddings (RoPE) with base frequency 1,000,000
• Attention: Grouped Query Attention (GQA) with configurable KV groups
• Routing: Top-k softmax-normalized expert selection
• Vocabulary: 151,936 tokens (Qwen3 tokenizer)
• Context Capacity: 262,144 tokens (256K)

3.2 Search Space

We explore three phases of architecture search:

Phase 1: Model Architecture (5 configurations)

Phase 2: Learning Rate (9 configurations)

Learning rates from 5e-6 to 1e-3 on the best Phase 1 architecture.

Phase 3: Batch Size x Context Length (13 configurations)

3.3 Training Configuration

• Optimizer: AdamW with beta=(0.9, 0.95), weight_decay=0.1
• Scheduler: Linear warmup (10%) followed by cosine decay
• Training: 2,000 steps per experiment
• Evaluation: Every 500 steps on held-out validation
• Dataset: NVIDIA Nemotron balanced pretraining data
• Hardware: NVIDIA A100 40GB GPU
• Precision: BFloat16 mixed precision

4 Results

4.1 Phase 1: Architecture Comparison

Testing five architectures with default training settings (bs=4, ctx=1024):

Key Finding 1: Top-8 Routing Outperforms Top-2

The large-moe-1.3b with top-8 routing achieved 34% lower loss than its top-2 counterpart (1.15 vs 1.75), despite having identical parameter counts. This suggests that at billion-parameter scale, aggressive expert activation provides better gradient flow and capacity utilization.

Key Finding 2: Active Parameter Ratio Correlates with Quality

Models with higher active parameter ratios consistently achieved lower loss. The best model activates 72% of its parameters per token, compared to 40-51% for lower-performing models.

Key Finding 3: Shallow-Wide Beats Deep-Narrow

The 8-layer, 2048-dim architecture outperformed the 16-layer, 1024-dim variant by 41% (1.15 vs 1.95 loss), while being faster (17.0 vs 13.2 tok/s) and using less memory.

4.2 Phase 3: Batch Size vs Context Length

Using the best Phase 1 model, we tested 13 batch/context combinations:

Critical Finding: Context Length Trumps Batch Size

The most surprising result is that bs=1, ctx=2048 achieves 4.2x lower loss than bs=4, ctx=1024, despite processing the same number of tokens per step. This challenges the conventional wisdom of maximizing batch size and suggests that for MoE models, longer context provides more valuable learning signal than larger batches.

We hypothesize this occurs because:

1. MoE routing benefits from longer sequences to learn token relationships
2. Expert load balancing improves with more diverse tokens per sequence
3. Gradient accumulation effectively compensates for smaller batch sizes

4.3 Inference Performance

Benchmarked on NVIDIA A100 40GB:

The best quality model (large-moe-1.3b-bs1-ctx2048) achieves 17.3 tokens/second, a reasonable trade-off for its 3.6x lower loss compared to faster alternatives.

4.4 Memory Analysis

The optimal configuration (bs=1, ctx=2048) is also one of the most memory-efficient, making it practical for A100 40GB deployment.

5 Optimal Configuration

Based on our experiments, we recommend the following configuration for 1.3B MoE models:

model:
  hidden_size: 2048
  num_hidden_layers: 8
  num_attention_heads: 32
  num_key_value_heads: 8
  head_dim: 128
  num_experts: 16
  num_experts_per_tok: 8      # Key: Top-8 routing
  moe_intermediate_size: 768
  vocab_size: 151936
  max_position_embeddings: 262144

training:
  batch_size: 1               # Key: Small batch
  context_length: 2048        # Key: Long context
  learning_rate: 1e-4
  gradient_accumulation: 4
  warmup_ratio: 0.1
  weight_decay: 0.1
  optimizer: adamw
  scheduler: cosine

5.1 Hardware-Specific Recommendations

6 Discussion

6.1 Why Top-8 Routing?

Our results show top-8 routing outperforms top-2 by 34% despite identical parameter counts. We attribute this to:

1. Better gradient distribution: 50% of experts receive gradients vs 12.5%
2. Reduced winner-take-all dynamics: Smoother routing prevents expert collapse
3. Higher effective capacity: More parameters contribute to each prediction

The inference cost increases ~4x compared to top-2, but training quality improvements justify this trade-off for many applications.

6.2 Context vs Batch Trade-off

The dramatic improvement from longer context (4.2x better loss) suggests MoE models benefit from:

1. Richer routing signals: Longer sequences provide more diverse tokens for router learning
2. Better expert specialization: Experts can specialize on token patterns within context
3. Improved load balancing: More tokens per batch leads to more uniform expert utilization

This finding has practical implications: practitioners should prioritize context length over batch size when memory-constrained.

6.3 Comparison with Prior Work

Our findings diverge from dense model guidelines, particularly around batch size, highlighting the need for MoE-specific training practices.

7 Limitations

1. Short training runs: 2,000 steps may not capture long-training dynamics
2. Single dataset: Results may vary with different pretraining corpora
3. Fixed learning rate in Phase 3: Optimal LR may differ for each batch/context combination
4. Memory constraints: Could not fully explore longer context configurations

8 Conclusion

We present a systematic architecture search for billion-parameter MoE language models, testing 22 configurations across three experimental phases. Our key findings are:

1. Top-8 routing is optimal for 1.3B MoE models, providing 34% lower loss than top-2
2. Context length matters more than batch size, with bs=1, ctx=2048 achieving 3.6x better results than bs=4, ctx=1024
3. Shallow-wide architectures outperform deep-narrow at the 1B+ scale
4. The optimal configuration is also memory-efficient, enabling deployment on 40GB GPUs

All 22 models are available at huggingface.co/collections/kshitijthakkar/large-moe-architecture-search-1b-2b.

References

[1] Jiang, A. Q., et al. "Mixtral of Experts." arXiv preprint arXiv:2401.04088 (2024).

[2] Fedus, W., Zoph, B., & Shazeer, N. "Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity." JMLR (2022).

[3] Dai, D., et al. "DeepSeek-MoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models." arXiv preprint arXiv:2401.06066 (2024).

[4] Jacobs, R. A., et al. "Adaptive mixtures of local experts." Neural computation 3.1 (1991): 79-87.

[5] Shazeer, N., et al. "Outrageously large neural networks: The sparsely-gated mixture-of-experts layer." ICLR (2017).

[6] HuggingFace. "The Smol Training Playbook." (2025).

[7] Hoffmann, J., et al. "Training compute-optimal large language models." NeurIPS (2022).

[8] Clark, A., et al. "Unified scaling laws for routed language models." ICML (2022).

---

Appendix A: Full Experiment Results

A.1 Phase 1: Model Architecture

A.2 Phase 3: Batch x Context

A.3 Inference Benchmarks

---

Resources

• Model Collection: Large MoE Architecture Search (1B-2B)
• Inference Benchmark Dataset: kshitijthakkar/large-moe-inference-benchmark

Citation

@misc{thakkar2026scalingmoe,
  title={Scaling Mixture of Experts: Architecture Search for Billion-Parameter Language Models},
  author={Thakkar, Kshitij},
  year={2026},
  url={https://huggingface.co/collections/kshitijthakkar/large-moe-architecture-search-1b-2b}
}