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 $\mathrm{S}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}$ 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 $\mathrm{S}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}$ 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 $\mathbf{\text{zero-additional-cost}}$ Kernel Inspired Router with Normalization (KERN), an FFN-style router function, as an alternative to $\mathrm{S}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}$. We demonstrate that this router generalizes both $\mathrm{S}\mathrm{i}\mathrm{g}\mathrm{m}\mathrm{o}\mathrm{i}\mathrm{d}$- and $\mathrm{S}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}$-based routers. $\mathbf{\text{Based}}\mathbf{\text{on}}\mathbf{\text{empirical}}\mathbf{\text{observations}}\mathbf{\text{and}}\mathbf{\text{established}}\mathbf{\text{practices}}\mathbf{\text{in}}\mathbf{\text{FFN}}\mathbf{\text{implementation,}}\mathbf{\text{we}}\mathbf{\text{recommend}}\mathbf{\text{the}}\mathbf{\text{use}}\mathbf{\text{of}}\mathrm{R}\mathrm{e}\mathrm{L}\mathrm{U}\mathbf{\text{activation}}\mathbf{\text{and}}{\mathrm{\ell }}_2\mathbf{\text{-normalization}}\mathbf{\text{in}}\mathrm{K}\mathrm{E}\mathrm{R}\mathrm{N}\mathbf{\text{router}}\mathbf{\text{function.}}$ Comprehensive experiments in MoE and LLM validate the effectiveness of the proposed FFN-style router function \methodNorm.