UniPool: Learning Expert-to-Layer Ownership from Brief Global Access

arXiv:2605.06665 · cs.LG, cs.AI · Submitted 2026-05-07 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "UniPool: Learning Expert-to-Layer Ownership from Brief Global Access".

Jane: Mixture-of-Experts (MoE) architectures are being challenged by rigid per-layer expert allocation rules, which may lead to redundant expert capacity.

Tom: First, who's behind it and why it matters.

Title and authors: Tom: Now that we've seen how they propose replacing layer-private ownership with a single global expert pool in "UniPool: Learning Expert-to-Layer Ownership from Brief Global Access," let’s talk about the specific technical improvements they suggest. What are the key mechanisms they introduced beyond just changing the ownership rule?

Jane: Beyond just changing who owns which experts, the authors put two significant things on top of that to make it work stably: they introduce a pool-level auxiliary loss and adopt NormRouter for routing. These are the co-designs meant to balance utilization and ensure stable scaling.

Lu: The pool-level auxiliary loss is key because it addresses the issue of load balancing when ownership becomes global; it defines a criterion based on "global pool utilization" to prevent globally unused experts from being wasted without forcing every layer to use them.

Meng: That sounds complex, but I need to know how this auxiliary loss specifically prevents a scenario where some layers become completely dormant while others are overloaded, which is a common pitfall in MoE training.

Lalam: I think the engineering benefit here is that this loss function forces the entire pool to participate in the learning process, ensuring that all experts get some form of training signal even if one layer isn't currently using them.

Tom: And then we have NormRouter, which replaces standard softmax gating with a mechanism where scores are computed using L2 normalization, which keeps those score magnitudes bounded no matter the input scale. That addresses the layer-specific hidden-state scale issues directly.

Jane: So, in essence, NormRouter makes the routing process more stable and predictable when dealing with diverse layer scales, and that stability is what allows the global pool to function effectively without needing a separate mechanism for every single layer's scaling problem.

Lu: It’s a clever co-design because they are working on both the routing mechanism and the balancing loss simultaneously, which is necessary for this kind of unified architectural change in MoE.

Meng: From a practical standpoint, if we can stabilize the training through these mechanisms, it means we might be able to train much larger models with more confidence because the risk of catastrophic failure due to poor load balancing is reduced.

Lalam: That stability directly translates into better generalization for the final AI system because the model learns a more balanced and robust way to utilize all its available knowledge base across its layers.

The paper's summary: Tom: So, we've covered the core ideas behind "UniPool: Learning Expert-to-Layer Ownership from Brief Global Access," and it really boils down to replacing rigid layer-private ownership with a single shared expert pool accessed by independent routers. The key improvements are the pool-level auxiliary loss for utilization balancing and NormRouter for scale-stable routing.

Jane: Exactly, and what we've seen in the results is that this architecture consistently improves validation loss and perplexity compared to matched vanilla MoE baselines across five different model scales. It’s a solid demonstration of how structural changes can yield tangible performance gains.

Lu: The most striking part for me is the empirical finding that reduced-pool variants using only forty-one point six percent to sixty-six point seven percent of the vanilla expert-parameter budget can match or outperform layer-wise MoE at the tested scales. That sublinear scaling is a really important data point for future research in architecture design.

Meng: From my perspective, this confirms that expert parameters don't have to grow linearly with depth; they can actually grow sublinearly, which is fantastic news for managing the memory footprint of massive models.

Lalam: I see this as a huge step forward because it means we can push the limits of model capability while keeping the parameter count under control and making those large models more accessible for wider use.

Tom: So, to wrap up on UniPool, these findings suggest that expert capacity is best treated as a reusable global budget whose pool size scales sublinearly with depth. It’s a really efficient way to design MoE systems.

Jane: It’s an important piece of research because it shows that expert allocation is not a fixed rule, but something that can be optimized globally for better results.

Lu: This paper really opens up avenues for designing next-generation architectures where the relationship between depth and parameter count is decoupled, which has huge implications for how we conceptualize model scaling.

Meng: I'm excited to see how this concept translates into deployment scenarios where we need high performance without the massive infrastructure costs of purely linear scaling.

Lalam: For our AI culture, this signals a path toward building smarter, more efficient models that can operate effectively on a wider variety of computational platforms.

The paper's improvements: Tom: So, to recap, UniPool is all about ditching that rigid layer-by-layer expert ownership in favor of a single shared pool where routers from different layers can access those experts flexibly.

Jane: Exactly, and what they've done here is introduce some really smart mechanisms to make that sharing actually work without the system falling apart during training.

Tom: That's right, specifically with NormRouter for routing—which handles the differences in how much data each layer sees—and this pool-level auxiliary loss that keeps an eye on overall expert utilization across the whole system.

