PaperlayerView exact revision on arXiv

Source-linked AI summary

UniPool: A Globally Shared Expert Pool for Mixture-of-Experts

Minbin Huang, Han Shi, Chuanyang Zheng, Yimeng Wu, Guoxuan Chen, Xintong Yu, Yichun Yin, Hong Cheng

arXiv:2605.06665v1cs.LGcs.AI

TL;DR

Mixture-of-Experts models typically assign separate expert sets to each layer, potentially wasting capacity that could be shared. UniPool replaces layer-private experts with a globally shared pool and shows consistent validation improvements over matched vanilla MoE baselines, including comparable or better results with 41.6%–66.7% of the expert-parameter budget.

  • Problem

    Per-layer expert ownership prevents cross-layer reuse, despite evidence that experts at different layers can learn similar transformations.

  • Method

    UniPool uses one globally shared expert pool with independent per-layer routers, pool-level utilization balancing, and NormRouter for sparse, scale-stable routing.

  • Results

    41.6%–66.7% of vanilla MoE’s expert-parameter budget suffices for reduced-pool UniPool variants to match or outperform layer-wise MoE across five model scales.

  • Takeaways & Limitations

    UniPool supports treating expert capacity as a reusable global budget whose pool size can scale sublinearly with depth.

  • Takeaways & Limitations

    The study does not report wall-clock throughput comparisons, leaving expert-dispatch efficiency under larger shared pools for future work.

Abstract

from arXiv · show

Modern Mixture-of-Experts (MoE) architectures allocate expert capacity through a rigid per-layer rule: each transformer layer owns a separate expert set. This convention couples depth scaling with linear expert-parameter growth and assumes that every layer needs isolated expert capacity. However, recent analyses and our routing probe challenge this allocation rule: replacing a deeper layer's learned top-k router with uniform random routing drops downstream accuracy by only 1.0-1.6 points across multiple production MoE models. Motivated by this redundancy, we propose UniPool, an MoE architecture that treats expert capacity as a global architectural budget by replacing per-layer expert ownership with a single shared pool accessed by independent per-layer routers. To enable stable and balanced training under sharing, we introduce a pool-level auxiliary loss that balances expert utilization across the entire pool, and adopt NormRouter to provide sparse and scale-stable routing into the shared expert pool. Across five LLaMA-architecture model scales (182M, 469M, 650M, 830M, and 978M parameters) trained on 30B tokens from the Pile, UniPool consistently improves validation loss and perplexity over the matched vanilla MoE baselines. Across these scales, UniPool reduces validation loss by up to 0.0386 relative to vanilla MoE. Beyond raw loss improvement, our results identify pool size as an explicit depth-scaling hyperparameter: reduced-pool UniPool variants using only 41.6%-66.7% of the vanilla expert-parameter budget match or outperform layer-wise MoE at the tested scales. This shows that, under a shared-pool design, expert parameters need not grow linearly with depth; they can grow sublinearly while remaining more efficient and effective than vanilla MoE. Further analysis shows that UniPool's benefits compose with finer-grained expert decomposition.

1 Introduction

UniPool replaces layer-private expert ownership with a globally shared expert pool accessed by independent per-layer routers, motivated by evidence that deeper layer-specific experts can be redundant. It combines pool-level balancing with NormRouter and achieves consistent gains over vanilla MoE while enabling sublinear expert scaling.

  • Architecture: UniPool replaces layer-private expert sets with a single shared expert pool accessed by independent per-layer routers, enabling cross-layer expert reuse while preserving layer-specific routing.The global pool decouples expert ownership from transformer depth.
  • Motivation: 1.0–1.6 points: Uniform random routing in a single deep-half MoE layer reduces downstream accuracy by only 1.0–1.6 points across three production MoE models.This probe suggests that deeper layer-private experts may duplicate functions and that standard expert capacity can be wastefully tied to individual layers.
  • Balancing: UniPool introduces a pool-level auxiliary loss that balances utilization across the shared pool rather than enforcing layer-local usage.An expert unused by one layer may still be frequently selected by others, so layer-local deadness is misaligned with global parameter allocation.
  • Routing: NormRouter provides sparse, scale-stable routing into the shared pool through normalization, ReLU-based competition, and a learnable routing scale.These mechanisms address layer-dependent logit scales and inconsistent routing sharpness across depths.

2 Related Work

