Understanding the Mixture-of-Experts with Nadaraya-Watson Kernel
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866911208256307200 |
|---|---|
| author | Zheng, Chuanyang Sun, Jiankai Gao, Yihang Xie, Enze Wang, Yuehao Wang, Peihao Xu, Ting Chang, Matthew Ren, Liliang Li, Jingyao Xiong, Jing Rasul, Kashif Schwager, Mac Schneider, Anderson Wang, Zhangyang Nevmyvaka, Yuriy |
| author_facet | Zheng, Chuanyang Sun, Jiankai Gao, Yihang Xie, Enze Wang, Yuehao Wang, Peihao Xu, Ting Chang, Matthew Ren, Liliang Li, Jingyao Xiong, Jing Rasul, Kashif Schwager, Mac Schneider, Anderson Wang, Zhangyang Nevmyvaka, Yuriy |
| contents | Mixture-of-Experts (MoE) has become a cornerstone in recent state-of-the-art large language models (LLMs). Traditionally, MoE relies on $\mathrm{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 $\mathrm{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 \textbf{zero-additional-cost} Kernel Inspired Router with Normalization (KERN), an FFN-style router function, as an alternative to $\mathrm{Softmax}$. We demonstrate that this router generalizes both $\mathrm{Sigmoid}$- and $\mathrm{Softmax}$-based routers. \textbf{Based on empirical observations and established practices in FFN implementation, we recommend the use of $\mathrm{ReLU}$ activation and $\ell_2$-normalization in $\mathrm{KERN}$ router function.} Comprehensive experiments in MoE and LLM validate the effectiveness of the proposed FFN-style router function \methodNorm. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2509_25913 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Understanding the Mixture-of-Experts with Nadaraya-Watson Kernel Zheng, Chuanyang Sun, Jiankai Gao, Yihang Xie, Enze Wang, Yuehao Wang, Peihao Xu, Ting Chang, Matthew Ren, Liliang Li, Jingyao Xiong, Jing Rasul, Kashif Schwager, Mac Schneider, Anderson Wang, Zhangyang Nevmyvaka, Yuriy Computation and Language Mixture-of-Experts (MoE) has become a cornerstone in recent state-of-the-art large language models (LLMs). Traditionally, MoE relies on $\mathrm{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 $\mathrm{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 \textbf{zero-additional-cost} Kernel Inspired Router with Normalization (KERN), an FFN-style router function, as an alternative to $\mathrm{Softmax}$. We demonstrate that this router generalizes both $\mathrm{Sigmoid}$- and $\mathrm{Softmax}$-based routers. \textbf{Based on empirical observations and established practices in FFN implementation, we recommend the use of $\mathrm{ReLU}$ activation and $\ell_2$-normalization in $\mathrm{KERN}$ router function.} Comprehensive experiments in MoE and LLM validate the effectiveness of the proposed FFN-style router function \methodNorm. |
| title | Understanding the Mixture-of-Experts with Nadaraya-Watson Kernel |
| topic | Computation and Language |
| url | https://arxiv.org/abs/2509.25913 |