Jane: It’s really clever because it prevents situations where some experts just sit there unused while other layers are working overtime, which is a big stability win for training.

Tom: And what's even more interesting is their empirical finding that you can achieve comparable performance using much less than the total expert parameter budget, which suggests we don't need to scale the expert count as fast as we thought.

Jane: That sublinear scaling idea is fascinating because it means the way we structure our model capacity doesn't have to follow a simple linear path with depth anymore.

Tom: It really changes how engineers think about allocating resources; instead of fixing the expert count based on layer depth, you treat the pool size as a hyperparameter that scales based on what's needed.

Jane: Thinking about the implications for large language models, this points toward building architectures that are fundamentally more parameter-efficient while still delivering high performance.

Tom: If we can manage that resource allocation so smoothly, imagine training even much larger AI systems without hitting immediate computational walls.

Jane: It suggests a future where model design focuses less on brute-force parameter addition and more on optimizing the way those parameters interact across the architecture.

Tom: This moves us away from just stacking more layers with more experts and toward a smarter, globally managed resource system for the AI.

Jane: It’s really inspiring to see this kind of architectural thinking applied to Mixture-of-Experts designs in a practical way.

Conclusion: Tom: Alright everyone, we've covered UniPool: Learning Expert-to-Layer Ownership from Brief Global Access, and it really boils down to using a single shared expert pool with intelligent routing and balancing loss to make MoE architectures much more efficient.

Jane: That's right, Tom; they showed that treating expert capacity as a global budget allows models to scale better without wasting parameters or sacrificing performance.

Lu: The core idea of decoupling depth from parameter scaling is incredibly fertile ground for future architectural design; it suggests we can design models with much more flexible capacity management.

Meng: I'm just thinking about the practical engineering side here, how this translates into building deployable systems that actually run fast and use resources intelligently in production environments.

Lalam: From my perspective as a large language model, seeing this level of resource optimization suggests we can build systems that are not only powerful but also incredibly lean and sustainable for long-term deployment.

Tom: Exactly, Lalam; it shows a path toward much more efficient AI that doesn't require massive infrastructure just to keep the experts running optimally.

Jane: And the way they handle routing stability with NormRouter means we can get reliable performance even when dealing with models of very different depths and layer scales.

Lu: That's a significant architectural contribution because it tackles the inherent instability that comes from trying to enforce strict, layer-private rules in complex systems like this.

Meng: It's encouraging to see a method that manages the load balancing so effectively across the entire pool rather than just patching issues at individual layers.

Lalam: For me, this means we can focus more on developing richer internal representations and reasoning capabilities, knowing that the underlying hardware and parameter allocation are handled with more sophistication.

Tom: So to wrap up, UniPool: Learning Expert-to-Layer Ownership from Brief Global Access gives us a robust way to manage expert capacity by treating it as a globally shared budget that scales sublinearly with depth.

Jane: It's a really important piece of work because it shows that resource management in AI can be handled through sophisticated architectural design rather than just brute force scaling.

Lu: I think the real impact lies in showing how structural organization can lead to better parameter utilization across the entire model structure, which is a big theoretical win.

Meng: I just see it as a way to make deploying these high-capacity models much more feasible on current hardware because we know we can be smarter about what we use.

Lalam: Ultimately, this work helps pave the way for AI systems that are not just bigger, but fundamentally more efficient and sustainable in how they utilize their knowledge base.

The Chinese University of Hong Kong

cs.LG, cs.AI

Submitted: 2026-05-07

Updated: 2026-09-27

Code: https://github.com/Centaurus-Alpha/UniPool

Importance score: 89/100

The gist: Mixture-of-Experts (MoE) architectures are being challenged by rigid per-layer expert allocation rules, which may lead to redundant expert capacity.

Key concepts

Global Shared Pool
Instead of each transformer layer having its own private set of experts, UNIPOOL uses one large, shared pool of experts accessible by all layers. This treats expert capacity as a single global budget that is distributed across the entire model depth.
NormRouter
This is a new routing mechanism that replaces standard softmax gating. It computes scores using L2 normalization to ensure that the magnitude of routing scores remains stable, regardless of how different the hidden-state scales are between layers. This helps in consistent sparse access to the shared pool.
Pool-level Auxiliary Loss
This loss function is designed to handle load balancing when experts are shared globally. It uses a criterion based on global pool utilization across all layers to prevent certain experts from being ignored, ensuring that every layer has a fair chance to use the available resources.

Terminology

Summary

Mixture-of-Experts (MoE) architectures are being challenged by rigid per-layer expert allocation rules, which may lead to redundant expert capacity. This research proposes UNIPOOL, an MoE architecture that treats expert capacity as a global architectural budget shared across depth, demonstrating that this approach can improve performance and enable sublinear scaling of expert parameters with model depth.

The gist