Prior work established sparse MoE as an effective scaling strategy and developed routing methods for expert selection and balance. UniPool builds on parameter-sharing research by selectively sharing expert FFNs while retaining layer-specific attention blocks and routers.

  • Sparse MoE and scaling: Sparse MoE evolved from sparsely gated expert layers through top-1 routing, expert-parallel training, and router-stability improvements.Large-scale systems including Mixtral and the DeepSeek series further demonstrate sparse expert capacity as an effective way to scale language models.
  • Routing and load balancing: Standard MoE routing uses softmax gating with the Switch auxiliary loss to balance per-expert token fractions and routing probabilities within each layer.Alternative designs include expert choice, BASE’s linear assignment, deterministic hash routing, and sigmoid gating.
  • Parameter sharing and expert reuse: UniPool selectively shares MoE expert FFNs across depth while retaining layer-specific attention blocks and routers, distinguishing it from broader parameter sharing in Universal Transformers and ALBERT.MoEUT provides a closer precedent by cyclically repeating a small group of shared transformer blocks across depth with per-layer entropy-based mechanisms.

3 Motivating Observation: Expert Redundancy in Deep MoE Layers

Analyses of trained MoEs reveal substantial redundancy among same-layer experts, especially in deep layers, while routing probes show that expert choice there carries limited local information. This motivates treating deep-layer expert capacity as less dependent on sharply isolated per-layer specializations.

  • Expert redundancy: Above 0.9 pairwise cosine similarity shows that same-layer expert weight matrices in Qwen and DeepSeek MoEs share a dominant subspace.Related analyses also report up to 2× decoding speedups after similarity-based rerouting and roughly 8% relative quality loss when pruning about half of Mixtral 8×7B’s experts.
  • Motivation: Prior work treats expert redundancy as a post-hoc compression target while retaining per-layer expert ownership.UniPool’s motivation instead uses routing redundancy to question whether isolated expert capacity is necessary at every layer.
  • Routing redundancy: 1.0–1.6 points is the accuracy drop from randomizing routing in a deep-half layer across three production MoE models.The probe tests whether deep-layer experts have distinct specializations by replacing learned routing decisions with randomization.
  • Routing redundancy: Limited local information in deep-layer expert choice indicates that per-layer routers are not forming sharply distinct functional partitions over private expert sets.This routing result aligns with evidence that same-layer expert parameters and outputs are highly similar, with strongest similarity concentrated in deep layers.

4 Method

UniPool replaces layer-private expert sets with a globally shared expert pool accessed by per-layer routers, and stabilizes shared routing with a pool-level auxiliary loss and NormRouter.

  • 4.1 Shared Expert Pool: UniPool replaces L separate expert sets with one global pool of M expert FFNs while retaining an independent router at each layer.The same expert module is reused regardless of which layer invokes it; pool size M is configurable and can match vanilla MoE’s expert-parameter budget while preserving dense-equivalent active FFN compute.
  • 4.2 Pool Auxiliary Loss: Under shared parameters, layer-local load balancing can misclassify experts as dead because an expert unused by one layer may be frequently used by others.The mismatch arises because expert ownership is global rather than tied to a single layer.
  • 4.2 Pool Auxiliary Loss: The pool-level auxiliary loss balances expert utilization using token fractions and routing probabilities averaged across all L sharing layers.Because the global token fraction is shared across layers, the objective decomposes into independently computable per-layer contributions.
  • 4.3 NormRouter: NormRouter replaces softmax routing with L2-normalized scores whose magnitudes remain bounded regardless of input scale.This addresses differing hidden-state norms and logit scales across depths that could otherwise produce inconsistent routing sharpness into the shared pool.
  • 4.3 NormRouter: After scoring, UniPool selects the top-k experts, and NormRouter supports both per-layer and pool-level auxiliary losses.These losses operate directly on routing scores rather than softmax probabilities.

5 Experiments

Across five LLaMA-scale models trained on approximately 30B Pile tokens, UniPool consistently improves validation loss over dense and vanilla MoE baselines. Its gains persist with substantially smaller shared pools, finer-grained experts, and the pool-level routing and balancing design.

  • Main results: UniPool consistently outperforms dense and vanilla MoE baselines across five model scales from 182M to 978M parameters.Models use LLaMA-style transformer backbones and are trained on the Pile for approximately 30B tokens.
  • Main results: 0.0386 validation-loss reduction is the largest UniPool gain over vanilla MoE, with reductions of 0.0288, 0.0346, 0.0308, 0.0386, and 0.0172 at 182M, 469M, 650M, 830M, and 978M.Both sparse MoE methods substantially outperform the dense baseline; at 182M, the comparison is 1.9029 vs. 2.042.
  • Efficiency scaling: 41.6%-66.7% of vanilla expert parameters is sufficient for reduced-pool UniPool variants to beat the layer-private baseline at tested scales.The smallest winning pools retain 66.7% at 182M, 50% at 469M and 650M, and 41.6% at 830M, while keeping top-1 routing and matched active parameters.
  • Granularity scaling: UniPool outperforms matched vanilla MoE under all tested 182M configurations: 8E/top-1, 16E/top-2, and 32E/top-4.Both methods improve with larger expert counts, indicating that UniPool’s gain composes with finer-grained expert decomposition.
  • Ablations: A shared pool with the original per-layer auxiliary loss underperforms vanilla MoE at 1.9480 vs. 1.9317, whereas the pool-level auxiliary loss improves validation loss to 1.9180.NormRouter alone slightly worsens validation loss to 1.9375 vs. 1.9317; global sharing performs best among sharing-scope variants.

