Flex-MoE: Scalable, Adaptive Mixture-of-Experts

• Flex-MoE is a class of dynamic Mixture-of-Experts systems defined by flexible expert design, adaptive scheduling, and efficient resource placement across modalities.
• The frameworks employ dynamic device placement, online profiling, and unified abstractions to optimize expert utilization and mitigate load imbalances in large-scale distributed training.
• They incorporate federated, personalized, and composable expert architectures to enhance performance for multimodal learning, edge computing, and non-IID data environments.

Flex-MoE refers to a family of Mixture-of-Experts (MoE) methods and systems characterized by flexibility in expert design, dynamic system scheduling, data placement, personalized adaptation, and federated or composable deployment. While Flex-MoE appears as a common element in several independently authored frameworks, key works converge on system-level innovation for scalable, efficient, and personalized MoE architectures across domains such as multimodal learning, federated language modeling, and efficient distributed training. The following provides a comprehensive technical overview of major Flex-MoE frameworks, their principles, system architectures, evaluation, and limitations.

1. Flex-MoE System Principles and Objectives

Flex-MoE systems aim to address bottlenecks and limitations inherent in large-scale sparse MoE models. Principal motivating factors include:

• Dynamic Workload and Routing Imbalance: Sparse expert activation commonly results in highly skewed load distributions; a small subset of experts handle most tokens, leading to compute bottlenecks and inefficiency (Nie et al., 2023).
• Scalability and System Adaptation: Efficiently mapping MoE computation to heterogeneous, large-scale hardware while adapting to dynamic routing patterns is critical for throughput and resource utilization.
• Flexible Expert Specialization: Varying domain, client, or modality demands require malleable expert assignments and, in some frameworks, support for personalized or modular expert construction (Yun et al., 2024, Liu et al., 1 Jun 2025).
• Communication and Memory Efficiency: Operating within stringent resource constraints (e.g., federated edge environments or high-expert-count settings) necessitates reductions in model transmission and activation footprints without sacrificing model performance.
• Modality and Federation Adaptation: For multimodal and federated tasks, Flex-MoE systems introduce dynamic routing mechanisms or device placements that tune expert activation and capacity assignment to observed data or end-user constraints (Zhang et al., 28 Dec 2025, Yun et al., 2024).

2. System Architectures and Scheduling Mechanisms

2.1 Dynamic Device Placement and Scheduling

The original FlexMoE system (Nie et al., 2023) innovates through a runtime “Scheduler” that continually reassigns experts across devices in response to observed routing statistics. The system employs a vExpert abstraction—lightweight expert replicas—that can be expanded or shrunk per-GPU, and re-mapped to hardware depending on load. At each iteration, token assignments $I_{e,g}$ are monitored and a “balance ratio” (max per-GPU vs. average load) is computed. If this ratio exceeds a tunable threshold, a cost-model–driven heuristic (using per-expert compute and communication models) proposes atomic Expand/Shrink/Migrate operations, seeking to minimize the step time $T(\mathcal{I},\mathcal{P})$ across all devices.

2.2 Unified Abstractions and Online Profiling

FSMoE (Pan et al., 18 Jan 2025) decomposes each MoE layer into four pluggable modules: token routing, token communication (ordering and dispatch/combine), expert computation, and expert parallelism. Each submodule conforms to an interface (e.g., GateBase, OrderBase), allowing subclassing for new gating or dispatch algorithms. Before training, FSMoE runs microbenchmarks of each kernel and fits linear time/performance models, enabling the system to select optimal scheduling and pipelining parameters with negligible overhead.

2.3 Adaptive Parallelism and Pipelining

The Flex system in Tutel (Hwang et al., 2022) and FlexMoE family supports switching between data-, model-, and expert-parallel (DP/MP/EP) schemes at each step, without any data or parameter migration. This is achieved through a unified data/parameter layout, allowing O(1) time switching. An adaptive pipelining controller dynamically partitions batches and selects between linear and hierarchical (2DH) All-to-Alls, using pre-built hash tables of optimal configuration per workload bucket. The result is near-optimal communication cost for a broad spectrum of MoE parallelization schemes.

3. Flexible Expert Specialization and Federated Adaptation

3.1 Personalized and Resource-Constrained Federated MoE

FLEX-MoE (Zhang et al., 28 Dec 2025) addresses expert assignment under strict client capacity and non-IID data in federated settings. The server holds parameters for all experts and a shared extractor, but each client is assigned only a capacity-constrained subset of experts in each round. Assignment is governed by a joint ILP that maximizes client-expert fitness (using accuracy- or loss-driven indicators) and enforces balanced expert utilization via dynamic load constraints. This design substantially reduces expert-load coefficient-of-variation (CV) and improves average accuracy compared to greedy or random assignment, even under highly non-IID data.

3.2 Composable and Rank-Heterogeneous Experts

FlexOlmo (Shi et al., 9 Jul 2025) and FlexMoRE (Pirchert et al., 9 Feb 2026) frameworks support independently trained domain experts—each on local data—which can be federatedly merged at inference by router concatenation. FlexMoRE extends this to rank-heterogeneous mixtures, where each expert is either a full-size FFN or a low-rank adapter obtained via SVD (LoRA-style). This enables parameter-efficient deployment while maintaining or exceeding the performance of full-expert mixtures, with domain/task-specific rank adaptation for knowledge- versus reasoning-intensive tasks.

