FlexMoRE: A Flexible Mixture of Rank-heterogeneous Experts for Efficient LLMs

Paper: 2602.08818 Authors: Annemette Brok Pirchert, Jacob Nielsen, Mogens Henrik From, Lukas Galke Poech, Peter Schneider-Kamp Categories: cs.LG

Abstract

FlexMoRE introduces a Mixture-of-Experts architecture where experts can be either full-sized or low-rank adapters of varying ranks. The paper systematically investigates the trade-off between expert rank and downstream performance using 6 experts at ranks from $2^0$ to $2^{14}$, evaluating 150 different mixtures across 120 tasks. A key finding is that the optimal rank is substantially higher for reasoning-heavy benchmarks compared to knowledge-heavy benchmarks. Using task-optimal ranks, FlexMoRE achieves improved performance (47.18 vs 45.46) at less than one-third the parameters (10.75B vs 33.27B).

Key Contributions

• Rank-heterogeneous experts: Experts within the same MoE layer can have different ranks, from full-size to extremely low-rank
• Systematic rank investigation: 150 mixtures across 120 tasks with ranks spanning $2^0$ to $2^{14}$
• Task-dependent optimal rank: Reasoning tasks require higher ranks than knowledge retrieval tasks
• Parameter efficiency: 3x parameter reduction with improved performance
• Federated learning compatibility: Designed for efficient federated training scenarios

Background: MoE and Low-Rank Adaptation

Mixture of Experts

MoE architectures route each input to a subset of specialized experts, enabling larger model capacity without proportional compute increase. A router network $G(x)$ selects the top-$k$ experts:

$y = \sum_{i \in \text{top-}k} G(x)_i \cdot E_i(x)$

Low-Rank Adaptation (LoRA)

LoRA approximates weight updates with low-rank matrices:

$W' = W + BA$

where $B \in \mathbb{R}^{d \times r}$ and $A \in \mathbb{R}^{r \times d}$, with rank $r \ll d$. This dramatically reduces trainable parameters.

The FlexMoRE Insight

FlexMoRE combines these ideas: each expert in the MoE can be a LoRA adapter of a different rank. This allows the system to allocate capacity where it matters most.

Experimental Design

Rank Space

The study evaluates experts at ranks: $r \in \{1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384\}$ (i.e., $2^0$ through $2^{14}$), plus full-rank experts.

Mixture Configurations

With 6 experts per layer, the study explores 150 different rank combinations, systematically varying:

• Homogeneous mixtures (all experts same rank)
• Heterogeneous mixtures (experts with different ranks)
• Graduated mixtures (ranks increasing or decreasing across experts)

Evaluation

Performance is measured across 120 tasks spanning:

• Reasoning benchmarks: Mathematical reasoning, logical deduction, code generation
• Knowledge benchmarks: Factual QA, commonsense reasoning, world knowledge

Key Finding: Task-Dependent Optimal Rank

The most significant finding is that optimal rank depends strongly on task type:

• Reasoning-heavy tasks: Require substantially higher ranks (often $r \geq 256$)
• Knowledge-heavy tasks: Perform well even at very low ranks ($r \leq 16$)

This makes intuitive sense: knoeval primarily requires accessing stored information (which can be captured in low-rank subspaces), while reasoning requires more complex transformations that need higher-dimensional representations.

Results

Configuration	Avg. Accuracy	Parameters
Full-rank baseline	45.46	33.27B
FlexMoRE (optimal)	47.18	10.75B
Homogeneous low-rank	43.21	8.92B

• 47.18 vs 45.46: FlexMoRE outperforms the full-rank baseline
• 10.75B vs 33.27B: Using less than 1/3 of the parameters
• Heterogeneous mixtures consistently outperform homogeneous ones

Significance

1. Not all experts need to be equal: Rank heterogeneity is a powerful design dimension for MoE architectures
2. Task-aware architecture design: The finding that reasoning needs higher ranks than knowledge retrieval has broad implications for model design
3. Massive parameter savings: 3x reduction with improved performance challenges the assumption that bigger is always better
4. Federated learning: Rank-heterogeneous experts are naturally suited to federated settings where different clients have different compute budgets
5. Practical guidance: The systematic study provides concrete guidance for practitioners choosing expert configurations