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

📄 arXiv: 2509.25913v2 📥 PDF

作者: Chuanyang Zheng, Jiankai Sun, Yihang Gao, Enze Xie, Yuehao Wang, Peihao Wang, Ting Xu, Matthew Chang, Liliang Ren, Jingyao Li, Jing Xiong, Kashif Rasul, Mac Schwager, Anderson Schneider, Zhangyang Wang, Yuriy Nevmyvaka

分类: cs.CL

发布日期: 2025-09-30 (更新: 2025-10-14)

备注: Tech Report

💡 一句话要点

提出KERN以替代Softmax解决MoE中的路由问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合专家模型 Nadaraya-Watson回归 路由函数 深度学习 大型语言模型

📋 核心要点

  1. 现有的MoE模型普遍依赖Softmax作为路由函数，但这一选择缺乏理论基础，限制了模型的灵活性和性能。
  2. 本文提出了一种新的FFN风格的路由函数KERN，基于Nadaraya-Watson回归的数学框架，旨在替代传统的Softmax路由。
  3. 实验结果表明，KERN在多个MoE和LLM任务中表现优越，能够有效提升模型的性能和稳定性。

📝 摘要（中文）

混合专家模型（MoE）已成为最新大型语言模型（LLMs）的基石。传统上，MoE依赖Softmax作为路由得分函数来聚合专家输出，这一设计选择在早期MoE模型中延续至今，并被广泛视为标准做法。然而，使用Softmax将路由权重投影到概率单纯形的必要性仍然是一个未经挑战的假设。本文首先重新审视经典的Nadaraya-Watson回归，发现MoE与Nadaraya-Watson回归具有相同的数学形式。我们提出了一种零额外成本的内核启发路由器KERN，作为Softmax的替代方案，并通过全面实验验证了其在MoE和LLM中的有效性。

🔬 方法详解

问题定义：本文旨在解决传统MoE模型中使用Softmax作为路由函数的局限性，探讨其理论基础不足的问题。

核心思路：通过重新审视Nadaraya-Watson回归，提出KERN作为一种新的路由函数，能够在不增加额外计算成本的情况下，替代Softmax并实现更好的性能。

技术框架：KERN的整体架构包括输入层、路由层和输出层，其中路由层采用基于ReLU激活和L2归一化的设计，确保了输出的有效性和稳定性。

关键创新：KERN的最大创新在于其将路由函数与Nadaraya-Watson回归相结合，提供了一种新的视角来理解和实现MoE模型的路由机制，突破了Softmax的限制。

关键设计：KERN路由函数使用ReLU激活和L2归一化，确保了输出的合理性和有效性，且在实验中表现出色，验证了其在实际应用中的可行性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，KERN在多个基准测试中相较于传统Softmax路由函数，性能提升幅度达到10%以上，验证了其在MoE和LLM中的有效性和优越性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、计算机视觉等需要高效模型的任务。KERN的提出为未来的MoE模型设计提供了新的思路，可能会在大规模模型训练和推理中产生深远影响。

📄 摘要（原文）

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.