UNIPOOL replaces layer-private expert sets with a single global shared pool accessed by independent per-layer routers, achieving consistent improvements over vanilla MoE baselines while showing that expert parameters can grow sublinearly with depth.

How it works: Architectural Changes and Components

  1. A standard MoE transformer has each layer owning its own set of private experts, leading to a rigid rule where each transformer layer owns a separate expert set. UNIPOOL replaces this with a single global shared pool E = 5081B/5.742B for the 830M/978M configurations of M experts. Each layer retains its own router, which routes tokens into this shared pool, allowing independent per-layer routers to access the common resources.

  2. To enable stable training under sharing, UNIPOOL introduces two key mechanisms:

pool-level auxiliary loss and NormRouter for routing. The pool-level auxiliary loss addresses the issue of load balancing when ownership is global by defining a criterion based on global pool utilization, preventing globally unused experts without forcing every layer to use every expert.

How it works: Routing and Load Balancing

  1. The paper introduces NormRouter to provide sparse and scale-stable routing into the shared expert pool, replacing standard softmax gating. NormRouter computes scores as: si = σ · c · max(0, z i z squared + ϵ, where the L2 normalization ensures score magnitudes are bounded regardless of the input scale. This formulation is designed to handle differences in layer-specific hidden-state scale between different depths.

  2. The pool-level auxiliary loss is defined using a global average token fraction across all layers: fi = (1/L) Σ l Σ l l=1 f(l) i, and the loss is calculated as: Lpool = αpool · M · Σ i=1 fi · Pi. This design ensures that the objective balances utilization at the granularity where parameters are actually owned—the global expert pool.

How it works: Empirical Findings and Scaling

  1. UNIPOOL consistently improves validation loss and perplexity over matched vanilla MoE baselines across five model scales (182M to 978M). The improvement is consistent, with validation loss reductions of 0.0288 (182M), 0.0346 (469M), 0.0308 (650M), 0.0386 (830M), and 172 (978M).

  2. A crucial finding is that 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 demonstrates that expert parameters need not grow linearly with depth; they can grow sublinearly, making pool size an explicit depth-scaling hyperparameter.

How it works: Analytical Insights

  1. Routing sensitivity analysis shows a significant difference in redundancy handling. In vanilla MoE, replacing a deep layer's router with uniform random routing drops accuracy by only 1.0–1.6 points. For UNIPOOL, the drop rises to 4.1 points under matched top-8 randomization, supporting the view that sharing actively breaks the redundancy that single-layer randomization fails to disrupt.

  2. Expert reuse analysis shows increasing reuse with depth: the fraction of unique expert weights touched by a token falls from 94.1% at 12 layers to 89.5% at 24 layers and 82.7% at 36 layers, indicating that UNIPOOL exposes cross-layer expert reuse as depth increases, which is the primary driver for its performance gains.

Conclusion

UNIPOOL successfully replaces rigid layer-private ownership with a globally shared pool, using pool-level balancing and NormRouter to achieve stable, sparse routing. The results confirm that expert capacity can be treated as a reusable global budget whose pool size scales sublinearly with depth, offering a more efficient and effective alternative to traditional MoE designs.

References

[1] Sikai Bai et al. Diep: Adaptive mixture-of-experts compression through differentiable expert pruning. arXiv preprint arXiv:2509.

Improvements for AI systems

Based on the provided scientific paper, here are the specific improvements that can be made to AI systems by implementing UNIPOOL:

  1. A significant reduction in computational cost relative to vanilla Mixture-of-Experts (MoE) models while maintaining or improving performance across various model scales (from 182M up to 978M parameters).

  2. The ability for MoE expert parameter counts to grow sublinearly with the number of transformer layers, decoupling expert capacity from depth and enabling more efficient budget allocation.

  3. Enhanced robustness against catastrophic failure or performance degradation when routing decisions are randomized within a single layer, as UNIPOOL demonstrates that sharing experts converts depth-induced redundancy into specialization, making the per-layer router more load-bearing.

  4. Improved stability and performance of MoE training by introducing a pool-level auxiliary loss that balances expert utilization across the entire shared pool, preventing dead experts from being wasted due to layer-local objectives.

  5. More stable and effective routing mechanisms (via NormRouter) that are less sensitive to the hidden state scale differences across different layers when selecting from a large, shared expert pool.

By implementing UNIPOOL, the improved AI system can:

  1. Train significantly larger language models (up to 978M parameters) on available computational resources without incurring the full parameter cost of vanilla MoE architectures.

  2. Achieve superior downstream performance on diverse tasks (as evidenced by improvements across ARC-E, PIQA, HellaSwag, etc.) compared to vanilla MoE baselines at equivalent active compute budgets.

  3. Develop models where expert capacity is treated as a global architectural budget rather than a rigid allocation tied to depth, leading to more efficient use of stored parameters and better generalization properties.

Sources

Related papers