6 Analysis

The analysis shows that expert sharing makes routing decisions more important, while cross-layer reuse and pool-level balancing explain UniPool’s behavior and efficiency. Shared-pool utilization requires the pool-level auxiliary loss to avoid collapse.

  • Analysis overview: The analysis covers routing randomization, expert reuse and budget allocation, and expert utilization and routing diversity under the shared pool.These are presented as three analytical lenses on UniPool’s behavior in Sections 6.1–6.3.
  • Routing randomization: 4.1 points at both 469M and 978M: cardinality-matched top-8 randomization reduces UniPool average accuracy more than vanilla MoE’s 1.3/1.5-point drops.The intervention samples from each layer’s top-8 most-used shared experts in UniPool and randomizes one deep-half layer under matched settings.
  • Expert reuse and budget allocation: UniPool enables the same expert functions to be reused across depths instead of restricting each depth to its own layer-private expert set.This reframes expert allocation as routed compositions with cross-layer sharing under matched top-1 compute.
  • Expert utilization and routing diversity: Pool-level auxiliary loss is critical because per-layer auxiliary loss with softmax routing causes aggregate shared-pool traffic to collapse.Figure 3 compares configurations with the same global expert pool but different auxiliary-loss and router designs.

7 Conclusion

UNIPOOL replaces layer-private experts with a globally shared pool trained using pool-level balancing and NormRouter. Across five model scales, it improves over matched vanilla MoE baselines, including with substantially reduced expert-parameter budgets.

  • Architecture: UNIPOOL replaces layer-private expert ownership with a global shared pool trained using pool-level balancing and NormRouter.The architecture treats expert capacity as shared rather than layer-private.
  • Results: Across five model scales, UNIPOOL improves validation loss and perplexity over matched vanilla MoE baselines.The improvement is reported consistently across the evaluated scales.
  • Efficiency: 41.6%–66.7% of its expert-parameter budget is sufficient for reduced-pool variants to outperform vanilla MoE.These variants use only the stated fraction of the vanilla MoE expert-parameter budget.

A Limitations and Future Work

The study’s validation is limited by its sub-billion-parameter scale, lack of wall-clock throughput comparisons, and evaluation restricted to seven zero-shot benchmarks. Future work should test larger models, longer training, and broader few-shot evaluation while clarifying efficiency gains in reduced-pool settings.

  • Scale of experiments: 182M–978M parameters and 30B training tokens constrain validation; billion-parameter models with longer training horizons remain an important direction.The reported improvements span five scales, but larger-scale and longer-horizon validation is still needed.
  • Throughput and memory: No wall-clock throughput comparisons are reported, and matched UNIPOOL (M = 8L) has the same total expert FFN count as vanilla MoE.At matched capacity, sharing changes expert access by reference rather than reducing parameter count; storage and memory savings arise only with reduced pools.
  • Downstream evaluation: Seven zero-shot benchmarks are evaluated, while broader testing including few-shot settings would further strengthen the findings.The downstream evaluation is described in Section 5.3.

B Model and MoE Configurations

The evaluated models span five backbone scales and use SwiGLU FFNs with intermediate size 4 × H. Vanilla MoE and UniPool configurations are compared under matched expert-capacity budgets, including total expert FFNs and per-token expert FLOPs.

  • Backbone configurations: Five evaluation scales use dense-width FFNs, or expert FFNs for MoE variants, with intermediate size 4 × H and SwiGLU activation.The scales are documented in the backbone configurations table.
  • Backbone configurations: “Active scale” denotes the dense-equivalent active parameter budget including embeddings, while “Total Params” reports stored UniPool parameters.MoE variants store additional expert parameters.
  • MoE configuration comparison: Vanilla MoE and UniPool are matched in total expert FFNs and per-token expert FLOPs.The comparison isolates the configuration difference under equal expert compute and capacity measures.

