Research map / Architectures and efficiency / Mixture of experts
Ada-K Routing: Boosting the Efficiency of MoE-based LLMs
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 Ada-K routing, which uses lightweight PPO-trained per-layer allocators on frozen MoE LLMs to choose a per-token number of activated experts, cutting activations 30-40% while slightly improving benchmark accuracy.
Problem
Static Top-K routing in MoE-based LLMs activates a fixed, equal number of experts for every token regardless of its importance or difficulty, wasting compute on simple tokens and under-serving complex ones.
Approach
Attach a small learnable allocator to each layer that samples a per-token expert count k from a distribution, trained end-to-end on a frozen MoE model with PPO using the log-likelihood as reward plus an expected-count regularizer, and warm-started by pseudo-labels from Top-P (nucleus) routing.
Key results
- On Mixtral-8x22B, Ada-K reaches 79.92 average accuracy vs 79.15 for Top-K (k=2) while cutting activated experts to 1.31 (34.4% reduction) and FLOPs by ~30% (14.08T vs 20.04T), with 1.31x speedup.
- On Qwen1.5-MoE-A2.7B, Ada-K averages 55.13 vs 54.43 for Top-K (k=4) with 2.58 activated experts (35.5% reduction).
- Training allocators takes under 8 hours (7.92h for Mixtral-8x22B) with about 1-3M trainable parameters.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace default Top-K routing with Ada-K dynamic per-token expert allocation | Mixtral-8x22B · Open LLM Leaderboard (average) | accuracy: 79.15 -> 79.92 | Helped |
| replace default Top-K routing with Ada-K dynamic per-token expert allocation | Mixtral-8x7B · Open LLM Leaderboard (average) | accuracy: 67.58 -> 68.19 | Helped |
| replace default Top-K routing with Ada-K dynamic per-token expert allocation | Qwen1.5-MoE-A2.7B · Open LLM Leaderboard (average) | accuracy: 54.43 -> 55.13 | Helped |
| replace default Top-K routing with Ada-K dynamic per-token expert allocation | DeepSeek-MoE-16B · Open LLM Leaderboard (average) | accuracy: 49.91 -> 50.55 | Helped |
| fine-tune routers for reduced static Top-K activation (ablation) | Qwen1.5-MoE-A2.7B · six-benchmark average | accuracy: 52.54 -> 52.85; 53.34 -> 53.44; 54.43 -> 53.89 | No clear effect |
| use threshold-based MoED or D2D routing instead of Ada-K (ablation) | Qwen1.5-MoE-A2.7B · six-benchmark average | accuracy: 55.13 vs. 53.45; 53.60; 53.73; 53.64 | Hurt |
| use activation count as PPO reward instead of regularization loss (ablation) | Ada-K on Qwen1.5-MoE-A2.7B · six-benchmark average | accuracy: 54.64 vs. 55.13 | No clear effect |
| train allocators on pretrain data instead of SFT data (ablation) | Ada-K on Qwen1.5-MoE-A2.7B · six-benchmark average | accuracy: 55.13 -> 55.78 | Helped |
| remove allocator warm-start pretraining (ablation) | Ada-K on Qwen1.5-MoE-A2.7B · six-benchmark average | accuracy: 55.13 -> 54.18 | Hurt |
| use K-Warm rather than P-Warm allocator initialization (ablation) | Ada-K on Qwen1.5-MoE-A2.7B · six-benchmark average | accuracy: 55.13 -> 54.97 | Hurt |
| increase the fraction of layers equipped with allocators (ablation) | Ada-K on Qwen1.5-MoE-A2.7B · six-benchmark evaluation | FLOPs: 1.19T -> 1.14T -> 1.04T -> 0.92T | Helped |
| increase activation reduction beyond the trade-off point (ablation) | Ada-K on Qwen1.5-MoE-A2.7B · trade-off curve | accuracy | Hurt |
Surprising findings
- Adaptive routing improves accuracy over the full Top-K baseline while activating fewer experts per token.
- Simply fine-tuning the routers of lower-k static Top-K models gives only modest gains, and threshold-based methods (MoED, D2DMoE) fall well short of Ada-K.
Limitations
- Evaluated only as a post-training method on frozen models; not validated during full pretraining from scratch.
- Reported gains on some benchmarks are small or negative (e.g., ARC-E and MMLU slight drops), and results are from single runs on six benchmarks.
- Trade-off curve shows accuracy declines sharply beyond about 44% activation reduction.
Experimental setup
Datasets: Alpaca GPT4 UltraChat 200k LIMA OpenPlatypus CodeAlpaca 20k Wiki QA MathInstruct Wiki Demo RedPajama V2 Wikipedia StarCoder
Benchmarks: ARC-C HellaSwag MMLU GSM8K TruthfulQA Winogrande BBH ARC-E
Metrics: accuracy average activated experts activation reduction rate FLOPs inference speedup
Baselines: Top-K routing MoED D2D
Builds on: Top-K routing PPO nucleus sampling MoE-D D2DMoE
Closest papers by meaning
- $\textit{Read-ME}$: Refactorizing LLMs as Router-Decoupled Mixture of Experts with System Co-Design (NeurIPS 2024 · similarity 0.91)
- Oracle-MoE: Locality-preserving Routing in the Oracle Space for Memory-constrained Large Language Model Inference (ICML 2025 · similarity 0.90)
- Advancing Expert Specialization for Better MoE (NeurIPS 2025 · similarity 0.90)
- Layerwise Recurrent Router for Mixture-of-Experts (ICLR 2025 · similarity 0.90)
- OpenMoE: An Early Effort on Open Mixture-of-Experts Language Models (ICML 2024 · similarity 0.90)
- Certain Head, Uncertain Tail: Expert-Sample for Test-Time Scaling in Fine-Grained MoE (ICML 2026 · similarity 0.89)
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.