MixtureKit: 一种用于组合、训练和可视化混合专家模型的通用框架
MixtureKit: A General Framework for Composing, Training, and Visualizing Mixture-of-Experts Models
AI总结:
MixtureKit提供了一种通用框架,用于组合、训练和可视化混合专家模型,通过三种方法提升模型性能与灵活性。
摘要:
我们介绍了MixtureKit,一种模块化的开源框架,用于从任意预训练或微调的模型中构建、训练和分析混合专家(MoE)模型。MixtureKit目前支持三种互补的方法:(i)传统MoE,使用每个Transformer块一个路由器选择专家;(ii)BTX(分支-训练-混合),为每个指定的子层引入单独的路由器,实现细粒度的token路由;(iii)BTS(分支-训练-缝合),保持专家完全完整,并引入可训练的缝合层,用于在hub和专家之间受控的信息交换。MixtureKit会自动修改模型配置,修补解码器和因果语言模型类,并保存一个统一的检查点,用于推理或微调。我们进一步提供了一个可视化界面,以检查每个token的路由决策、专家权重分布和层间贡献。使用多语言代码切换数据(例如阿拉伯语-拉丁语)的实验表明,使用MixtureKit训练的BTX模型在多个基准上优于基线密集模型。我们发布MixtureKit作为研究和开发MoE系统在不同领域的实用基础。
英文摘要:
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) \emph{Traditional MoE}, which uses a single router per transformer block to select experts, (ii) \emph{BTX} (Branch-Train-Mix), which introduces separate routers for each specified sub-layer enabling fine-grained token routing, and (iii) \emph{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.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。