C Additional Training Curves … F Additional Routing-Randomization Details

The additional analyses show that UniPool’s training advantages persist throughout optimization, while shared routing enables cross-layer expert reuse and requires carefully matched randomization protocols. They also document the common training configuration and formalize how distinct experts are counted under vanilla MoE and UniPool.

  • C Additional Training Curves: Across 182M, 469M, and 650M scales, UNIPOOL remains below matched vanilla MoE validation loss after warmup, with a roughly 0.035 gap at 469M.At 182M the gap widens steadily, while at 650M UNIPOOL maintains a clear advantage throughout training.
  • C Additional Training Curves: Global sharing (G=1) has the lowest-loss trajectory, vanilla MoE (G=12) the highest-loss endpoint, and grouped sharing (G=2, 4, 6) generally interpolates between them.The ordering suggests broader sharing improves optimization trajectories rather than only final-checkpoint selection.
  • D Hyperparameter Details: Table 8 provides complete hyperparameter details for all experimental configurations.The models use RMSNorm, SwiGLU, RoPE, grouped-query attention with 4 KV heads, and untied input/output embeddings.
  • D Hyperparameter Details: Activation checkpointing with MoE layer recompute is enabled for the 469M, 650M, 830M, and 978M scales.Training uses Megatron-LM with sequence parallelism and a distributed optimizer.
  • E Distinct-Expert Accounting: In vanilla MoE, each layer owns disjoint experts, so a token touches exactly L unique expert tensors even when local expert indices repeat across layers.In UniPool, all layers route into a shared pool of M experts, and the token-level count is U(x)=|{e_l(x): l=1,...,L}| with 1≤U(x)≤L.
  • E Distinct-Expert Accounting: The reported metrics E_x[U(x)] and E_x[U(x)]/L quantify how much cross-layer expert reuse emerges in UniPool’s shared pool.These metrics are reported on the validation set.
  • F Additional Routing-Randomization Details: Routing-randomization results report per-task accuracy for TOP-K learned routing versus RANDOM after randomizing one deep-half MoE layer at a time and averaging across layers.Drops are measured relative to TOP-K, with Avg denoting the unweighted mean.
  • F Additional Routing-Randomization Details: For comparable randomization, vanilla MoE samples from 8 private experts, whereas UniPool samples from each layer’s held-out-validation top-8 shared experts; full-pool randomization is also reported.UniPool’s full shared pool has M=L×8 experts, so uniform sampling over all shared experts is not directly comparable to vanilla MoE.

G Pool Auxiliary Loss: Detailed Derivation · H NormRouter: Monte Carlo Initialization Details

The pool-level auxiliary loss decomposes into independently computable per-layer contributions, implemented with one-step-behind global token statistics to preserve activation checkpointing. NormRouter calibrates its initial score scale by Monte Carlo estimation based on expert count and top-k routing.

  • G Pool Auxiliary Loss: Detailed Derivation: The full derivation shows that the pool-level auxiliary loss decomposes into per-layer terms.The derivation starts from the pool-loss definition and rewrites its terms using f_i · P_i and f_i · P^(l).
  • G Pool Auxiliary Loss: Detailed Derivation: Each summand is a per-layer pool-loss contribution that can be computed independently because f_i does not depend on layer l.This independence enables layer-wise computation of the global pool objective.
  • G Pool Auxiliary Loss: Detailed Derivation: The one-step-behind scheme uses f_i from the previous micro-batch because current global token statistics require the completed forward pass.This avoids cross-layer tensor dependencies that would interfere with activation checkpointing.
  • G Pool Auxiliary Loss: Detailed Derivation: The global token distribution f_i is accumulated without gradients and updated after all layers complete their forward pass.Only the routing probabilities P^(l) carry gradients through the pool-loss path.
  • H NormRouter: Monte Carlo Initialization Details: NormRouter selects calibration factor c so initial selected scores have approximately unit magnitude for E experts and top-k routing.The calibration targets the expected norm of the selected scores.
  • H NormRouter: Monte Carlo Initialization Details: Algorithm 1 estimates c by Monte Carlo sampling at initialization with N = 10^5 samples.The procedure samples normalized Gaussian vectors, applies ReLU, sorts components, selects top-k values, and averages their inverse norm.
  • H NormRouter: Monte Carlo Initialization Details: For each sample, Algorithm 1 computes y = ReLU(x/∥x∥_2), extracts the top-k components, and appends their inverse L2 norm before averaging.The sampled x follows a standard E-dimensional Gaussian distribution.
Loading 2605.06665v1…