SoftMoE:可微分专家路由的混合专家模型

SoftMoE: Soft Differentiable Routing for Mixture-of-Experts in LLMs

精选理由

稀疏MoE的top-k路由不灵活还浪费算力,SoftMoE用可微路由让模型自己学会少用专家,性能却不输,代码开源了。

AI 摘要

SoftMoE用截断的soft top-k LapSum松弛替换传统稀疏MoE的离散top-k路由,实现专家路由的可微分化。模型参数化每层平均激活专家数并施加全局预算,使容量分配可学习。在语言建模和下游任务上,SoftMoE性能与稀疏MoE相当或更优,但激活专家数量更少。实验显示分配呈高度非均匀性,后层激活更多专家。

原文 · arXiv cs.AI

SoftMoE: Soft Differentiable Routing for Mixture-of-Experts in LLMs

Sparse Mixture-of-Experts (MoE) architectures enable scaling LLM parameters under a fixed inference budget by activating only a small subset of experts via top-$k$ routing. While this preserves causality and suits autoregressive language models, the discrete top-$k$ operator is not differentiable, forcing a fixed number of active experts per input and resulting in inefficient use of computation. We propose SoftMoE, which replaces discrete routing with a truncated soft top-$k$ LapSum relaxation, allowing gradient-based optimization of expert routing. We further parameterize the mean number of active experts per layer and impose a global budget constraint, enabling the model to learn how to allocate expert capacity across layers. SoftMoE remains fully compatible with autoregressive modeling and achieves performance comparable to or better than sparse MoE on language modeling and downstream tasks, while activating significantly fewer experts. Notably, the learned allocation is highly non-uniform, with later layers activating more experts. The source code is publicly available$^\dagger$.