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

Paper: 2605.06665 Authors: Minbin Huang, Han Shi, Chuanyang Zheng, Yimeng Wu, Guoxuan Chen, Xingtong Yu, Yichun Yin, Hong Cheng Categories: cs.LG Code: https://github.com/Centaurus-Alpha/UniPool

Abstract

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, this allocation rule can be wasteful—experts at different layers cannot be shared even when they learn similar transformations. This paper introduces 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. Across five model scales trained on 30B tokens, UniPool consistently improves validation loss and enables sublinear expert scaling—using only 41.6%–66.7% of vanilla MoE’s expert parameters while matching or exceeding performance.

Key Contributions

• Global Expert Pool Architecture: Replaces the conventional per-layer expert ownership with a single shared pool, enabling cross-layer expert reuse while preserving layer-specific routing capabilities.

• Pool-Level Balancing: Introduces a novel auxiliary loss that balances expert utilization at the global pool level rather than per-layer, preventing globally dead experts while allowing layers to specialize on different expert subsets.

• NormRouter Integration: Adopts NormRouter (L2-normalize-then-ReLU) for sparse and scale-stable routing into the shared expert pool, addressing stability challenges that arise from multiple layers accessing the same expert budget.

• Sublinear Expert Scaling: Demonstrates that expert parameters need not grow linearly with depth—reduced-pool variants using 41.6%–66.7% of vanilla MoE’s expert budget match or outperform layer-wise MoE across tested scales.

• Consistent Performance Gains: Achieves up to 0.0386 validation loss reduction across five LLaMA-architecture model scales (182M to 978M parameters) trained on 30B tokens from the Pile dataset.

Methodology

The Problem with Vanilla MoE

Traditional MoE architectures follow a rigid allocation rule:

• Each transformer layer owns a private set of expert FFNs
• A layer-specific router selects a sparse subset of those private experts
• This creates a linear relationship between depth and total expert parameters
• Experts cannot be shared across layers, even when learning similar transformations

The authors’ routing-randomization probe on production MoE models revealed significant redundancy: replacing a deep layer’s learned router with uniform random routing drops accuracy by only 1.0–1.6 points, suggesting that standard MoE training may duplicate expert functions across layers.

UniPool Architecture

UniPool fundamentally rethinks expert allocation:

1. Single Shared Pool: All layers access the same global expert pool instead of maintaining private expert sets
2. Independent Per-Layer Routers: Each layer retains its own router to select experts from the shared pool
3. Pool-Level Load Balancing: Aggregates token-to-expert assignments across all layers and applies a single balancing objective over the shared pool
4. NormRouter for Stability: Uses L2-normalize-then-ReLU scoring with learnable scaling factors to maintain stable routing across different depths

Key Design Decisions

Why Pool-Level Balancing?

• Layer-local balancing (forcing every layer to use every expert) conflicts with cross-layer reuse
• Pool-level balancing prevents globally dead experts while allowing layer specialization
• Aggregates utilization statistics across layers for a single global objective

Why NormRouter?

• Softmax-based routers designed for layer-specific experts become unstable with shared pools
• L2 normalization makes scores less sensitive to layer-specific hidden-state scales
• ReLU induces sparse competition over the large pool
• Learnable scaling factors let each router adjust routing strength during training

Results

Performance Improvements

Across five model scales (182M, 469M, 650M, 830M, 978M parameters):

• Consistent validation loss reduction: Up to 0.0386 improvement over vanilla MoE
• Better perplexity: Improvements across all tested scales
• Training stability: No collapse or instability issues with shared pool design

Sublinear Expert Scaling

The most striking finding is that UniPool enables efficient parameter reduction:

Expert Budget	Performance vs Vanilla MoE
66.7%	Matches or exceeds vanilla MoE
50%	Competitive performance
41.6%	Still matches vanilla at some scales

This demonstrates that under a shared-pool design, expert parameters can grow sublinearly with depth while remaining more efficient than vanilla MoE.

Pool Size as Hyperparameter

UniPool identifies pool size as an explicit depth-scaling hyperparameter:

• Larger pools provide more capacity but may be underutilized
• Smaller pools force more efficient expert reuse
• Optimal pool size depends on model depth and task requirements

Composability with Expert Decomposition

Further analysis shows UniPool’s benefits compose with finer-grained expert decomposition, suggesting the approach is compatible with other MoE optimization techniques.

Takeaways

• Rethinking Expert Allocation: The conventional per-layer expert ownership is a rigid design choice that may be suboptimal. Treating expert capacity as a global budget enables more efficient parameter usage.

• Redundancy in Deep Layers: The routing-randomization probe reveals that deeper layers in vanilla MoE may have substantial redundancy, suggesting opportunities for parameter sharing.

• Sublinear Scaling Unlocked: UniPool demonstrates that expert parameters need not scale linearly with model depth—a finding with significant implications for scaling large MoE models efficiently.

• Practical Efficiency Gains: Using only 41.6%–66.7% of vanilla MoE’s expert parameters while matching performance translates to substantial memory and computational savings in production deployments.

• New Design Space: Pool-level balancing and shared expert pools open a new design space for MoE architectures, with pool size emerging as a key hyperparameter for depth scaling.

• Production Viability: The approach maintains training stability and shows consistent improvements across multiple scales, suggesting readiness for production deployment.