Skip to content

Repository files navigation

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

Arxiv Blog

This software project accompanies the research paper, Omni-router: Sharing Routing Decisions in Sparse Mixture-of-Experts for Speech Recognition, accepted in ASRU 2025.

Model Architecture

Omni-router Transformer is a new Mixture-of-Experts (MoE) architecture that explicitly couples routing across layers using a shared router to learn strong and specialized experts. Omni-router's routing decisions appear to form consistent temporal segments and strutured usage across model depth, suggesting meaningful coordination between layers.

Model Architecture

Omni-router ASR achieves better performance compared with similar size dense model and Whisper-small.en model, with half of the activate parameters.

Omni-router Dense Whisper-small.en
246M (140M activate) 210M 244M
AMI-IHM 17.8 17.8 17.6
Chime6 28.2 27.9 27.7
CommonVoice 16.4 17.6 15.2
Fleurs 8.4 8.8 7.6
Callhome 15.0 15.2 20
Switchboard 13.5 13.8 14.9
Librispeech (clean) 3.3 3.7 3.1
Librispeech (other) 7.3 8 7.4
WSJ (nov92) 3.7 4.1 3.5
Tedlium 4.2 4.3 4
Voxpopuli 8.5 9 8.2
Average 11.48 11.84 11.75

🚀 Getting Started

1. Installation

pip install .

2. Prepare Checkpoints

Model checkpoints can be downloaded from Huggingface with

export HF_TOKEN="your_actual_token_here"
python download_model.py

Model Cards

Models trained with the internal SpeechCrawl data (click to download manually).

Model Streaming
omni-router-speechcrawl-asr-0.25b-v1 no
omni-router-speechcrawl-streaming-asr-0.6b-v1 yes

For reproducibility, we also train our models with Libriheavy dataset and make dense, switch and omni-router checkpoints available.

Model MoE
dense-asr-libriheavy-0.08b no
moe-asr-libriheavy-0.5b switch
omni-router-asr-libriheavy-0.5b omni-router

3. Batched Inference

The model can be used for batched inference.

from evaluate import load_model_and_extractor, read_audio, extract_feature, forward_model, merge_transcriptions
import torch

device = torch.device("cpu")
asr, tokenizer, mel_func = load_model_and_extractor("omni-router-speechcrawl-asr-0.25b-v1", device)

filenames = ["example/7976-110124-0000.wav", "example/7976-110124-0001.wav"]
sample, sample_length, file_chunk_mapping = read_audio(filenames, device)
sample, sample_length = extract_feature(mel_func, sample, sample_length)
transcriptions = forward_model(asr, tokenizer, sample, sample_length)
transcriptions = merge_transcriptions(file_chunk_mapping, transcriptions, len(filenames))
print(transcriptions)

# ['Long ago, there lived a merchant who had three daughters.', 'Every year, at a certain day of a certain month, he went away to a distant city to collect money on an account.']

📚 Citation

If you find this work useful, please cite our paper:

@article{gu2025omnirouter,
  title={Omni-router: Sharing Routing Decisions in Sparse Mixture-of-Experts for Speech Recognition},
  author={Gu, Zijin and Likhomanenko, Tatiana and Jaitly, Navdeep},
  journal={arXiv preprint arXiv:2507.05724},
  year={2025}
}

📝 License

  • Please check out the repository LICENSE before using the provided code and LICENSE_MODEL for the released models.

About

No description, website, or topics provided.

Resources

Code of conduct

Contributing

Security policy

Stars

18 stars

Watchers

1 watching

Forks

Releases

Packages

Used by

Contributors

Languages