MixtureKit: A General Framework for Composing, Training, and Visualizing Mixture-of-Experts Models
Abstract
MixtureKit is an open-source framework for building and analyzing Mixture-of-Experts models via traditional routing, fine-grained branch routing, and stitch-based expert merging, with visualization tools and multilingual evaluation.
We introduce MixtureKit, a modular open-source framework for constructing, training, and analyzing Mixture-of-Experts (MoE) models from arbitrary pre-trained or fine-tuned models. MixtureKit currently supports three complementary methods: (i) Traditional MoE, which uses a single router per transformer block to select experts, (ii) BTX (Branch-Train-Mix), which introduces separate routers for each specified sub-layer enabling fine-grained token routing, and (iii) BTS (Branch-Train-Stitch), which keeps experts fully intact and introduces trainable stitch layers for controlled information exchange between hub and experts. MixtureKit automatically modifies the model configuration, patches decoder and causal LM classes, and saves a unified checkpoint ready for inference or fine-tuning. We further provide a visualization interface to inspect per-token routing decisions, expert weight distributions, and layer-wise contributions. Experiments with multilingual code-switched data (e.g. Arabic-Latin) show that a BTX-based model trained using MixtureKit can outperform baseline dense models on multiple benchmarks. We release MixtureKit as a practical foundation for research and development of MoE-based systems across diverse domains.
Community
Code: https://github.com/MBZUAI-Paris/MixtureKit
Install: pip install mixturekit
Docs: https://mbzuai-paris.github.io/MixtureKit/
Colab (build, train and visualize an MoE in 15 min): https://colab.research.google.com/github/MBZUAI-Paris/MixtureKit/blob/main/examples/build_your_own_moe.ipynb
Models citing this paper 0
No model linking this paper
Datasets citing this paper 0
No dataset linking this paper
Spaces citing this paper 0
No Space linking this paper
Collections including this paper 0
No Collection including this paper