Omni-Router: Sharing Routing Decisions in Sparse Mixture-of-Experts for Speech Recognition

TL;DR

The paper tackles the limited cross-layer coordination in sparse MoE transformers for ASR by introducing Omni-router, a shared routing mechanism across MoE layers. This inductive bias leads to more coherent and specialized expert usage, improving training stability and achieving lower WER than both dense and Switch Transformer baselines across diverse, out-of-domain benchmarks. Empirical results on a large pseudo-labeled SpeechCrawl dataset and Libriheavy-based tests show average relative WER reductions of up to 11.2% versus dense and 8.2% versus Switch Transformer, along with demonstrated robustness and structured routing patterns. The approach is simple to implement (no extra embedding networks or auxiliary losses beyond load balancing) and scales across different numbers of experts and model sizes, highlighting the importance of inter-layer routing coordination in MoE-based ASR systems.

Abstract

Mixture-of-experts (MoE) architectures have expanded from language modeling to automatic speech recognition (ASR). Traditional MoE methods, such as the Switch Transformer, route experts independently within each layer. Our analysis reveals that routers in most layers make expert choices that are not strongly correlated with the choices of the routers in other layers. To increase the cooperation between experts in different layers and encourage greater specialization, we use a shared router across different MoE layers. We call this model Omni-router Transformer. Extensive experiments on a large-scale pseudo-labeled dataset and evaluations across 10 diverse, out-of-domain ASR benchmarks demonstrate that the Omni-router Transformer is able to achieve lower training loss and consistently outperform dense and Switch Transformer models, reducing average word error rates by 11.2% and 8.2%, respectively, while providing structured expert usage and improved robustness to diverse data.

Figures (7)

• Figure 1: Illustration of Switch (left) and Omni-router (right) Transformers. $E_i$ represents an expert.
• Figure 2: Visualization of randomly selected audio samples from the Librispeech dev-other dataset, showing (top) their corresponding mel-spectrograms; (middle) expert usage patterns in a four-expert (blue, orange, green, pink) Switch Transformer ASR model; and (bottom) expert usage patterns in a four-expert Omni-router Transformer ASR model. Expert assignments are displayed across transformer layers ($l_i$, y-axis) and time ($t$, x-axis).
• Figure 3: Contingency tables illustrate the expert correspondence between layers 14 and 15 for both Switch Transformer and Omni-router Transformer as representative examples.
• Figure 4: Expert correlation (Cramér's V) between adjacent layers in Switch Transformer and Omni-router Transformer.
• Figure 5: Expert entropy across layers of Switch Transformer and Omni-router Transformer.