Research map / Architectures and efficiency / Mixture of experts

FlexMoE: One-for-All Nested Intra-Expert Pruning for MoE Language Models

NeurIPS 2026 · Architectures and efficiency › Mixture of experts · approach: Expert routing and specialization · 0 citations (OpenAlex)

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.

Converts a pretrained MoE LLM into a nested family of pruned subnetworks by ranking expert FFN channels, learning per-expert retention actions via straight-through Gumbel-Softmax with increasing cost pressure, and a single mid-budget LoRA recovery reused across budgets, plus a kernel co-design for online budget switching.

Problem

Existing MoE compression methods produce one compressed model per fixed target budget, requiring separate pipelines or checkpoints for each deployment budget that varies across devices, users, and workloads.

Approach

Importance-rank expert FFN channels once, then learn a token-independent discrete retention action per expert under gradually increasing cost pressure so a single action-training run exports nested masks across budgets; one recovery LoRA at a 40% mid-budget mask is merged and reused for all budgets.

Key results

• On Qwen2-57B-A14B, without fine-tuning, the pruned model retains ~99.8% of base average accuracy at 50% routed-expert pruning (58.57 vs 58.71).
• On Mixtral-8x7B at 40% pruning with recovery fine-tune, average accuracy 58.00 vs TD-MoE 57.00 and MoE-SVD 50.00.
• Offline pruned Qwen2-57B-A14B at 60% prune reaches 5607.8 tok/s vs 1228.0 base (x4.57) under SGLang on a single H200.

Surprising findings

• Qwen2-57B-A14B retains about 92.9% of base average accuracy at 80% pruning without fine-tuning.
• Fine-tuning at the 20% high-budget point transfers poorly to tighter budgets, while the mid-budget (40%) recovery point generalizes best.

Limitations

• Online budget switching remains slower than offline pruned checkpoints due to slicing and dispatch overhead.
• Action learning is costly (about 48 hours on Mixtral, 30 on Phi-3.5-MoE).
• Single-point recovery sacrifices some accuracy at high budgets relative to per-budget fine-tuning because tail channels are not updated.

Experimental setup

Datasets: Zyda-2

Benchmarks: ARC-Challenge ARC-Easy HellaSwag OpenBookQA PIQA WinoGrande MathQA

Metrics: zero-shot accuracy average accuracy throughput (tok/s) checkpoint size (GB)

Baselines: NAEE MoE-I^2 MoE-SVD TD-MoE CP-FT

Builds on: MoE-SVD TD-MoE NAEE AmoebaLLM US-Nets Gumbel-Softmax RECAP LoRA

Closest papers by meaning

• Pruning and Distilling Mixture-of-Experts into Dense Language Models (NeurIPS 2026 · similarity 0.94)
• LightMoE: Reducing Mixture-of-Experts Redundancy through Expert Replacing (NeurIPS 2026 · similarity 0.93)
• Mixture Compressor for Mixture-of-Experts LLMs Gains More (ICLR 2025 · similarity 0.91)
• Advancing Expert Specialization for Better MoE (NeurIPS 2025 · similarity 0.91)
• MoE-SVD: Structured Mixture-of-Experts LLMs Compression via Singular Value Decomposition (ICML 2025 · similarity 0.90)
• MoEQuant: Enhancing Quantization for Mixture-of-Experts Large Language Models via Expert-Balanced Sampling and Affinity Guidance (ICML 2025 · similarity 0.90)