Understanding the Mixture-of-Experts with Nadaraya-Watson Kernel

Mixture-of-Experts (MoE) has become a cornerstone in recent state-of-the-art large language models (LLMs). Traditionally, MoE relies on Softmax as the router score function to aggregate expert output, a designed choice that has persisted from the earliest MoE models to modern LLMs, and is now widely regarded as standard practice. However, the necessity of using Softmax to project router weights into a probability simplex remains an unchallenged assumption rather than a principled design choice. In this work, we first revisit the classical Nadaraya-Watson regression and observe that MoE shares the same mathematical formulation as Nadaraya-Watson regression. Furthermore, we show that both feed-forward neural network (FFN) and MoE can be interpreted as a special case of Nadaraya-Watson regression, where the kernel function corresponds to the input neurons of the output layer. Motivated by these insights, we propose the zero-additional-cost Kernel Inspired Router with Normalization (KERN), an FFN-style router function, as an alternative to Softmax. We demonstrate that this router generalizes both Sigmoid- and Softmax-based routers. textbfBased on empirical observations and established practices in FFN implementation, we recommend the use of ReLU activation and ℓ2-normalization in KERN router function. Comprehensive experiments in MoE and LLM validate the effectiveness of the proposed FFN-style router function methodNorm.