Source-linked AI summary
FlexMoRE: A Flexible Mixture of Rank-heterogeneous Experts for Efficient Federatedly-trained Large Language Models
Annemette Brok Pirchert, Jacob Nielsen, Mogens Henrik From, Lukas Galke Poech, Peter Schneider-Kamp
TL;DR
Domain-specific data is often difficult to centralize, while maintaining multiple full-sized domain-adapted experts is resource-intensive. FlexMoRE combines full-sized and low-rank experts, studies rank sensitivity, and finds that task-dependent ranks can improve performance while reducing parameters.
Problem
Privacy, legal, and proprietary constraints limit centralized access to domain-specific data, while maintaining multiple domain-adapted models is costly.
Method
FlexMoRE combines independently trained full-sized and low-rank experts within a shared mixture-of-experts framework and analyzes rank sensitivity using regression across tasks.
Results
Optimal expert rank is task-dependent: reasoning-heavy benchmarks require higher ranks than knowledge-oriented tasks, while low-rank experts can improve performance and reduce parameters by 67%.
Takeaways & Limitations
Rank should be allocated selectively rather than uniformly to achieve substantial memory savings without sacrificing performance.
Takeaways & Limitations
Low-rank expert derivation introduces approximation error when the decomposition is truncated at rank r.
Abstract
from arXiv · showhide
Recent advances in mixture-of-experts architectures have shown that individual experts models can be trained federatedly, i.e., in isolation from other experts by using a common base model to facilitate coordination. However, we hypothesize that full-sized experts may not be necessary for all domains and that instead low-rank adapters may be sufficient. Here, we introduce FlexMoRE, a Flexible Mixture of Rank-heterogenous Experts, which may be either full-sized experts or adapters of a suitable rank. We systematically investigate the trade-off between expert rank and downstream task performance by evaluating $6$ experts with ranks $2^0$ to $2^{14}$ resulting in experiments covering 150 mixtures (96 with 2 experts, 54 with 7 experts) that are evaluated across $120$ tasks. For our experiments, we build on FlexOlmo and turn its pre-trained experts into low-rank versions. Our regression analysis from expert rank to downstream task performance reveals that the best-performing rank is substantially higher for reasoning-heavy benchmarks than for knowledge-heavy benchmarks. These findings on rank sensitivity come with direct implications for memory efficiency: Using optimal ranks, FlexMoRE yields improved downstream task performance (average score $47.18$) compared to the baseline FlexOlmo-style mixture of full-sized experts (average score $45.46$) at less than one third the parameters ($10.75$B for FlexMoRE vs. $33.27$B for FlexOlmo). All code will be made available.
1 Introduction
FlexMoRE addresses privacy-constrained domain adaptation by combining full-size and rank-heterogeneous low-rank experts in a federatedly composable MoE. The architecture targets lower memory use while preserving or improving benchmark performance.
- Motivation: Privacy, legal, and proprietary constraints limit centralized domain-data sharing, motivating architectures that incorporate expertise without direct data exchange.These constraints arise in settings including healthcare, law, and enterprise systems.
- Motivation: Full-size experts make FlexOlmo difficult to scale because of high accelerator-memory requirements.
- FlexMoRE: FlexMoRE combines full-size experts with low-rank experts within the same MoE routing framework.Low-rank experts can be trained as adapters or derived through post-hoc low-rank factorization of fully fine-tuned experts.
- FlexMoRE: FlexMoRE extends decentralized training and inference-time composition to experts parameterized by low-rank approximations.
- Results: At one third of the memory requirements, post-hoc LoRA experts retain and can improve performance across most benchmarks without additional training.
- Results: Reasoning-heavy tasks require higher LoRA ranks than knowledge-heavy tasks.
2 Related Work
Related work addresses MoE scaling, parameter-efficient adaptation, and decentralized composition separately. FlexMoRE targets the unresolved question of whether FlexOlmo-style decentralized composition remains effective with low-rank rather than full-size experts.
- Mixture-of-Experts: MoE architectures scale capacity through sparse routing but commonly rely on centrally trained full-size experts with high parameter and memory costs.
- Low-Rank Adapters and Mixtures: LoRA introduces lightweight trainable modules while keeping the backbone frozen, reducing adaptation cost and trainable parameters.
- Low-Rank Adapters and Mixtures: Most MoE–LoRA approaches assume centralized training and joint optimization over a shared adapter pool.This limits their applicability when data sharing is restricted and experts are trained federally.
- Decentralized Expert Composition: FlexOlmo enables inference-time composition of independently trained dense experts under data-governance constraints.
- Decentralized Expert Composition: FlexMoRE explores whether FlexOlmo-style decentralized expert composition remains effective with low-rank rather than full-size experts.
3 Methods
FlexMoRE extends decentralized expert composition by combining a full-size base expert with independently trained or extracted low-rank experts of potentially different ranks. It also evaluates how rank relates to task performance using regression and observed peak scores.
- 3.1 The FlexMoRE architecture: FlexMoRE is a rank-heterogeneous MoE architecture that combines one full-size expert with multiple low-rank experts.The paper focuses on compositions with a single full-sized expert and low-rank experts that may have varying ranks.
- 3.1 The FlexMoRE architecture: Low-rank adapters can attach to an arbitrary full-size expert, with the public model serving as the natural base in FlexOlmo.Routing to a low-rank expert also triggers its corresponding full-size base expert.
- 3.2 Deriving Experts through Adapter Extraction: Each domain expert is approximated by subtracting the public base weights, truncating the difference with SVD to rank r, and adding the approximation back to the base expert.Truncation introduces an approximation error, while the reconstructed update is used as the rank-tuned expert.
- 3.2 Deriving Experts through Adapter Extraction: The extracted SVD factors are split into low-rank adapter components so the resulting experts can integrate with LoRA libraries and the MoE computation.The paper also allows low-rank adapters to be trained from scratch alongside the public base model.
- 3.3 Rank Sensitivity Analysis: Rank sensitivity is estimated by regressing evaluation score on log2 expert rank, while typical peak rank is selected directly from observed scores.Positive regression slopes indicate improved performance with increasing rank; near-zero or negative slopes indicate diminishing returns.
4 Experimental Setup
The experiments test whether low-rank experts preserve performance relative to full-sized experts across a FlexOlmo-aligned suite of 120 tasks. They evaluate individual experts, homogeneous low-rank mixtures, and FlexOlmo full-sized baselines under matched public-base conditions.
- 4.1 Evaluation Tasks: 120 tasks from general-purpose and domain-specific benchmark groups are used for evaluation.The groups include MC9, GEN5, AGIEval, BBH, MMLU, and MMLU-Pro.
- 4.2 Evaluation Procedure and Baselines: Average performance is the unweighted mean of the six evaluation-group means.The task-level scores are first aggregated within each group before averaging across groups.
- 4.2 Evaluation Procedure and Baselines: The setup evaluates six available FlexOlmo experts, homogeneous FlexMoRE mixtures with ranks 2^0 through 2^11, and configurations with the public model plus low-rank adapters.The Educational Text expert is excluded because it is not publicly available.
- 4.2 Evaluation Procedure and Baselines: All expert evaluations use a 2x7B setup with each expert alongside the public base expert, unlike FlexOlmo's isolated single-expert evaluation.This procedure is intended to match how the experts were trained and used in the proposed mixtures.
- 4.2 Evaluation Procedure and Baselines: FlexOlmo mixtures with full-sized experts and 2, 4, or 7 active experts provide the main final-model baselines.Relative improvement is quantified against the corresponding full-size baseline using the reported average score difference formula.
5 Results
FlexMoRE experiments show that rank-heterogeneous experts outperform homogeneous alternatives and the FlexOlmo baseline while reducing memory requirements. Rank sensitivity varies by task: reasoning benchmarks favor higher ranks, whereas knowledge benchmarks peak earlier or show weaker sensitivity.
- Single-expert results: Low-rank experts improve over FlexOlmo by 0.73%–7.13% for Code, Creative Writing, Academic, and Reddit, while Math decreases by 1.34%.The best ranks are 29, 24, 26, 211, and 29 for the reported experts, respectively.
- Single-expert results: No single LoRA rank is consistently best across experts, because evaluation groups differ in rank sensitivity and some exhibit flat performance across wide rank ranges.Some groups favor higher-rank specialization.
- Mixture-of-experts results: Rank-heterogeneous FlexMoRE models outperform both homogeneous FlexMoRE models and the FlexOlmo baseline on average.The comparisons include mixtures with 2, 4, and 7 active experts and rank-selection strategies based on MC9 or all benchmarks.
- Mixture-of-experts results: MC9-calibrated rank selection consistently outperforms selection based on average performance across all six evaluation groups.The advantage is most pronounced with 2 or 7 active experts and smaller, though non-negligible, with 4 active experts.
- Efficiency: 31.39%–99.96% memory reductions are reported for the single experts in the best MC9-based heterogeneous model.The reductions vary by expert, with the largest savings for Academic and the smallest for Math.
- Rank sensitivity: Reasoning-heavy benchmarks peak at higher ranks than knowledge-oriented benchmarks, while FlexMoRE rank sensitivity is strongest with four active experts.BBH peaks at median log2 r = 11.5, MMLU at 2, and FlexMoRE a4 has median slope 0.0030.
6 Discussion
The discussion interprets optimal rank as task-specific and argues that rank heterogeneity improves performance while reducing the need for full-size experts. It also describes MC9 as a favorable proxy for selecting expert ranks and identifies router tuning as an untested limitation.
- Task-dependency of optimal ranks: Five of six experts achieve superior downstream performance after post-hoc low-ranking, while the Math expert is sensitive to lower rank and declines slightly.The discussion associates lower ranks with knowledge-heavy datasets and higher ranks with reasoning-based tasks.
- Rank heterogeneity: Rank-heterogeneous models consistently outperform homogeneous models regardless of the homogeneous model’s rank.The result supports allocating capacity unevenly across experts rather than using one uniform rank.
- Expert rank selection: MC9-based rank selection performs favorably because global averaging can select higher-than-necessary ranks, whereas MC9 prioritizes rank-efficient experts.The selection is based on individual-expert performance, not the performance of the complete mixture.
- Limitations: Router tuning was not applied after post-hoc LoRA extraction, although its effect was marginal in FlexOlmo and the authors expect further improvement.This is the discussion’s stated methodological limitation.
7 Conclusion
FlexMoRE combines rank-heterogeneous low-rank adapters with full-size experts for federated LLM training. The conclusion reports improved performance, a 67% parameter reduction, and task-dependent rank requirements.
- Conclusion: FlexMoRE supports rank-heterogeneous low-rank adapters alongside full-size experts for federated LLM training.The architecture is presented as a flexible mixture-of-experts design.
- Conclusion: 67% parameter reductions accompany improved performance from low-rank experts.The conclusion reports this as an empirical finding.
- Conclusion: Reasoning-heavy benchmarks require higher ranks than knowledge-oriented tasks, indicating that rank should be allocated selectively rather than uniformly.The conclusion frames selective allocation as the route to memory savings without sacrificing performance.
A Top-k Experts Benchmark Performance Analysis
The top-k analysis evaluates expert–rank combinations and unique experts independently for each benchmark. These complementary views distinguish peak performance under rank multiplicity from expert-level comparisons without repeated experts.
- Expert–rank combinations: The analysis reports top-k expert models per benchmark, with each entry representing a specific expert–rank pair.Multiple ranks from the same expert may appear in this view.
- Ranking procedure: Rankings are computed independently per benchmark, with ties resolved by selecting the lowest LoRA rank.The rank constraint is rank ≤211 in the unique-expert analysis.
- Expert–rank combinations: The unrestricted top-k view captures peak standalone performance under the rank constraint independently of mixture composition.It emphasizes which expert–rank combinations score highest on each benchmark.
- Unique experts: The unique-expert view allows each expert to appear once at its best-performing rank, facilitating comparison across domains.This removes rank multiplicity from the benchmark comparison.
B.1 Rank Sensitivity Analysis
Rank sensitivity is estimated by regressing evaluation score on log2 LoRA rank, with the slope summarizing performance changes as rank increases.
- s(r) = α + β log2 r models evaluation score as a function of LoRA rank.Here, s(r) is the observed score at rank r.
- Positive β indicates consistent gains from increasing rank, whereas near-zero or negative β suggests diminishing returns.The regression is applied to individual experts and combined mixture-of-experts models.
B.2 Typical Peak Rank Per Task
Peak performance generally occurs at moderate ranks, but the preferred rank varies substantially by benchmark type: knowledge-oriented tasks peak earlier, while reasoning-heavy tasks peak later.
- Peak rank is selected directly from observed scores for each expert and evaluation group, choosing the lowest rank in case of ties.The analysis avoids regression or smoothing when identifying the peak.
- Figure 4 compares expert specialization across benchmarks: Code and Math favor BBH, while Creative Writing and Reddit are competitive on GEN5.MC9 favors rank-efficient generalist experts with broad cross-domain utility.
- Table 4 reports MC9 scores within 1.1 percentage points of the maximum across unique experts despite ranks varying by over three orders of magnitude.
- Rank sensitivity differs across evaluation groups and experts, with BBH showing strong positive sensitivity in FlexMoRE models while many expert-only effects are weak or negative.The per-expert median sensitivity is −0.0004, with values ranging from −0.0014 to 0.0028.
- Median peak rank for aggregated average performance is log2 r = 9, with IQR [6.75, 10.50], corresponding to ranks ≈27–210.
- Knowledge-oriented benchmarks peak earlier, with MMLU at median log2 r = 2 and GEN5 at median log2 r = 5.Their interquartile ranges are [1.25, 8.00] for MMLU and [4.25, 5.75] for GEN5.
- Reasoning-heavy BBH peaks at a substantially higher median log2 r = 11.5, with IQR [7.25, 12.00].