UniPool: Breaking the Per-Layer Expert Ownership Rule in Mixture-of-Experts

Paper: 2605.06665 Authors: Minbin Huang, Han Shi, Chuanyang Zheng, Yimeng Wu, Guoxuan Chen, Xintong Yu, Yichun Yin, Hong Cheng Categories: cs.LG, cs.AI

The Gap

Mixture-of-Experts (MoE) models follow a rigid allocation rule: each transformer layer owns a separate set of experts. This means if you have 32 layers and 8 experts per layer, you need 256 expert networks total. The field assumed this isolation was necessary—that each layer needs its own specialized capacity to handle different abstraction levels.

But the authors ran a simple probe: they replaced learned routers in deeper layers with uniform random routing. The accuracy drop? Only 1.0-1.6 points across production models. This suggests massive redundancy—layers aren’t using their exclusive expert capacity as efficiently as assumed. The gap: expert parameters scale linearly with depth, but the actual need for isolated per-layer capacity is questionable.

Problem: Per-layer expert ownership
   |
   v
Assumption: Each layer needs isolated expert capacity
   |
   v
Probe: Replace deep-layer routers with random routing
   |
   v
Evidence: Only 1.0-1.6 point accuracy drop
   |
   v
Conclusion: Expert capacity is redundantly allocated
   |
   v
Method: Share a global expert pool across all layers
   |
   v
Result: Better performance with 41.6%-66.7% fewer experts

The Increment

One sentence: Before this paper, expert parameters grew linearly with model depth because each layer owned separate experts; after, expert capacity becomes a global budget that can grow sublinearly while outperforming the linear baseline.

Core Mechanism

UniPool replaces per-layer expert ownership with a single shared pool. Instead of Layer 1 having Experts 1-8, Layer 2 having Experts 9-16, etc., all layers route into the same pool of experts. Each layer still has its own learned router that decides which experts to activate, but the experts themselves are shared infrastructure.

The challenge is training stability. When multiple layers compete for the same experts, you can get collapse—some experts become overused while others are ignored. UniPool addresses this with two components: (1) a pool-level auxiliary loss that balances expert utilization across the entire pool rather than per-layer, and (2) NormRouter, which provides sparse and scale-stable routing by normalizing router logits before the top-k selection.

Data flows like this: a token enters a layer, the layer’s router computes affinity scores for all experts in the pool, top-k experts are selected, the token is processed by those experts, and outputs are combined. The next layer repeats this process with the same expert pool but its own independent router. The pool-level loss monitors which experts are being used across all layers and penalizes imbalance.

Token Input
    |
    v
Layer 1 Router ---> [Expert Pool: E1 E2 E3 ... En]
    |                      ^  ^  ^
    +----------------------+  |  |
    |                         |  |
    v                         |  |
Layer 2 Router ---------------+  |
    |                            |
    +----------------------------+
    |
    v
Layer 3 Router
    |
    v
Output

Pool-Level Loss: monitors expert usage across all layers
NormRouter: normalizes logits before top-k selection

Think of it like a city’s public transportation system versus private cars. In the old MoE design, each neighborhood (layer) owns its own fleet of buses (experts). If you live in Neighborhood 5, you can only use buses from Fleet 5. This means every neighborhood needs to maintain a full fleet, even if some buses sit idle while others are packed.

UniPool is like converting to a citywide bus system. All buses belong to a shared pool. Neighborhood 5’s dispatch center (router) still decides which buses to call based on local demand, but those buses come from the shared fleet. Neighborhood 12’s dispatch center makes independent decisions, but draws from the same pool. The pool-level loss is like a central coordinator ensuring no bus is perpetually idle or overworked. NormRouter is the standardized dispatch protocol that prevents any single neighborhood from monopolizing resources.

The key insight: neighborhoods don’t need exclusive bus ownership—they need routing authority. Sharing the fleet reduces total buses needed while improving service quality because high-demand routes can access more capacity.

Key Concepts

• Expert Pool as Architectural Budget: Traditional MoE treats expert count as a per-layer decision: “Each layer gets 8 experts.” UniPool reframes this as a global budget: “The entire model gets 64 experts, allocated across 8 layers.” This decouples expert capacity from depth. You can have a 32-layer model with only 64 experts instead of 256. The pool size becomes an explicit hyperparameter you tune independently of depth, like choosing how many buses a city needs regardless of how many neighborhoods it has.

