Research map / Architectures and efficiency / Mixture of experts
ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes ReMoE, a fully differentiable MoE that replaces TopK+Softmax routing with ReLU routing, controlled by an adaptive L1 regularizer with load-balancing refinement, and reports consistent gains over TopK MoE on LLaMA-style models trained on The Pile.
Problem
Vanilla TopK MoE routers are discontinuous and non-differentiable, which limits MoE performance and scalability, while existing fully differentiable MoE variants either break causality or underperform TopK routing.
Approach
Replace TopK+Softmax with a ReLU router whose natural zeros give sparse, continuous, fully differentiable expert gating; hit a target sparsity with an adaptively tuned L1 regularizer whose per-expert weights integrate load balancing.
Key results
- ReMoE (182M, E=8, k=1) achieves average zero-shot accuracy 40.03 vs 39.53 for TopK MoE and validation loss 1.921 vs 1.936.
- ReMoE scales with expert count: at E=128 avg accuracy 42.12 vs 41.10 for MoE, with a steeper slope as E increases.
- Fine-grained ReMoE with G=32 and G=64 reaches the Dense x8 upper bound performance while using far fewer FLOPs.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace TopK+Softmax routing with ReLU routing | MoE · Downstream tasks, 120B tokens | Avg. accuracy: 44.12 -> 45.49 | Helped |
| replace TopK+Softmax routing with ReLU routing | MoE · The Pile, 120B tokens | Valid Loss: 1.716 -> 1.689 | Helped |
| add load-balancing refinement to the L1 regularization (ablation) | ReMoE · Downstream tasks, load balancing ablation | Avg. accuracy: 39.52 -> 40.03 | Helped |
| add load balancing (ablation) | TopK MoE · Downstream tasks, load balancing ablation | Avg. accuracy: 37.67 -> 39.67 | Helped |
| add a near-dense warmup phase (ablation) | TopK MoE · Downstream tasks, near-dense warmup | Avg. accuracy: 39.53 -> 39.29 | Hurt |
| add a near-dense warmup phase (ablation) | TopK MoE · Valid loss, near-dense warmup | Valid Loss: 1.936 -> 1.928 | Helped |
| vary initial adaptive regularization coefficient λ0 across the tested values (ablation) | ReMoE · ReMoE, λ0 sweep | Valid Loss: 2.031, 2.029, 2.032, 2.036, 2.032 | No clear effect |
| increase the adaptive coefficient update multiplier α to 1.5 (ablation) | ReMoE · ReMoE, α sweep | Valid Loss: 2.032 -> 2.057$^*$ | Hurt |
| increase expert count from 8 to 128 | ReMoE · Downstream tasks, active parameters 182M | Avg. accuracy: 40.03 -> 42.12 | Helped |
| increase expert count from 8 to 128 | MoE · Downstream tasks, active parameters 182M | Avg. accuracy: 39.53 -> 41.10 | Helped |
| increase granularity from 1 to 64 | ReMoE · Downstream tasks, active parameters 182M | Avg. accuracy: 40.03 -> 41.69 | Helped |
| increase granularity from 1 to 64 | MoE · Downstream tasks, active parameters 182M | Avg. accuracy: 39.53 -> 41.35 | Helped |
Surprising findings
- ReLU routing without load balancing already matches a well-tuned TopK MoE with load balancing, indicating differentiability reduces the need for balancing.
- Fixing the L1 coefficient would cause routing collapse toward zero, so the adaptive coefficient is required.
Limitations
- Stages I and II activate more experts, adding temporary compute and memory overhead.
- Load imbalance is not fully eliminated and remains a hardware utilization issue.
- Experiments are limited to small models (up to 978M active parameters) trained on about 30B tokens.
Experimental setup
Datasets: The Pile
Benchmarks: ARC-c ARC-e BoolQ HellaSwag LAMBADA PIQA RACE
Metrics: validation loss zero-shot accuracy TFLOPS throughput
Baselines: dMoE (token-choice TopK) Expert-choice TopK routing Hash routing Lory SparseMixer-v2 Dense Dense x8
Builds on: TopK routing Switch Transformer load balancing loss Megatron-LM LLaMA Lory
Closest papers by meaning
- ReMoE: Boosting Expert Reuse through Router Fine-Tuning in Memory-Constrained MoE LLM Inference (ICML 2026 · similarity 0.90)
- Layerwise Recurrent Router for Mixture-of-Experts (ICLR 2025 · similarity 0.90)
- Dense Backpropagation Improves Training for Sparse Mixture-of-Experts (NeurIPS 2025 · similarity 0.89)
- MoE++: Accelerating Mixture-of-Experts Methods with Zero-Computation Experts (ICLR 2025 · similarity 0.88)
- Ada-K Routing: Boosting the Efficiency of MoE-based LLMs (ICLR 2025 · similarity 0.87)
- Shortcut-connected Expert Parallelism for Accelerating Mixture of Experts (ICML 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.