4. Multimodal and Modality-Combination Flexibility

The Flex-MoE framework for multimodal learning (Yun et al., 2024) introduces a “missing modality bank” for contextualized embedding of absent modalities (indexed by observed combinations), allowing models to flexibly handle arbitrary incomplete input sets across $2^M-1$ possible nonempty modality combinations. The routing mechanism employs two stages: (1) a generalized router (𝒢-Router) trained on complete data to inject broad knowledge, and (2) a specialized router (𝒮-Router) that locks top-1 assignment to the correct expert matching the observed combination, with remaining top-K slots allocated by standard gate-based scores. This enables superior performance, especially in medical and clinical settings characterized by frequent missing modalities.

5. Empirical Performance and Evaluation

Flex-MoE systems consistently outperform baselines across efficiency, load balance, and task accuracy metrics. Notable quantitative results include:

• FlexMoE achieves up to 1.70×–2.10× end-to-end throughput gains over DeepSpeed and up to 1.45× over FasterMoE on a range of BERT-MoE, GPT-MoE, and Swin-MoE models (Nie et al., 2023).
• FSMoE attains 1.18×–3.01× speedup over DeepSpeed-MoE and Tutel on real-world LLMs (GPT-2, Mixtral) (Pan et al., 18 Jan 2025).
• FLEX-MoE reduces expert-load CV by two orders of magnitude and recovers accuracy under non-IID/disjoint partitions in federated learning, reaching up to +6% accuracy improvement over best baselines (Zhang et al., 28 Dec 2025).
• FlexOlmo reports a 41% mean improvement over public-expert-only models on 31 diverse NLP benchmarks, with opt-out/merging capability for regulated data (Shi et al., 9 Jul 2025).
• Flex-MoE in multimodal settings attains +4.7% ACC on ADNI and +12.2% on MIMIC-IV over best multimodal baselines, with significant parameter reduction (–89.2% vs. FuseMoE) (Yun et al., 2024).

6. Design Trade-Offs, Limitations, and Future Directions

Key trade-offs in Flex-MoE involve:

• System Overheads: Dynamic vExpert expansion/migration, scheduler profiling, and per-step placement adjustments incur minor overhead versus static baselines but deliver substantially improved system efficiency.
• Federated Coordination versus Centralization: Decentralized expert training and adapter-based parameterization reduce privacy risk and memory, but introduce rank-selection sensitivity and potential for degraded transfer when routing or base model quality is suboptimal (Pirchert et al., 9 Feb 2026).
• Routing Load-Balance vs. Model Quality: System-level device rebalancing or scheduling can restore utilization harmed by over-regularized balance losses, but requires accurate cost modeling and slow adaptation windows to mitigate oscillations.
• Modality Explosion and Routing Overhead: In multimodal settings, the number of modality combinations (exponential in modalities) bounds expert counts, suggesting future scaling challenges as M→10+ (Yun et al., 2024).

Future work emphasizes:

• Extension to exascale and heterogeneous hardware (TPUs, CPUs) (Hwang et al., 2022).
• More principled, possibly reinforcement learning–based, scheduling and adaptation for changing network or compute circumstances (Pan et al., 18 Jan 2025).
• Fine-grained, learned gating or adapter routing per user/domain in federated and regulated settings (Shi et al., 9 Jul 2025, Pirchert et al., 9 Feb 2026).
• Support for mixed expert types, dynamic expert merging/creation, and scaling laws for modality and expert count (Yun et al., 2024, Pirchert et al., 9 Feb 2026).

7. Comparative Summary of Flex-MoE Variants

Framework / Paper	Key Innovation	Domain/Application
FlexMoE (Nie et al., 2023)	Dynamic device placement, vExpert scheduling	Large-scale vision/NLP
FSMoE (Pan et al., 18 Jan 2025)	Modular abstraction, pipelined scheduling	General MoE (LLMs, Mixtral)
Flex (Tutel) (Hwang et al., 2022)	Unified tensor layout: O(1) parallelism switching	Vision (Swin-MoE)
FLEX-MoE (Zhang et al., 28 Dec 2025)	Federated, optimized/balanced expert assignment	Edge/Federated LMs
FlexOlmo (Shi et al., 9 Jul 2025)	Opt-in composable federated MoE, data-flexible inference	Regulated LMs
FlexMoRE (Pirchert et al., 9 Feb 2026)	Rank-heterogeneous, parameter-optimized Experts	Memory-bounded LMs
Flex-MoE (multimodal) (Yun et al., 2024)	Modality-combination–aware routing, missing modality bank	Biomedical MM LMs

Flex-MoE frameworks collectively represent a progression toward practical, highly configurable, and efficient MoE systems that adapt their computation to dynamic workload, data, hardware, and privacy constraints. Each variant addresses a specific subdomain of the broad MoE scaling and deployment problem, with system-level and architectural flexibility as unifying design principles.