• Pool-Level Auxiliary Loss: Standard MoE uses per-layer load balancing—Layer 3 tries to balance its own 8 experts. But with sharing, this breaks down. If Layer 3 balances its usage but Layer 7 ignores half the pool, you still have collapse. Pool-level loss monitors expert utilization across all layers simultaneously. It’s computed by tracking how often each expert is selected by any layer, then penalizing deviation from uniform usage. This prevents the scenario where Layers 1-10 love Expert 5 while Layers 11-20 never touch it.

• Sublinear Expert Scaling: The paper’s most striking finding. In traditional MoE, doubling depth means doubling expert parameters. UniPool shows you can grow expert count sublinearly—41.6% to 66.7% of the vanilla budget—and still outperform. Why? Because layers share capacity, redundancy is eliminated. It’s like discovering that a city with twice the population doesn’t need twice the buses if you optimize routes and sharing. The relationship between depth and expert capacity is revealed to be flexible, not fixed.

Framework Shift

Before (Vanilla MoE):                After (UniPool):

Layer 1: [E1 E2 E3 E4]              Layer 1 Router
         |  |  |  |                      |
Layer 2: [E5 E6 E7 E8]              Layer 2 Router ----+
         |  |  |  |                      |             |
Layer 3: [E9 E10 E11 E12]           Layer 3 Router    |
         |  |  |  |                      |             v
                                         +-----> [E1 E2 E3 E4 E5 E6]
Experts = Layers * Experts/Layer                 (Shared Pool)
(Linear scaling)                     
                                     Experts = Independent Budget
                                     (Sublinear scaling possible)

From isolated per-layer ownership to centralized capacity with distributed routing, the core shift is treating expert parameters as a global resource rather than a per-layer allocation.

Expert Assessment

Problem choice: This is a real gap. The field inherited the per-layer expert convention from early MoE designs without questioning it. The random routing probe is simple but devastating—it directly challenges a foundational assumption. The problem sits at a critical juncture: MoE models are scaling up, and parameter efficiency is becoming a bottleneck. Questioning the linear scaling assumption is timely and important.

Method maturity: The solution is elegant—share the pool, balance globally, use stable routing. But there’s a missed opportunity: the paper doesn’t deeply explore why random routing works so well in deeper layers. Is it because deeper layers need less specialization? Because the routing learned patterns are redundant? A mechanistic analysis of what the shared pool learns versus per-layer experts would strengthen the contribution. The method feels like 80% insight, 20% unexplored depth.

Experimental integrity: Baselines are fair—matched parameter counts, same training data, consistent architecture. The experiments span five model scales (182M to 978M parameters) and show consistent improvements. The reduced-pool variants (41.6%-66.7% of vanilla budget) are particularly convincing because they demonstrate efficiency gains, not just performance gains. One concern: all experiments use LLaMA architecture on the Pile dataset. Generalization to other architectures (T5, GPT) and domains (code, multilingual) is untested. The 30B token training budget is also modest by modern standards.

Writing quality: The abstract and introduction are strong—they clearly state the problem and contribution. The method section is dense but complete. The weakness is in the analysis section: the paper shows that UniPool works but doesn’t deeply explain why. The routing visualization (Figure 4) hints at interesting patterns but isn’t thoroughly interpreted. Rewriting Section 5 to include mechanistic analysis—what do shared experts specialize in? how do layer routers differentiate?—would elevate this from a solid empirical paper to a foundational one.

Verdict: weak accept — Challenges a fundamental assumption with strong empirical evidence, but lacks the mechanistic depth to fully explain why sharing works better than isolation.

Takeaways

Decouple resource allocation from architectural depth: The core idea transfers beyond MoE. Anytime you have a resource that scales linearly with depth (attention heads, hidden dimensions, expert sets), ask: does each layer need exclusive ownership, or can we share a global pool with per-layer routing? This framing applies to neural architecture search, federated learning (shared model pool across clients), and even database query optimization (shared execution plan pool).

Global balancing for shared resources: When multiple agents (layers, clients, processes) compete for a shared resource, per-agent balancing fails. You need a global loss or coordinator that monitors aggregate usage. This pattern appears in load balancing, resource scheduling, and multi-agent systems. The pool-level auxiliary loss is a clean implementation of this principle.

Probe assumptions with minimal interventions: The random routing experiment is a masterclass in assumption testing. Instead of building a complex alternative, they made the smallest possible change (replace learned routing with random) and measured the gap. If the gap is small, the assumption is weak. This technique—minimal intervention probing—is underused in ML research. Apply it to other conventions: do we need layer normalization at every layer? do we need separate embeddings for encoder and decoder?